promptfoo--promptfoo/test/blobs/extractor.test.ts

import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
import {
  extractAndStoreBinaryData,
  isBlobStorageEnabled,
  normalizeAudioMimeType,
} from '../../src/blobs/extractor';
import { sha256 } from '../../src/util/createHash';

import type { ProviderResponse } from '../../src/types/providers';

// Mock the remoteUpload module
vi.mock('../../src/blobs/remoteUpload', () => ({
  uploadBlobRemote: vi.fn(),
}));

// Mock the blob index module
vi.mock('../../src/blobs/index', () => ({
  storeBlob: vi.fn().mockResolvedValue({
    ref: {
      uri: 'promptfoo://blob/abc123',
      hash: 'abc123',
    },
  }),
  recordBlobReference: vi.fn().mockResolvedValue(undefined),
}));

describe('normalizeAudioMimeType', () => {
  describe('undefined and empty inputs', () => {
    it('should return audio/wav for undefined format', () => {
      expect(normalizeAudioMimeType(undefined)).toBe('audio/wav');
    });

    it('should return audio/wav for empty string', () => {
      expect(normalizeAudioMimeType('')).toBe('audio/wav');
    });

    it('should return audio/wav for whitespace-only string', () => {
      expect(normalizeAudioMimeType('   ')).toBe('audio/wav');
    });
  });

  describe('already valid MIME types', () => {
    it('should pass through valid audio MIME types unchanged', () => {
      expect(normalizeAudioMimeType('audio/wav')).toBe('audio/wav');
      expect(normalizeAudioMimeType('audio/mpeg')).toBe('audio/mpeg');
      expect(normalizeAudioMimeType('audio/ogg')).toBe('audio/ogg');
      expect(normalizeAudioMimeType('audio/flac')).toBe('audio/flac');
      expect(normalizeAudioMimeType('audio/aac')).toBe('audio/aac');
      expect(normalizeAudioMimeType('audio/mp4')).toBe('audio/mp4');
      expect(normalizeAudioMimeType('audio/webm')).toBe('audio/webm');
    });

    it('should pass through vendor MIME types without periods', () => {
      expect(normalizeAudioMimeType('audio/x-custom')).toBe('audio/x-custom');
      expect(normalizeAudioMimeType('audio/x-wav')).toBe('audio/x-wav');
    });

    it('should reject MIME types with periods to prevent injection attacks', () => {
      // Periods could enable attacks like "audio/wav.html" being interpreted as HTML
      expect(normalizeAudioMimeType('audio/vnd.company.format')).toBe('audio/wav');
      expect(normalizeAudioMimeType('audio/wav.html')).toBe('audio/wav');
    });
  });

  describe('short format normalization', () => {
    it('should normalize common short formats to MIME types', () => {
      expect(normalizeAudioMimeType('wav')).toBe('audio/wav');
      expect(normalizeAudioMimeType('mp3')).toBe('audio/mpeg');
      expect(normalizeAudioMimeType('ogg')).toBe('audio/ogg');
      expect(normalizeAudioMimeType('flac')).toBe('audio/flac');
      expect(normalizeAudioMimeType('aac')).toBe('audio/aac');
      expect(normalizeAudioMimeType('m4a')).toBe('audio/mp4');
      expect(normalizeAudioMimeType('webm')).toBe('audio/webm');
    });

    it('should be case-insensitive for known formats', () => {
      expect(normalizeAudioMimeType('WAV')).toBe('audio/wav');
      expect(normalizeAudioMimeType('Mp3')).toBe('audio/mpeg');
      expect(normalizeAudioMimeType('OGG')).toBe('audio/ogg');
      expect(normalizeAudioMimeType('FLAC')).toBe('audio/flac');
    });

    it('should handle whitespace around format', () => {
      expect(normalizeAudioMimeType('  wav  ')).toBe('audio/wav');
      expect(normalizeAudioMimeType('\tmp3\n')).toBe('audio/mpeg');
    });
  });

  describe('unknown formats', () => {
    it('should prefix unknown alphanumeric formats with audio/', () => {
      expect(normalizeAudioMimeType('opus')).toBe('audio/opus');
      expect(normalizeAudioMimeType('pcm')).toBe('audio/pcm');
      expect(normalizeAudioMimeType('raw')).toBe('audio/raw');
    });

    it('should allow dash and underscore in unknown formats', () => {
      expect(normalizeAudioMimeType('x-custom')).toBe('audio/x-custom');
      expect(normalizeAudioMimeType('codec_v2')).toBe('audio/codec_v2');
      expect(normalizeAudioMimeType('x-wav-hd')).toBe('audio/x-wav-hd');
    });

    it('should reject unknown formats with periods to prevent MIME injection', () => {
      // Periods could enable attacks like "wav.html" -> "audio/wav.html"
      expect(normalizeAudioMimeType('vnd.company.format')).toBe('audio/wav');
      expect(normalizeAudioMimeType('format-1.0')).toBe('audio/wav');
      expect(normalizeAudioMimeType('wav.html')).toBe('audio/wav');
    });
  });

  describe('security: invalid formats', () => {
    it('should reject formats with path traversal attempts', () => {
      expect(normalizeAudioMimeType('../../etc/passwd')).toBe('audio/wav');
      expect(normalizeAudioMimeType('../../../etc/shadow')).toBe('audio/wav');
      expect(normalizeAudioMimeType('path/to/file')).toBe('audio/wav');
    });

    it('should reject formats with newline injection', () => {
      expect(normalizeAudioMimeType('audio\nwav')).toBe('audio/wav');
      expect(normalizeAudioMimeType('wav\r\nContent-Type: text/html')).toBe('audio/wav');
    });

    it('should reject formats with header injection characters', () => {
      expect(normalizeAudioMimeType('audio;wav')).toBe('audio/wav');
      expect(normalizeAudioMimeType('audio:wav')).toBe('audio/wav');
      expect(normalizeAudioMimeType('audio wav')).toBe('audio/wav');
    });

    it('should reject formats with backslashes', () => {
      expect(normalizeAudioMimeType('audio\\wav')).toBe('audio/wav');
      expect(normalizeAudioMimeType('..\\..\\etc\\passwd')).toBe('audio/wav');
    });

    it('should reject formats with multiple path components', () => {
      expect(normalizeAudioMimeType('audio/wav/extra')).toBe('audio/wav');
    });
  });
});

