Codex changes

screenshot/src/server.ts

@@ -1,10 +1,14 @@
-import express from 'express';
+import express, { type Request } from 'express';
 import { ScreenshotCache } from './cache.js';
 import { takeScreenshot, checkWebGL, closeBrowser, initialize } from './screenshot.js';
+import { buildScreenshotRequest, ValidationError } from './validation.js';
 
 const PORT = parseInt(process.env.PORT || '8002', 10);
 const APP_URL = process.env.APP_URL;
 const CACHE_DIR = process.env.CACHE_DIR;
+const SCREENSHOT_CONCURRENCY = parsePositiveIntEnv('SCREENSHOT_CONCURRENCY', 3);
+const RATE_LIMIT_WINDOW_MS = parsePositiveIntEnv('SCREENSHOT_RATE_WINDOW_MS', 60_000);
+const RATE_LIMIT_MAX = parsePositiveIntEnv('SCREENSHOT_RATE_LIMIT', 30);
 
 if (!APP_URL) {
   console.error('Error: APP_URL environment variable is required');
@@ -18,6 +22,58 @@ if (!CACHE_DIR) {
 
 const cache = new ScreenshotCache(CACHE_DIR);
 const app = express();
+app.set('trust proxy', true);
+
+let activeScreenshots = 0;
+let lastRateLimitPrune = 0;
+const rateLimitBuckets = new Map<string, { count: number; resetAt: number }>();
+
+function parsePositiveIntEnv(name: string, fallback: number): number {
+  const value = Number.parseInt(process.env[name] || '', 10);
+  return Number.isFinite(value) && value > 0 ? value : fallback;
+}
+
+function acquireScreenshotSlot(): (() => void) | null {
+  if (activeScreenshots >= SCREENSHOT_CONCURRENCY) {
+    return null;
+  }
+  activeScreenshots += 1;
+  let released = false;
+  return () => {
+    if (released) return;
+    released = true;
+    activeScreenshots = Math.max(0, activeScreenshots - 1);
+  };
+}
+
+function rateLimitKey(req: Request): string {
+  const forwardedFor = req.get('x-forwarded-for')?.split(',')[0]?.trim();
+  return forwardedFor || req.ip || req.socket.remoteAddress || 'unknown';
+}
+
+function allowScreenshotRequest(req: Request): boolean {
+  const now = Date.now();
+  if (now - lastRateLimitPrune > RATE_LIMIT_WINDOW_MS) {
+    lastRateLimitPrune = now;
+    for (const [key, bucket] of rateLimitBuckets) {
+      if (bucket.resetAt <= now) {
+        rateLimitBuckets.delete(key);
+      }
+    }
+  }
+
+  const key = rateLimitKey(req);
+  let bucket = rateLimitBuckets.get(key);
+  if (!bucket || bucket.resetAt <= now) {
+    bucket = { count: 0, resetAt: now + RATE_LIMIT_WINDOW_MS };
+    rateLimitBuckets.set(key, bucket);
+  }
+  if (bucket.count >= RATE_LIMIT_MAX) {
+    return false;
+  }
+  bucket.count += 1;
+  return true;
+}
 
 app.get('/health', (_req, res) => {
   res.status(200).send('ok');
@@ -33,28 +89,9 @@ app.get('/debug', async (_req, res) => {
 });
 
 app.get('/screenshot', async (req, res) => {
+  let releaseSlot: (() => void) | null = null;
   try {
-    const qs = new URLSearchParams();
-    for (const key of ['lat', 'lon', 'zoom', 'tab', 'og']) {
-      const val = req.query[key];
-      if (typeof val === 'string' && val) {
-        qs.set(key, val);
-      }
-    }
-    // Repeated params: filter, poi, tt (travel time)
-    for (const key of ['filter', 'poi', 'tt']) {
-      const val = req.query[key];
-      if (typeof val === 'string' && val) {
-        qs.append(key, val);
-      } else if (Array.isArray(val)) {
-        for (const v of val) {
-          if (typeof v === 'string' && v) qs.append(key, v);
-        }
-      }
-    }
-
-    // Extract path param (used for non-root pages like /invite/CODE)
-    const pagePath = typeof req.query.path === 'string' && req.query.path ? req.query.path : '/';
+    const { pagePath, qs } = buildScreenshotRequest(req.query as Record<string, unknown>);
     if (pagePath !== '/') qs.set('path', pagePath);
 
     // Include auth status in cache key so authenticated screenshots
@@ -75,6 +112,17 @@ app.get('/screenshot', async (req, res) => {
       return;
     }
 
+    if (!allowScreenshotRequest(req)) {
+      res.status(429).json({ error: 'Screenshot rate limit exceeded' });
+      return;
+    }
+
+    releaseSlot = acquireScreenshotSlot();
+    if (!releaseSlot) {
+      res.status(503).json({ error: 'Screenshot service busy' });
+      return;
+    }
+
     // Build the URL for the frontend in screenshot mode
     qs.set('screenshot', '1');
     const url = `${APP_URL}${pagePath}?${qs}`;
@@ -90,8 +138,14 @@ app.get('/screenshot', async (req, res) => {
     res.setHeader('X-Cache', 'MISS');
     res.send(jpeg);
   } catch (err) {
+    if (err instanceof ValidationError) {
+      res.status(err.status).json({ error: err.message });
+      return;
+    }
     console.error('Screenshot failed:', err);
     res.status(500).json({ error: 'Screenshot failed' });
+  } finally {
+    releaseSlot?.();
   }
 });
 
@@ -99,6 +153,8 @@ const server = app.listen(PORT, () => {
   console.log(`Screenshot service listening on port ${PORT}`);
   console.log(`  APP_URL: ${APP_URL}`);
   console.log(`  CACHE_DIR: ${CACHE_DIR}`);
+  console.log(`  SCREENSHOT_CONCURRENCY: ${SCREENSHOT_CONCURRENCY}`);
+  console.log(`  SCREENSHOT_RATE_LIMIT: ${RATE_LIMIT_MAX}/${RATE_LIMIT_WINDOW_MS}ms`);
 
   // Pre-warm browser and populate network cache in background.
   // The health endpoint is available immediately; screenshot requests