Browse Source

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
handoff-20260429-1057
Claude 6 months ago
parent
commit
7c3b32ae58
Failed to extract signature
  1. 9
      src/app/api/providers/fal/models/route.ts
  2. 6
      src/app/api/save-generation/__tests__/route.test.ts
  3. 52
      src/app/api/save-generation/route.ts
  4. 13
      src/app/globals.css
  5. 33
      src/components/__tests__/AnnotationModal.test.tsx
  6. 34
      src/components/__tests__/AnnotationNode.test.tsx
  7. 12
      src/components/__tests__/CostDialog.test.tsx
  8. 3
      src/components/__tests__/EdgeToolbar.test.tsx

9
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<NextResponse<ModelsResponse>> {
// 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");

6
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 () => {

52
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);
// 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;
}
} else {
// Handle base64 data URL
const dataUrlMatch = content.match(/^data:([\w/+-]+);base64,/);

13
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;

33
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 }) => <div data-testid="konva-stage">{children}</div>,
Layer: ({ children }: { children: React.ReactNode }) => <div data-testid="konva-layer">{children}</div>,
Stage: ({ children }: { children: ReactNode }) => <div data-testid="konva-stage">{children}</div>,
Layer: ({ children }: { children: ReactNode }) => <div data-testid="konva-layer">{children}</div>,
Image: () => <div data-testid="konva-image" />,
Rect: () => <div data-testid="konva-rect" />,
Ellipse: () => <div data-testid="konva-ellipse" />,
@ -74,7 +75,12 @@ const createMockAnnotationStore = (overrides = {}) => ({
let mockAnnotationStore = createMockAnnotationStore();
vi.mock("@/store/annotationStore", () => ({
useAnnotationStore: () => mockAnnotationStore,
useAnnotationStore: (selector?: (state: ReturnType<typeof createMockAnnotationStore>) => 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,10 +286,12 @@ describe("AnnotationModal", () => {
// so we just verify the button exists and triggers the handler
fireEvent.click(screen.getByText("Done"));
// closeModal should be called after save
// closeModal should be called after save (async operation)
await waitFor(() => {
expect(mockCloseModal).toHaveBeenCalled();
});
});
});
describe("Color Picker", () => {
it("should render color label", () => {

34
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 <ReactFlowProvider>{children}</ReactFlowProvider>;
}
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<FileReader>) => void) | null = null;
@ -384,7 +392,7 @@ describe("AnnotationNode", () => {
}, 0);
}
}
global.FileReader = MockFileReader as unknown as typeof FileReader;
vi.stubGlobal("FileReader", MockFileReader);
render(
<TestWrapper>
@ -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();

12
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,

3
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 <ReactFlowProvider>{children}</ReactFlowProvider>;
}

Loading…
Cancel
Save