From 7c3b32ae5850a8efb5e7a473fd13e4c5f143b65f Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 15 Jan 2026 02:12:58 +0000 Subject: [PATCH] fix: address PR review comments for security, CSS, and test isolation Security & API route fixes: - Remove API key query param support in fal/models route to prevent credential leakage - Fix MIME type fallback in save-generation route (now uses prefix-based fallback: image/* -> png, video/* -> mp4, unknown -> bin) - Add timeout (60s) and size limits (500MB) for remote fetch in save-generation route using AbortController CSS fix: - Consolidate duplicate :root blocks in globals.css - Change handle selector from data-handletype="prompt" to data-handletype="text" - Update --handle-color-text CSS variable from #a3a3a3 to #3b82f6 Test file global mock isolation: - Fix useAnnotationStore mock to support Zustand selector pattern - Use vi.stubGlobal() instead of direct global assignments to prevent test leakage - Restore global.Image in afterAll hook in AnnotationModal.test.tsx - Wrap async Done assertion in waitFor() in AnnotationModal.test.tsx - Use idiomatic fireEvent.dragOver() in AnnotationNode.test.tsx - Fix global.confirm stub with proper beforeAll/afterAll lifecycle in CostDialog.test.tsx - Import ReactNode type properly in EdgeToolbar.test.tsx - Update save-generation route test to expect AbortController signal in fetch call --- src/app/api/providers/fal/models/route.ts | 9 ++- .../save-generation/__tests__/route.test.ts | 6 +- src/app/api/save-generation/route.ts | 68 ++++++++++++++++--- src/app/globals.css | 13 ++-- .../__tests__/AnnotationModal.test.tsx | 35 +++++++--- .../__tests__/AnnotationNode.test.tsx | 34 ++++++---- src/components/__tests__/CostDialog.test.tsx | 12 +++- src/components/__tests__/EdgeToolbar.test.tsx | 3 +- 8 files changed, 132 insertions(+), 48 deletions(-) diff --git a/src/app/api/providers/fal/models/route.ts b/src/app/api/providers/fal/models/route.ts index 8adc71dd..5e7836e9 100644 --- a/src/app/api/providers/fal/models/route.ts +++ b/src/app/api/providers/fal/models/route.ts @@ -98,17 +98,20 @@ type ModelsResponse = ModelsSuccessResponse | ModelsErrorResponse; * Fetches available models from fal.ai API. * API key is optional - fal.ai works without but with rate limits. * + * Headers: + * - X-API-Key: API key for authentication (recommended) + * - Authorization: Alternative auth header + * * Query params: * - search: Optional search query to filter models - * - api_key: Alternative to X-API-Key header */ export async function GET( request: NextRequest ): Promise> { - // Get optional API key from header or query param + // Get optional API key from header only (never from query params to avoid credential leakage) const apiKey = request.headers.get("X-API-Key") || - request.nextUrl.searchParams.get("api_key"); + request.headers.get("Authorization")?.replace(/^Key\s+/i, ""); const searchQuery = request.nextUrl.searchParams.get("search"); diff --git a/src/app/api/save-generation/__tests__/route.test.ts b/src/app/api/save-generation/__tests__/route.test.ts index b9293d5c..27434004 100644 --- a/src/app/api/save-generation/__tests__/route.test.ts +++ b/src/app/api/save-generation/__tests__/route.test.ts @@ -264,7 +264,11 @@ describe("/api/save-generation route", () => { expect(data.success).toBe(true); expect(data.filename).toContain(expectedHash); - expect(global.fetch).toHaveBeenCalledWith("https://example.com/image.png"); + // Fetch is called with URL and options object containing AbortController signal + expect(global.fetch).toHaveBeenCalledWith( + "https://example.com/image.png", + expect.objectContaining({ signal: expect.any(AbortSignal) }) + ); }); it("should handle failed HTTP fetch", async () => { diff --git a/src/app/api/save-generation/route.ts b/src/app/api/save-generation/route.ts index c92aee6e..9140afb5 100644 --- a/src/app/api/save-generation/route.ts +++ b/src/app/api/save-generation/route.ts @@ -11,11 +11,27 @@ function getExtensionFromMime(mimeType: string): string { "image/jpeg": "jpg", "image/gif": "gif", "image/webp": "webp", + "image/svg+xml": "svg", "video/mp4": "mp4", "video/webm": "webm", "video/quicktime": "mov", }; - return mimeToExt[mimeType] || "mp4"; + + // Check explicit mapping first + if (mimeToExt[mimeType]) { + return mimeToExt[mimeType]; + } + + // Fallback based on MIME type prefix + if (mimeType.startsWith("image/")) { + return "png"; + } + if (mimeType.startsWith("video/")) { + return "mp4"; + } + + // Unknown type - use generic binary extension + return "bin"; } // Helper to detect if a string is an HTTP URL @@ -106,16 +122,48 @@ export async function POST(request: NextRequest) { // Handle HTTP URL (common for large video files from providers) logger.info('file.save', 'Fetching content from URL', { url: content.substring(0, 100) }); - const response = await fetch(content); - if (!response.ok) { - throw new Error(`Failed to fetch content: ${response.status} ${response.statusText}`); + // Set up timeout to prevent hanging requests (60 seconds for large video files) + const FETCH_TIMEOUT_MS = 60000; + const MAX_CONTENT_SIZE = 500 * 1024 * 1024; // 500MB max + + const controller = new AbortController(); + const timeoutId = setTimeout(() => controller.abort(), FETCH_TIMEOUT_MS); + + try { + const response = await fetch(content, { signal: controller.signal }); + clearTimeout(timeoutId); + + if (!response.ok) { + throw new Error(`Failed to fetch content: ${response.status} ${response.statusText}`); + } + + // Check content-length before downloading to avoid excessive bandwidth usage + const contentLength = response.headers.get("content-length"); + if (contentLength) { + const size = parseInt(contentLength, 10); + if (size > MAX_CONTENT_SIZE) { + throw new Error(`Content size ${size} bytes exceeds maximum allowed ${MAX_CONTENT_SIZE} bytes`); + } + } + + const contentType = response.headers.get("content-type") || (isVideo ? "video/mp4" : "image/png"); + extension = getExtensionFromMime(contentType); + + const arrayBuffer = await response.arrayBuffer(); + + // Double-check actual size after download + if (arrayBuffer.byteLength > MAX_CONTENT_SIZE) { + throw new Error(`Downloaded content size ${arrayBuffer.byteLength} bytes exceeds maximum allowed ${MAX_CONTENT_SIZE} bytes`); + } + + buffer = Buffer.from(arrayBuffer); + } catch (fetchError) { + clearTimeout(timeoutId); + if (fetchError instanceof Error && fetchError.name === 'AbortError') { + throw new Error(`Fetch timed out after ${FETCH_TIMEOUT_MS}ms`); + } + throw fetchError; } - - const contentType = response.headers.get("content-type") || (isVideo ? "video/mp4" : "image/png"); - extension = getExtensionFromMime(contentType); - - const arrayBuffer = await response.arrayBuffer(); - buffer = Buffer.from(arrayBuffer); } else { // Handle base64 data URL const dataUrlMatch = content.match(/^data:([\w/+-]+);base64,/); diff --git a/src/app/globals.css b/src/app/globals.css index ab09c649..9fd8482f 100644 --- a/src/app/globals.css +++ b/src/app/globals.css @@ -7,6 +7,9 @@ :root { --background: #0a0a0a; --foreground: #fafafa; + /* Handle colors (used by labels) */ + --handle-color-image: #10b981; + --handle-color-text: #3b82f6; } html { @@ -70,17 +73,11 @@ body { background: #10b981; } -/* Prompt/text handles - soft blue */ -.react-flow__handle[data-handletype="prompt"] { +/* Text handles - soft blue */ +.react-flow__handle[data-handletype="text"] { background: #3b82f6; } -/* CSS variables for handle colors (used by labels) */ -:root { - --handle-color-image: #10b981; - --handle-color-text: #a3a3a3; -} - .react-flow__edge-path { stroke: #94a3b8; stroke-width: 3; diff --git a/src/components/__tests__/AnnotationModal.test.tsx b/src/components/__tests__/AnnotationModal.test.tsx index d90d92a5..f974e160 100644 --- a/src/components/__tests__/AnnotationModal.test.tsx +++ b/src/components/__tests__/AnnotationModal.test.tsx @@ -1,12 +1,13 @@ -import { describe, it, expect, vi, beforeEach, afterEach } from "vitest"; -import { render, screen, fireEvent } from "@testing-library/react"; +import { describe, it, expect, vi, beforeEach, afterEach, beforeAll, afterAll } from "vitest"; +import { render, screen, fireEvent, waitFor } from "@testing-library/react"; +import type { ReactNode } from "react"; import { AnnotationModal } from "@/components/AnnotationModal"; import { ToolType, ToolOptions } from "@/types"; // Mock Konva and react-konva - canvas internals cannot be tested in jsdom vi.mock("react-konva", () => ({ - Stage: ({ children }: { children: React.ReactNode }) =>
{children}
, - Layer: ({ children }: { children: React.ReactNode }) =>
{children}
, + Stage: ({ children }: { children: ReactNode }) =>
{children}
, + Layer: ({ children }: { children: ReactNode }) =>
{children}
, Image: () =>
, Rect: () =>
, Ellipse: () =>
, @@ -74,7 +75,12 @@ const createMockAnnotationStore = (overrides = {}) => ({ let mockAnnotationStore = createMockAnnotationStore(); vi.mock("@/store/annotationStore", () => ({ - useAnnotationStore: () => mockAnnotationStore, + useAnnotationStore: (selector?: (state: ReturnType) => unknown) => { + if (selector) { + return selector(mockAnnotationStore); + } + return mockAnnotationStore; + }, })); // Mock workflow store @@ -104,12 +110,17 @@ class MockImage { } } -Object.defineProperty(global, "Image", { - value: MockImage, - writable: true, -}); +// Store original Image to restore later +const OriginalImage = global.Image; describe("AnnotationModal", () => { + beforeAll(() => { + vi.stubGlobal("Image", MockImage); + }); + + afterAll(() => { + vi.stubGlobal("Image", OriginalImage); + }); beforeEach(() => { vi.clearAllMocks(); mockAnnotationStore = createMockAnnotationStore(); @@ -275,8 +286,10 @@ describe("AnnotationModal", () => { // so we just verify the button exists and triggers the handler fireEvent.click(screen.getByText("Done")); - // closeModal should be called after save - expect(mockCloseModal).toHaveBeenCalled(); + // closeModal should be called after save (async operation) + await waitFor(() => { + expect(mockCloseModal).toHaveBeenCalled(); + }); }); }); diff --git a/src/components/__tests__/AnnotationNode.test.tsx b/src/components/__tests__/AnnotationNode.test.tsx index 14fd2e41..7b687409 100644 --- a/src/components/__tests__/AnnotationNode.test.tsx +++ b/src/components/__tests__/AnnotationNode.test.tsx @@ -1,5 +1,6 @@ -import { describe, it, expect, vi, beforeEach, afterEach } from "vitest"; +import { describe, it, expect, vi, beforeEach, afterEach, beforeAll, afterAll } from "vitest"; import { render, screen, fireEvent, waitFor } from "@testing-library/react"; +import type { ReactNode } from "react"; import { AnnotationNode } from "@/components/nodes/AnnotationNode"; import { ReactFlowProvider } from "@xyflow/react"; import { AnnotationNodeData } from "@/types"; @@ -32,7 +33,6 @@ vi.mock("@/store/annotationStore", () => ({ // Mock alert const mockAlert = vi.fn(); -global.alert = mockAlert; // Mock DataTransfer class MockDataTransfer { @@ -50,14 +50,22 @@ class MockDataTransfer { }; } } -global.DataTransfer = MockDataTransfer as unknown as typeof DataTransfer; // Wrapper component for React Flow context -function TestWrapper({ children }: { children: React.ReactNode }) { +function TestWrapper({ children }: { children: ReactNode }) { return {children}; } describe("AnnotationNode", () => { + beforeAll(() => { + vi.stubGlobal("alert", mockAlert); + vi.stubGlobal("DataTransfer", MockDataTransfer); + }); + + afterAll(() => { + vi.unstubAllGlobals(); + }); + beforeEach(() => { vi.clearAllMocks(); @@ -372,7 +380,7 @@ describe("AnnotationNode", () => { }); it("should process valid image file and call updateNodeData", async () => { - // Mock FileReader + // Mock FileReader using vi.stubGlobal for proper cleanup const mockReadAsDataURL = vi.fn(); class MockFileReader { onload: ((event: ProgressEvent) => void) | null = null; @@ -384,7 +392,7 @@ describe("AnnotationNode", () => { }, 0); } } - global.FileReader = MockFileReader as unknown as typeof FileReader; + vi.stubGlobal("FileReader", MockFileReader); render( @@ -405,6 +413,12 @@ describe("AnnotationNode", () => { annotations: [], }); }); + + // Restore FileReader + vi.unstubAllGlobals(); + // Re-stub the globals we need for other tests + vi.stubGlobal("alert", mockAlert); + vi.stubGlobal("DataTransfer", MockDataTransfer); }); it("should reject non-image file types", () => { @@ -452,13 +466,9 @@ describe("AnnotationNode", () => { ); const dropZone = screen.getByText("Drop, click, or connect").parentElement!; - const dragOverEvent = new Event("dragover", { bubbles: true }); - Object.assign(dragOverEvent, { - preventDefault: vi.fn(), - stopPropagation: vi.fn(), - }); - fireEvent(dropZone, dragOverEvent); + // Use fireEvent.dragOver for idiomatic testing + fireEvent.dragOver(dropZone); // Should handle without error expect(dropZone).toBeInTheDocument(); diff --git a/src/components/__tests__/CostDialog.test.tsx b/src/components/__tests__/CostDialog.test.tsx index 35faba17..bbf48ae2 100644 --- a/src/components/__tests__/CostDialog.test.tsx +++ b/src/components/__tests__/CostDialog.test.tsx @@ -1,4 +1,4 @@ -import { describe, it, expect, vi, beforeEach, afterEach } from "vitest"; +import { describe, it, expect, vi, beforeEach, afterEach, beforeAll, afterAll } from "vitest"; import { render, screen, fireEvent } from "@testing-library/react"; import { CostDialog } from "@/components/CostDialog"; import { PredictedCostResult } from "@/utils/costCalculator"; @@ -18,11 +18,19 @@ vi.mock("@/store/workflowStore", () => ({ // Mock confirm const mockConfirm = vi.fn(() => true); -global.confirm = mockConfirm; describe("CostDialog", () => { + beforeAll(() => { + vi.stubGlobal("confirm", mockConfirm); + }); + + afterAll(() => { + vi.unstubAllGlobals(); + }); + beforeEach(() => { vi.clearAllMocks(); + mockConfirm.mockReturnValue(true); // Reset default return value mockUseWorkflowStore.mockImplementation((selector) => { const state = { resetIncurredCost: mockResetIncurredCost, diff --git a/src/components/__tests__/EdgeToolbar.test.tsx b/src/components/__tests__/EdgeToolbar.test.tsx index 1ea571b9..9720da29 100644 --- a/src/components/__tests__/EdgeToolbar.test.tsx +++ b/src/components/__tests__/EdgeToolbar.test.tsx @@ -1,5 +1,6 @@ import { describe, it, expect, vi, beforeEach, afterEach } from "vitest"; import { render, screen, fireEvent } from "@testing-library/react"; +import type { ReactNode } from "react"; import { EdgeToolbar } from "@/components/EdgeToolbar"; import { ReactFlowProvider } from "@xyflow/react"; @@ -18,7 +19,7 @@ vi.mock("@/store/workflowStore", () => ({ })); // Wrapper component for React Flow context -function TestWrapper({ children }: { children: React.ReactNode }) { +function TestWrapper({ children }: { children: ReactNode }) { return {children}; }