diff --git a/storage-resize-images/functions/__tests__/content-filter.test.ts b/storage-resize-images/functions/__tests__/content-filter.test.ts index 1a92ee86b..f8abc7c34 100644 --- a/storage-resize-images/functions/__tests__/content-filter.test.ts +++ b/storage-resize-images/functions/__tests__/content-filter.test.ts @@ -18,15 +18,6 @@ jest.mock("@genkit-ai/vertexai", () => ({ gemini: jest.fn((version: string) => ({ name: `vertexai/${version}` })), })); -// Mock the sleep function to avoid actual waiting in tests -jest.mock("fs", () => ({ - readFileSync: jest.fn().mockReturnValue(Buffer.from("mockImageData")), -})); - -jest.mock("mime", () => ({ - lookup: jest.fn().mockReturnValue("image/png"), -})); - describe("checkImageContent with mocks", () => { // Test image path - using the same path as in your original test suite const imagePath = path.join(__dirname, "gun-image.png"); diff --git a/storage-resize-images/functions/__tests__/unit/generateResizedImageHandler.test.ts b/storage-resize-images/functions/__tests__/unit/generateResizedImageHandler.test.ts new file mode 100644 index 000000000..db255d02f --- /dev/null +++ b/storage-resize-images/functions/__tests__/unit/generateResizedImageHandler.test.ts @@ -0,0 +1,221 @@ +import * as path from "path"; +import * as fs from "fs"; +import { config as loadEnv } from "dotenv"; + +const envLocalPath = path.resolve( + __dirname, + "../../../../_emulator/extensions/storage-resize-images.env.local" +); + +loadEnv({ path: envLocalPath, debug: true, override: true }); + +jest.mock("fs", () => ({ + ...jest.requireActual("fs"), + copyFileSync: jest.fn(), +})); + +jest.mock("../../src/filters", () => ({ + shouldResize: jest.fn(), +})); + +jest.mock("../../src/file-operations", () => ({ + downloadOriginalFile: jest.fn(), + handleFailedImage: jest.fn(), + deleteTempFile: jest.fn().mockResolvedValue(undefined), + deleteRemoteFile: jest.fn().mockResolvedValue(undefined), +})); + +jest.mock("../../src/content-filter", () => ({ + checkImageContent: jest.fn(), +})); + +jest.mock("../../src/placeholder", () => ({ + replacePlaceholder: jest.fn().mockResolvedValue(undefined), +})); + +jest.mock("../../src/resize-image", () => ({ + resizeImages: jest.fn(), +})); + +jest.mock("../../src/events", () => ({ + setupEventChannel: jest.fn(), + recordStartResizeEvent: jest.fn().mockResolvedValue(undefined), + recordSuccessEvent: jest.fn().mockResolvedValue(undefined), + recordErrorEvent: jest.fn().mockResolvedValue(undefined), + recordStartEvent: jest.fn().mockResolvedValue(undefined), + recordCompletionEvent: jest.fn().mockResolvedValue(undefined), +})); + +jest.mock("../../src/logs", () => ({ + init: jest.fn(), + start: jest.fn(), + failed: jest.fn(), + complete: jest.fn(), + error: jest.fn(), + contentFilterErrored: jest.fn(), + contentFilterRejected: jest.fn(), + placeholderReplaceError: jest.fn(), +})); + +jest.mock("firebase-admin", () => ({ + initializeApp: jest.fn(), + storage: jest.fn(() => ({ + bucket: jest.fn(() => ({})), + })), +})); + +import { generateResizedImageHandler } from "../../src/index"; +import { shouldResize } from "../../src/filters"; +import { + downloadOriginalFile, + handleFailedImage, +} from "../../src/file-operations"; +import { checkImageContent } from "../../src/content-filter"; +import { replacePlaceholder } from "../../src/placeholder"; +import { resizeImages } from "../../src/resize-image"; +import * as logs from "../../src/logs"; +import exp from "constants"; + +describe("generateResizedImageHandler", () => { + beforeEach(() => { + jest.clearAllMocks(); + }); + + const mockObject = { + bucket: "demo-bucket", + name: "images/test.jpg", + contentType: "image/jpeg", + } as any; + + const parsedPathMatcher = expect.objectContaining({ + dir: "images", + base: "test.jpg", + name: "test", + ext: ".jpg", + }); + + test("routes blocked-by-filter images to the failed-image path with blockedByFilter=true", async () => { + (shouldResize as jest.Mock).mockReturnValue(true); + (downloadOriginalFile as jest.Mock).mockResolvedValue([ + "/tmp/test.jpg", + {}, + ]); + (checkImageContent as jest.Mock).mockResolvedValue(false); + (resizeImages as jest.Mock).mockResolvedValue([ + { + status: "fulfilled", + value: { success: true }, + }, + ]); + + await generateResizedImageHandler(mockObject, false); + + expect(handleFailedImage).toHaveBeenCalledWith( + expect.anything(), + "/tmp/test.jpg", + mockObject, + parsedPathMatcher, + true + ); + expect(handleFailedImage).toHaveBeenCalledTimes(1); + expect(fs.copyFileSync).toHaveBeenCalledWith( + "/tmp/test.jpg", + "/tmp/test.jpg-placeholder" + ); + expect(replacePlaceholder).toHaveBeenCalledWith( + "/tmp/test.jpg-placeholder", + {}, + null + ); + expect(resizeImages).toHaveBeenCalledWith( + expect.anything(), + "/tmp/test.jpg-placeholder", + parsedPathMatcher, + mockObject + ); + }); + + test("resizes when the content filter passes", async () => { + (shouldResize as jest.Mock).mockReturnValue(true); + (downloadOriginalFile as jest.Mock).mockResolvedValue([ + "/tmp/test.jpg", + {}, + ]); + (checkImageContent as jest.Mock).mockResolvedValue(true); + (resizeImages as jest.Mock).mockResolvedValue([ + { + status: "fulfilled", + value: { success: true }, + }, + ]); + + await generateResizedImageHandler(mockObject, false); + + expect(replacePlaceholder).not.toHaveBeenCalled(); + expect(resizeImages).toHaveBeenCalledWith( + expect.anything(), + "/tmp/test.jpg", + parsedPathMatcher, + mockObject + ); + expect(handleFailedImage).not.toHaveBeenCalled(); + }); + + test("treats filter errors as failures and skips resizing", async () => { + (shouldResize as jest.Mock).mockReturnValue(true); + (downloadOriginalFile as jest.Mock).mockResolvedValue([ + "/tmp/test.jpg", + {}, + ]); + (checkImageContent as jest.Mock).mockRejectedValue( + new Error("filter boom") + ); + + await generateResizedImageHandler(mockObject, false); + + expect(replacePlaceholder).not.toHaveBeenCalled(); + expect(resizeImages).not.toHaveBeenCalled(); + expect(handleFailedImage).toHaveBeenCalledWith( + expect.anything(), + "/tmp/test.jpg", + mockObject, + parsedPathMatcher, + false + ); + }); + + test("still routes blocked images to the failed path when placeholder swap errors", async () => { + (shouldResize as jest.Mock).mockReturnValue(true); + (downloadOriginalFile as jest.Mock).mockResolvedValue([ + "/tmp/test.jpg", + {}, + ]); + (checkImageContent as jest.Mock).mockResolvedValue(false); + + const swapErr = new Error("swap boom"); + (replacePlaceholder as jest.Mock).mockRejectedValue(swapErr); + + await generateResizedImageHandler(mockObject, false); + + expect(handleFailedImage).toHaveBeenCalledWith( + expect.anything(), + "/tmp/test.jpg", + mockObject, + parsedPathMatcher, + true + ); + expect(handleFailedImage).toHaveBeenCalledTimes(1); + expect(fs.copyFileSync).toHaveBeenCalledWith( + "/tmp/test.jpg", + "/tmp/test.jpg-placeholder" + ); + expect(replacePlaceholder).toHaveBeenCalledWith( + "/tmp/test.jpg-placeholder", + {}, + null + ); + expect(logs.placeholderReplaceError).toHaveBeenCalledWith(swapErr); + expect(logs.contentFilterErrored).not.toHaveBeenCalled(); + expect(resizeImages).not.toHaveBeenCalled(); + }); +}); diff --git a/storage-resize-images/functions/src/content-filter.ts b/storage-resize-images/functions/src/content-filter.ts index c16b0ae98..ab550aa02 100644 --- a/storage-resize-images/functions/src/content-filter.ts +++ b/storage-resize-images/functions/src/content-filter.ts @@ -2,36 +2,18 @@ import vertexAI, { gemini } from "@genkit-ai/vertexai"; import { genkit, z } from "genkit"; import type { SafetyThreshold } from "./config"; import * as fs from "fs"; -import * as path from "path"; -import { Bucket } from "@google-cloud/storage"; -import { ObjectMetadata } from "firebase-functions/v1/storage"; -import { - replaceWithConfiguredPlaceholder, - replaceWithDefaultPlaceholder, -} from "./util"; -// Import the logging functions from your log.ts module import * as log from "./logs"; import { globalRetryQueue } from "./global"; /** * Creates a data URL from an image file - * @param filePath Path to the image file + * @param imageBuffer Raw image file contents + * @param contentType MIME type for the image, for example "image/jpeg" * @returns Data URL string */ -function createImageDataUrl(filePath: string): string { - const imageBuffer = fs.readFileSync(filePath); +function createImageDataUrl(imageBuffer: Buffer, contentType: string): string { const base64Image = imageBuffer.toString("base64"); - const mimeType = getMimeType(filePath); - return `data:${mimeType};base64,${base64Image}`; -} - -/** - * Determines MIME type based on file extension - * @param filePath Path to the file - * @returns MIME type string - */ -function getMimeType(filePath: string): string { - return path.extname(filePath).toLowerCase(); + return `data:${contentType};base64,${base64Image}`; } /** @@ -43,20 +25,12 @@ function sleep(ms: number): Promise { return new Promise((resolve) => setTimeout(resolve, ms)); } -type SafetyCategory = - | "HARM_CATEGORY_UNSPECIFIED" - | "HARM_CATEGORY_HATE_SPEECH" - | "HARM_CATEGORY_DANGEROUS_CONTENT" - | "HARM_CATEGORY_HARASSMENT" - | "HARM_CATEGORY_SEXUALLY_EXPLICIT"; - -const HARM_CATEGORIES: ReadonlyArray = [ +const HARM_CATEGORIES = [ "HARM_CATEGORY_HATE_SPEECH", "HARM_CATEGORY_DANGEROUS_CONTENT", - "HARM_CATEGORY_UNSPECIFIED", "HARM_CATEGORY_SEXUALLY_EXPLICIT", "HARM_CATEGORY_HARASSMENT", -]; +] as const; /** * Creates safety settings based on filter level @@ -67,26 +41,34 @@ function createSafetySettings(filterLevel: SafetyThreshold) { return HARM_CATEGORIES.map((category) => ({ category, threshold: filterLevel, - })); + })) as any; } /** - * Performs the actual content check with the AI model + * Entry point for image moderation: short-circuits when disabled, otherwise runs a single + * Vertex/Gemini call per attempt with retries and queue-backed backoff on transient errors. + * * @param localOriginalFile Path to the local image file - * @param filterLevel The content filter level to apply + * @param filterLevel The content filter level to apply ('LOW', 'MEDIUM', 'HIGH', or null to disable) * @param prompt Optional custom prompt to use for content checking * @param contentType The content type of the image + * @param maxAttempts Maximum number of retry attempts in case of errors * @returns Promise - true if the image passes the filter, false otherwise */ -async function performContentCheck( +export async function checkImageContent( localOriginalFile: string, filterLevel: SafetyThreshold | null, prompt: string | null, - contentType: string + contentType: string, + maxAttempts = 3 ): Promise { - const dataUrl = createImageDataUrl(localOriginalFile); + if (filterLevel === null && prompt === null) { + return true; + } + + const imageBuffer = fs.readFileSync(localOriginalFile); + const dataUrl = createImageDataUrl(imageBuffer, contentType); - // Initialize Vertex AI client const ai = genkit({ plugins: [ vertexAI({ @@ -96,17 +78,19 @@ async function performContentCheck( ], }); - // Determine the effective safety settings and prompt to use - const effectiveFilterLevel: SafetyThreshold = - filterLevel === null ? "BLOCK_NONE" : filterLevel; - const effectivePrompt = - prompt !== null + /** One Gemini moderation call (no retries). */ + async function moderateImageOnce(): Promise { + const effectiveFilterLevel: SafetyThreshold = + filterLevel === null ? "BLOCK_NONE" : filterLevel; + + const hasCustomPrompt = prompt !== null; + + const effectivePrompt = hasCustomPrompt ? prompt + '\n\n Please respond in json with either { "response": "yes" } or { "response": "no" }' : "Is this image appropriate?"; - const effectiveOutput = - prompt !== null + const effectiveOutput = hasCustomPrompt ? { format: "json", schema: z.object({ @@ -115,80 +99,53 @@ async function performContentCheck( } : undefined; - // Determine max tokens based on whether we're using custom prompt - const maxOutputTokens = prompt !== null ? 100 : 1; + const maxOutputTokens = hasCustomPrompt ? 100 : 1; - try { - const result = await ai.generate({ - model: gemini("gemini-2.5-flash"), - messages: [ - { - role: "user", - content: [ - { - text: effectivePrompt, - }, - { - media: { - url: dataUrl, - contentType, + try { + const result = await ai.generate({ + model: gemini("gemini-2.5-flash"), + messages: [ + { + role: "user", + content: [ + { + text: effectivePrompt, + }, + { + media: { + url: dataUrl, + contentType, + }, }, - }, - ], + ], + }, + ], + output: effectiveOutput, + config: { + temperature: 0.1, + maxOutputTokens, + safetySettings: createSafetySettings(effectiveFilterLevel), }, - ], - output: effectiveOutput, - config: { - temperature: 0.1, - maxOutputTokens, - safetySettings: createSafetySettings(effectiveFilterLevel), - }, - }); + }); - if (result.output?.response === "yes" && prompt !== null) { - log.customFilterBlocked(); - return false; - } + if (result.output?.response === "yes" && hasCustomPrompt) { + log.customFilterBlocked(); + return false; + } - return true; - } catch (error) { - if (error.detail?.response?.finishReason === "blocked") { - log.contentFilterBlocked(); - return false; + return true; + } catch (error) { + if (error.detail?.response?.finishReason === "blocked") { + log.contentFilterBlocked(); + return false; + } + throw error; } - throw error; - } -} - -/** - * Checks if an image content is appropriate based on the provided filter level and optional custom prompt - * @param localOriginalFile Path to the local image file - * @param filterLevel The content filter level to apply ('LOW', 'MEDIUM', 'HIGH', or null to disable) - * @param prompt Optional custom prompt to use for content checking - * @param contentType The content type of the image - * @param maxAttempts Maximum number of retry attempts in case of errors - * @returns Promise - true if the image passes the filter, false otherwise - */ -export async function checkImageContent( - localOriginalFile: string, - filterLevel: SafetyThreshold | null, - prompt: string | null, - contentType: string, - maxAttempts = 3 -): Promise { - if (filterLevel === null && prompt === null) { - return true; } - // Helper function that handles retries using the queue const attemptWithRetry = async (attemptNumber: number): Promise => { try { - return await performContentCheck( - localOriginalFile, - filterLevel, - prompt, - contentType - ); + return await moderateImageOnce(); } catch (error) { // If we have attempts left, schedule a retry via the queue if (attemptNumber < maxAttempts) { @@ -219,53 +176,3 @@ export async function checkImageContent( // Start the first attempt (not via queue) return await attemptWithRetry(1); } - -/** - * Processes content filtering and handles placeholder replacement if needed - */ -export async function processContentFilter( - localFile: string, - object: ObjectMetadata, - bucket: Bucket, - _verbose: boolean, - config: any -): Promise<{ passed: boolean; failed: boolean | null }> { - let filterResult = true; // Default to true (pass) - let failed = null; // No failures yet - - try { - filterResult = await checkImageContent( - localFile, - config.contentFilterLevel, - config.customFilterPrompt, - object.contentType - ); - } catch (err) { - log.contentFilterErrored(err); - failed = true; - } - - if (filterResult === false) { - log.contentFilterRejected(object.name); - - try { - if (config.placeholderImagePath) { - log.replacingWithConfiguredPlaceholder(config.placeholderImagePath); - await replaceWithConfiguredPlaceholder( - localFile, - bucket, - config.placeholderImagePath - ); - } else { - log.replacingWithDefaultPlaceholder(); - await replaceWithDefaultPlaceholder(localFile); - } - log.placeholderReplaceComplete(localFile); - } catch (err) { - log.placeholderReplaceError(err); - failed = true; - } - } - - return { passed: filterResult, failed }; -} diff --git a/storage-resize-images/functions/src/index.ts b/storage-resize-images/functions/src/index.ts index 6a4a380cb..206947a0c 100644 --- a/storage-resize-images/functions/src/index.ts +++ b/storage-resize-images/functions/src/index.ts @@ -22,6 +22,7 @@ import * as path from "path"; import * as sharp from "sharp"; import { File } from "@google-cloud/storage"; import { ObjectMetadata } from "firebase-functions/v1/storage"; +import * as fs from "fs"; import { resizeImages } from "./resize-image"; import { config, deleteImage } from "./config"; @@ -29,7 +30,8 @@ import * as logs from "./logs"; import { shouldResize } from "./filters"; import * as events from "./events"; import { convertToObjectMetadata } from "./util"; -import { processContentFilter } from "./content-filter"; +import { checkImageContent } from "./content-filter"; +import { replacePlaceholder } from "./placeholder"; import { deleteRemoteFile, deleteTempFile, @@ -50,7 +52,7 @@ logs.init(config); * When an image is uploaded in the Storage bucket, we generate a resized image automatically using * the Sharp image converting library. */ -const generateResizedImageHandler = async ( +export const generateResizedImageHandler = async ( object: ObjectMetadata, verbose = true ): Promise => { @@ -67,9 +69,9 @@ const generateResizedImageHandler = async ( const bucket = admin.storage().bucket(object.bucket); const filePath = object.name; // File path in the bucket. const parsedPath = path.parse(filePath); - const objectMetadata = object; - let failed = null; + let localOriginalFile: string; + let localProcessingFile: string | undefined; let remoteOriginalFile: File; try { @@ -79,22 +81,57 @@ const generateResizedImageHandler = async ( verbose ); - // Check content filter and replace with placeholder if needed - const filterResult = await processContentFilter( - localOriginalFile, - object, - bucket, - verbose, - config - ); + let blockedByFilter = false; + let filterErrored = false; + let blockedImageStored = false; + + try { + const passed = await checkImageContent( + localOriginalFile, + config.contentFilterLevel, + config.customFilterPrompt, + object.contentType + ); + if (!passed) { + blockedByFilter = true; + logs.contentFilterRejected(object.name); + + await handleFailedImage( + bucket, + localOriginalFile, + object, + parsedPath, + true + ); + blockedImageStored = true; + + localProcessingFile = `${localOriginalFile}-placeholder`; + fs.copyFileSync(localOriginalFile, localProcessingFile); + try { + await replacePlaceholder( + localProcessingFile, + bucket, + config.placeholderImagePath + ); + } catch (err) { + logs.placeholderReplaceError(err); + filterErrored = true; + } + } + } catch (err) { + logs.contentFilterErrored(err); + filterErrored = true; + } + + const fileToResize = localProcessingFile ?? localOriginalFile; - // Process image resizing if content filter didn't fail - if (filterResult.failed !== true) { + let resizeFailed = false; + if (!filterErrored) { const resizeResults = await resizeImages( bucket, - localOriginalFile, + fileToResize, parsedPath, - objectMetadata + object ); await events.recordSuccessEvent({ @@ -102,31 +139,29 @@ const generateResizedImageHandler = async ( data: { input: object, outputs: resizeResults, - contentFilterPassed: filterResult.passed, + contentFilterPassed: !blockedByFilter, }, }); - // Only update failed status if it's still null (not already failed from content filter) - failed = - filterResult.failed === null - ? resizeResults.some( - (result) => - result.status === "rejected" || result.value.success === false - ) - : filterResult.failed; - } else { - failed = true; + resizeFailed = resizeResults.some( + (result) => + result.status === "rejected" || result.value.success === false + ); } + const failed = filterErrored || resizeFailed; + if (failed) { logs.failed(); - await handleFailedImage( - bucket, - localOriginalFile, - object, - parsedPath, - filterResult.passed === false - ); + if (!blockedImageStored) { + await handleFailedImage( + bucket, + localOriginalFile, + object, + parsedPath, + blockedByFilter + ); + } } else { if (config.deleteOriginalFile === deleteImage.onSuccess) { await deleteRemoteFile(remoteOriginalFile, filePath); @@ -142,6 +177,10 @@ const generateResizedImageHandler = async ( await deleteTempFile(localOriginalFile, filePath, verbose); } + if (localProcessingFile) { + await deleteTempFile(localProcessingFile, filePath, verbose); + } + if ( config.deleteOriginalFile === deleteImage.always && remoteOriginalFile diff --git a/storage-resize-images/functions/src/placeholder.ts b/storage-resize-images/functions/src/placeholder.ts new file mode 100644 index 000000000..9e298b009 --- /dev/null +++ b/storage-resize-images/functions/src/placeholder.ts @@ -0,0 +1,31 @@ +import { Bucket } from "@google-cloud/storage"; + +import * as log from "./logs"; +import { + replaceWithConfiguredPlaceholder, + replaceWithDefaultPlaceholder, +} from "./util"; + +/** + * Swaps the local file with a placeholder image. Uses the configured + * placeholder at `placeholderImagePath` when provided, otherwise the bundled + * default. + */ +export async function replacePlaceholder( + localFile: string, + bucket: Bucket, + placeholderImagePath: string | null +): Promise { + if (placeholderImagePath) { + log.replacingWithConfiguredPlaceholder(placeholderImagePath); + await replaceWithConfiguredPlaceholder( + localFile, + bucket, + placeholderImagePath + ); + } else { + log.replacingWithDefaultPlaceholder(); + await replaceWithDefaultPlaceholder(localFile); + } + log.placeholderReplaceComplete(localFile); +}