describe('Audio MIME type normalization (integration)', () => {
  // Skip tests if blob storage is disabled
  const maybeIt = isBlobStorageEnabled() ? it : it.skip;

  describe('extractAndStoreBinaryData with audio', () => {
    maybeIt('should process response with top-level audio', async () => {
      const response: ProviderResponse = {
        output: 'test',
        audio: {
          data: 'SGVsbG8gV29ybGQ=',
          format: 'wav',
        },
      };

      const result = await extractAndStoreBinaryData(response);
      expect(result).toBeDefined();
      expect(result?.audio).toBeDefined();
    });

    maybeIt('should process response with turns audio', async () => {
      const response: ProviderResponse = {
        output: 'test',
        turns: [
          {
            audio: {
              data: 'SGVsbG8gV29ybGQ=',
              format: 'mp3',
            },
          },
        ],
      } as any;

      const result = await extractAndStoreBinaryData(response);
      expect(result).toBeDefined();
      expect((result as any)?.turns).toBeDefined();
    });

    maybeIt('should handle response without audio', async () => {
      const response: ProviderResponse = {
        output: 'test',
      };

      const result = await extractAndStoreBinaryData(response);
      expect(result).toEqual(response);
    });

    maybeIt('should handle null response', async () => {
      const result = await extractAndStoreBinaryData(null);
      expect(result).toBeNull();
    });

    maybeIt('should handle undefined response', async () => {
      const result = await extractAndStoreBinaryData(undefined);
      expect(result).toBeUndefined();
    });
  });
});

