From cd55127bbd0e0bd1e437061b9775a05d76191638 Mon Sep 17 00:00:00 2001
From: shrimbly <falloon.willie@gmail.com>
Date: Sat, 7 Feb 2026 20:55:51 +1300
Subject: [PATCH] fix: add SSRF validation for fal.ai queue URLs, fix CDN
 upload test mocks

Validate fal-provided status_url and response_url against validateMediaUrl()
and domain allowlist (queue.fal.run) before using them for polling. Falls back
to constructed URLs on validation failure.

Also updates mockFalCdnUpload test helper to match the two-step CDN upload
flow (initiate + PUT) introduced in 3fea13e.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
---
 src/app/api/generate/__tests__/route.test.ts | 11 +++++++--
 src/app/api/generate/route.ts                | 26 +++++++++++++++++---
 2 files changed, 32 insertions(+), 5 deletions(-)

diff --git a/src/app/api/generate/__tests__/route.test.ts b/src/app/api/generate/__tests__/route.test.ts
index bc1ca5db..079f7f6e 100644
--- a/src/app/api/generate/__tests__/route.test.ts
+++ b/src/app/api/generate/__tests__/route.test.ts
@@ -1650,11 +1650,18 @@ describe("/api/generate route", () => {
       });
     }
 
-    // Helper: mock CDN upload for a base64 image
+    // Helper: mock two-step CDN upload for a base64 image
+    // Step 1: POST initiate → { upload_url, file_url }
+    // Step 2: PUT binary data → ok
     function mockFalCdnUpload(cdnUrl = "https://fal.ai/cdn/uploaded.png") {
+      // Initiate upload
+      mockFetch.mockResolvedValueOnce({
+        ok: true,
+        json: () => Promise.resolve({ upload_url: "https://fal.ai/cdn/put-target", file_url: cdnUrl }),
+      });
+      // PUT binary data
       mockFetch.mockResolvedValueOnce({
         ok: true,
-        json: () => Promise.resolve({ url: cdnUrl }),
       });
     }
 
diff --git a/src/app/api/generate/route.ts b/src/app/api/generate/route.ts
index e712ca17..59f121d3 100644
--- a/src/app/api/generate/route.ts
+++ b/src/app/api/generate/route.ts
@@ -1002,9 +1002,29 @@ async function generateWithFalQueue(
     };
   }
 
-  // Use URLs from response if provided, otherwise construct them
-  const statusUrl = submitResult.status_url || `https://queue.fal.run/${modelId}/requests/${falRequestId}/status`;
-  const responseUrl = submitResult.response_url || `https://queue.fal.run/${modelId}/requests/${falRequestId}`;
+  // Use URLs from response if provided, with SSRF validation; fall back to constructed URLs
+  const fallbackStatusUrl = `https://queue.fal.run/${modelId}/requests/${falRequestId}/status`;
+  const fallbackResponseUrl = `https://queue.fal.run/${modelId}/requests/${falRequestId}`;
+  let statusUrl = fallbackStatusUrl;
+  let responseUrl = fallbackResponseUrl;
+
+  if (submitResult.status_url) {
+    const statusCheck = validateMediaUrl(submitResult.status_url);
+    if (statusCheck.valid && submitResult.status_url.startsWith('https://queue.fal.run/')) {
+      statusUrl = submitResult.status_url;
+    } else {
+      console.warn(`[API:${requestId}] fal.ai provided invalid status URL: ${submitResult.status_url} — falling back to constructed URL`);
+    }
+  }
+  if (submitResult.response_url) {
+    const responseCheck = validateMediaUrl(submitResult.response_url);
+    if (responseCheck.valid && submitResult.response_url.startsWith('https://queue.fal.run/')) {
+      responseUrl = submitResult.response_url;
+    } else {
+      console.warn(`[API:${requestId}] fal.ai provided invalid response URL: ${submitResult.response_url} — falling back to constructed URL`);
+    }
+  }
+
   console.log(`[API:${requestId}] Queue request submitted: ${falRequestId}, status URL: ${statusUrl}`);
 
   // Poll for completion