describe('Local blob extraction', () => {
  let mockUploadBlobRemote: ReturnType<typeof vi.fn>;
  let mockStoreBlob: ReturnType<typeof vi.fn>;

  beforeEach(async () => {
    vi.resetAllMocks();

    // Get the mocked functions
    const remoteUploadModule = await import('../../src/blobs/remoteUpload');
    mockUploadBlobRemote = vi.mocked(remoteUploadModule.uploadBlobRemote);

    const blobIndexModule = await import('../../src/blobs/index');
    mockStoreBlob = vi.mocked(blobIndexModule.storeBlob);

    // Default mock implementations
    mockStoreBlob.mockResolvedValue({
      ref: {
        uri: 'promptfoo://blob/abc123def456',
        hash: 'abc123def456',
      },
    });
  });

  afterEach(() => {
    vi.resetAllMocks();
  });

  it('stores extracted blobs locally without eagerly uploading them', async () => {
    // Create a large enough data URL to trigger externalization (>1KB)
    const largeBase64 = Buffer.alloc(2000).toString('base64');
    const response: ProviderResponse = {
      output: `data:image/png;base64,${largeBase64}`,
    };

    await extractAndStoreBinaryData(response);

    // Should store locally
    expect(mockStoreBlob).toHaveBeenCalledTimes(1);
    expect(mockUploadBlobRemote).not.toHaveBeenCalled();
  });

  it('should externalize image data URIs to blobRefs', async () => {
    const largeBase64 = Buffer.alloc(2000).toString('base64');
    const response: ProviderResponse = {
      output: 'text output',
      images: [{ data: `data:image/png;base64,${largeBase64}`, mimeType: 'image/png' }],
    };

    const result = await extractAndStoreBinaryData(response);
    expect(result?.images?.[0].data).toBeUndefined();
    expect(result?.images?.[0].blobRef).toBeDefined();
    expect(result?.images?.[0].blobRef?.uri).toContain('promptfoo://blob/');
  });

  it('should reuse the same image blob when output and images contain identical data URIs', async () => {
    const largeBase64 = Buffer.alloc(2000).toString('base64');
    const dataUri = `data:image/png;base64,${largeBase64}`;
    const response: ProviderResponse = {
      output: dataUri,
      images: [{ data: dataUri, mimeType: 'image/png' }],
    };

    const result = await extractAndStoreBinaryData(response);

    expect(result?.output).toBe(result?.images?.[0].blobRef?.uri);
    expect(result?.images?.[0].data).toBeUndefined();
    expect(mockStoreBlob).toHaveBeenCalledTimes(1);
    expect(mockStoreBlob).toHaveBeenCalledWith(
      expect.any(Buffer),
      'image/png',
      expect.objectContaining({
        location: 'response.images[0].data',
        kind: 'image',
      }),
    );
  });

  it('should reuse the same image blob for concurrent identical images[] siblings', async () => {
    const largeBase64 = Buffer.alloc(2000).toString('base64');
    const dataUri = `data:image/png;base64,${largeBase64}`;
    const response: ProviderResponse = {
      output: 'text output',
      images: [
        { data: dataUri, mimeType: 'image/png' },
        { data: dataUri, mimeType: 'image/png' },
      ],
    };

    const result = await extractAndStoreBinaryData(response);

    expect(result?.images?.[0].data).toBeUndefined();
    expect(result?.images?.[1].data).toBeUndefined();
    expect(result?.images?.[0].blobRef).toEqual(result?.images?.[1].blobRef);
    expect(mockStoreBlob).toHaveBeenCalledTimes(1);
  });

  it('should pass through images without data URIs unchanged', async () => {
    const response: ProviderResponse = {
      output: 'text output',
      images: [
        {
          blobRef: {
            uri: 'promptfoo://blob/existing',
            hash: 'existing',
            mimeType: 'image/png',
            sizeBytes: 100,
            provider: 'local',
          },
        },
      ],
    };

    const result = await extractAndStoreBinaryData(response);
    expect(result).toEqual(response);
  });

  it('should handle mixed images: externalize data URIs but keep existing blobRefs', async () => {
    const largeBase64 = Buffer.alloc(2000).toString('base64');
    const response: ProviderResponse = {
      output: 'text output',
      images: [
        { data: `data:image/png;base64,${largeBase64}`, mimeType: 'image/png' },
        {
          blobRef: {
            uri: 'promptfoo://blob/existing',
            hash: 'existing',
            mimeType: 'image/jpeg',
            sizeBytes: 100,
            provider: 'local',
          },
        },
      ],
    };

    const result = await extractAndStoreBinaryData(response);
    // First image externalized
    expect(result?.images?.[0].data).toBeUndefined();
    expect(result?.images?.[0].blobRef).toBeDefined();
    // Second image kept as-is
    expect(result?.images?.[1].blobRef?.hash).toBe('existing');
  });

  it('should record blob references embedded in output text', async () => {
    const blobIndexModule = await import('../../src/blobs/index');
    const mockRecordBlobReference = vi.mocked(blobIndexModule.recordBlobReference);
    const firstHash = 'a'.repeat(64);
    const secondHash = 'b'.repeat(64);
    const response: ProviderResponse = {
      output:
        `First promptfoo://blob/${firstHash} repeated promptfoo://blob/${firstHash} ` +
        `second promptfoo://blob/${secondHash}`,
    };

    await extractAndStoreBinaryData(response, {
      evalId: 'eval-mixed-output',
      testIdx: 1,
      promptIdx: 2,
    });

    expect(mockRecordBlobReference).toHaveBeenCalledTimes(2);
    expect(mockRecordBlobReference).toHaveBeenCalledWith(
      firstHash,
      expect.objectContaining({
        evalId: 'eval-mixed-output',
        testIdx: 1,
        promptIdx: 2,
        location: 'response.output',
      }),
    );
    expect(mockRecordBlobReference).toHaveBeenCalledWith(
      secondHash,
      expect.objectContaining({
        evalId: 'eval-mixed-output',
        testIdx: 1,
        promptIdx: 2,
        location: 'response.output',
      }),
    );
  });

  it('should reuse the top-level audio blob for mirrored realtime metadata audio data', async () => {
    const largeBase64 = Buffer.alloc(2000).toString('base64');
    const response: ProviderResponse = {
      output: 'test',
      audio: {
        data: largeBase64,
        format: 'wav',
        transcript: 'hello',
      },
      metadata: {
        audio: {
          data: largeBase64,
          format: 'wav',
          transcript: 'hello',
        },
      },
    };

    const result = await extractAndStoreBinaryData(response);

    expect(result?.audio?.data).toBeUndefined();
    expect(result?.audio?.blobRef).toBeDefined();
    expect(result?.metadata?.audio?.data).toBeUndefined();
    expect(result?.metadata?.audio?.blobRef).toEqual(result?.audio?.blobRef);
    expect(mockStoreBlob).toHaveBeenCalledTimes(1);
    expect(mockStoreBlob).toHaveBeenCalledWith(
      expect.any(Buffer),
      'audio/wav',
      expect.objectContaining({
        location: 'response.audio.data',
        kind: 'audio',
      }),
    );
  });

  it('should externalize metadata-only audio data', async () => {
    const largeBase64 = Buffer.alloc(2000).toString('base64');
    const response: ProviderResponse = {
      output: 'test',
      metadata: {
        audio: {
          data: largeBase64,
          format: 'wav',
        },
      },
    };

    const result = await extractAndStoreBinaryData(response);

    expect(result?.metadata?.audio?.data).toBeUndefined();
    expect(result?.metadata?.audio?.blobRef).toBeDefined();
    expect(mockStoreBlob).toHaveBeenCalledTimes(1);
    expect(mockStoreBlob).toHaveBeenCalledWith(
      expect.any(Buffer),
      'audio/wav',
      expect.objectContaining({
        location: 'response.metadata.audio.data',
        kind: 'audio',
      }),
    );
  });

  it('should reuse the same audio blob across turns[].audio and metadata.audio', async () => {
    // Realtime providers can mirror an audio chunk into both turns[N].audio
    // and metadata.audio. Without a shared cache, the metadata path stored a
    // separate blob.

    const largeBase64 = Buffer.alloc(2000).toString('base64');
    const response = {
      output: 'test',
      turns: [
        {
          role: 'assistant',
          audio: { data: largeBase64, format: 'wav', transcript: 'hi' },
        },
      ],
      metadata: {
        audio: { data: largeBase64, format: 'wav', transcript: 'hi' },
      },
    } as unknown as ProviderResponse;

    const result = await extractAndStoreBinaryData(response);

    const turnsBlob = (result as any)?.turns?.[0]?.audio?.blobRef;
    expect(turnsBlob).toBeDefined();
    expect(result?.metadata?.audio?.blobRef).toEqual(turnsBlob);
    expect(mockStoreBlob).toHaveBeenCalledTimes(1);
    expect(mockStoreBlob).toHaveBeenCalledWith(
      expect.any(Buffer),
      'audio/wav',
      expect.objectContaining({
        location: 'response.turns[0].audio.data',
        kind: 'audio',
      }),
    );
  });

  it('should reuse the same image blob across images[] and a metadata data-URL of the same bytes', async () => {
    // E.g. a provider that mirrors a generated image into both `images[]` and
    // a metadata field. The shared per-response cache canonicalizes on the
    // parsed bytes so both paths land on one blob.
    const largeBase64 = Buffer.alloc(2000).toString('base64');
    const dataUri = `data:image/png;base64,${largeBase64}`;
    const response: ProviderResponse = {
      output: 'text',
      images: [{ data: dataUri, mimeType: 'image/png' }],
      metadata: {
        previewImage: dataUri,
      },
    };

    const result = await extractAndStoreBinaryData(response);

    expect(result?.images?.[0].blobRef?.uri).toBeDefined();
    expect(result?.metadata?.previewImage).toBe(result?.images?.[0].blobRef?.uri);
    expect(mockStoreBlob).toHaveBeenCalledTimes(1);
    expect(mockStoreBlob).toHaveBeenCalledWith(
      expect.any(Buffer),
      'image/png',
      expect.objectContaining({
        location: 'response.images[0].data',
        kind: 'image',
      }),
    );
  });

  it('should reuse the same image blob when output is a data URL and metadata holds raw base64', async () => {
    // Different encodings of identical bytes (data: URL vs. raw base64) must
    // hit the same cache slot — canonicalization happens on the parsed buffer.
    const largeBase64 = Buffer.alloc(2000).toString('base64');
    const dataUri = `data:image/png;base64,${largeBase64}`;
    const response: ProviderResponse = {
      output: dataUri,
      images: [{ data: `data:image/png;base64,${largeBase64}`, mimeType: 'image/png' }],
    };

    const result = await extractAndStoreBinaryData(response);

    expect(result?.output).toBe(result?.images?.[0].blobRef?.uri);
    expect(mockStoreBlob).toHaveBeenCalledTimes(1);
  });

  it('should store small raw SVG outputs for the media library without replacing text output', async () => {
    const svg = '<svg xmlns="http://www.w3.org/2000/svg"><circle r="4"/></svg>';
    const response: ProviderResponse = { output: svg };

    const result = await extractAndStoreBinaryData(response);

    expect(result?.output).toBe(svg);
    expect(result?.metadata?.blobUris).toEqual(['promptfoo://blob/abc123def456']);
    expect(mockStoreBlob).toHaveBeenCalledTimes(1);
    expect(mockStoreBlob).toHaveBeenCalledWith(
      Buffer.from(svg),
      'image/svg+xml',
      expect.objectContaining({
        location: 'response.output',
        kind: 'image',
      }),
    );
  });

  it('should not store the same raw SVG preview twice when metadata already references it', async () => {
    const svg = '<svg xmlns="http://www.w3.org/2000/svg"><circle r="4"/></svg>';
    const response: ProviderResponse = {
      output: svg,
      metadata: {
        blobUris: [`promptfoo://blob/${sha256(Buffer.from(svg))}`],
      },
    };

    const result = await extractAndStoreBinaryData(response);

    expect(result).toBe(response);
    expect(mockStoreBlob).not.toHaveBeenCalled();
  });

  it('should not store multiple root SVG documents as one media blob', async () => {
    const response: ProviderResponse = {
      output: '<svg><circle/></svg>\n\n<svg><rect/></svg>',
    };

    const result = await extractAndStoreBinaryData(response);

    expect(result).toBe(response);
    expect(mockStoreBlob).not.toHaveBeenCalled();
  });
});