openmaic/OpenMAIC/tests/api/proxy-media.test.ts

import { beforeEach, describe, expect, it, vi } from 'vitest';

const mocks = vi.hoisted(() => {
  class MockPublicUrlFetchError extends Error {
    constructor(
      readonly code: 'INVALID_URL' | 'BLOCKED_TARGET' | 'DNS_FAILURE' | 'UPSTREAM_FAILURE',
      message: string,
    ) {
      super(message);
    }
  }

  return {
    fetchPinnedPublicUrl: vi.fn(),
    PublicUrlFetchError: MockPublicUrlFetchError,
  };
});

vi.mock('@/lib/server/public-url-fetch', () => ({
  fetchPinnedPublicUrl: mocks.fetchPinnedPublicUrl,
  PUBLIC_MEDIA_MAX_RESPONSE_BYTES: 25 * 1024 * 1024,
  PublicUrlFetchError: mocks.PublicUrlFetchError,
}));
vi.mock('@/lib/logger', () => ({
  createLogger: () => ({
    info: vi.fn(),
    warn: vi.fn(),
    error: vi.fn(),
    debug: vi.fn(),
  }),
}));

function pinned(response: Response) {
  return { response, dispose: vi.fn().mockResolvedValue(undefined) };
}

async function postProxy(body: Record<string, unknown>) {
  const { POST } = await import('@/app/api/proxy-media/route');
  return POST(
    new Request('http://localhost/api/proxy-media', {
      method: 'POST',
      headers: { 'Content-Type': 'application/json' },
      body: JSON.stringify(body),
    }),
  );
}

describe('POST /api/proxy-media', () => {
  beforeEach(() => {
    vi.resetModules();
    vi.restoreAllMocks();
    mocks.fetchPinnedPublicUrl.mockReset();
  });

  it('returns upstream bytes and content type, then disposes the pinned connection', async () => {
    const bodyBytes = new Uint8Array([1, 2, 3, 4]);
    const upstream = pinned(
      new Response(bodyBytes, {
        status: 200,
        headers: { 'Content-Type': 'image/png' },
      }),
    );
    mocks.fetchPinnedPublicUrl.mockResolvedValue(upstream);

    const res = await postProxy({ url: 'https://example.com/image.png' });

    expect(res.status).toBe(200);
    expect(res.headers.get('Content-Type')).toBe('image/png');
    expect(new Uint8Array(await res.arrayBuffer())).toEqual(bodyBytes);
    expect(mocks.fetchPinnedPublicUrl).toHaveBeenCalledWith(
      'https://example.com/image.png',
      expect.objectContaining({ maxResponseBytes: 25 * 1024 * 1024 }),
    );
    expect(upstream.dispose).toHaveBeenCalledOnce();
  });

  it('follows a redirect through a newly resolved and pinned hop', async () => {
    const first = pinned(
      new Response(null, {
        status: 302,
        headers: { location: 'https://cdn.example.com/final.png' },
      }),
    );
    const second = pinned(
      new Response(new Uint8Array([10, 20]), {
        status: 200,
        headers: { 'Content-Type': 'image/png' },
      }),
    );
    mocks.fetchPinnedPublicUrl.mockResolvedValueOnce(first).mockResolvedValueOnce(second);

    const res = await postProxy({ url: 'https://example.com/redirect' });

    expect(res.status).toBe(200);
    expect(mocks.fetchPinnedPublicUrl).toHaveBeenCalledTimes(2);
    expect(mocks.fetchPinnedPublicUrl.mock.calls[0][0]).toBe('https://example.com/redirect');
    expect(mocks.fetchPinnedPublicUrl.mock.calls[1][0]).toBe('https://cdn.example.com/final.png');
    expect(first.dispose).toHaveBeenCalledOnce();
    expect(second.dispose).toHaveBeenCalledOnce();
  });

  it('rejects a redirect whose next hop resolves to a blocked address', async () => {
    const first = pinned(
      new Response(null, {
        status: 302,
        headers: { location: 'http://100.100.100.200/latest/meta-data/' },
      }),
    );
    mocks.fetchPinnedPublicUrl
      .mockResolvedValueOnce(first)
      .mockRejectedValueOnce(
        new mocks.PublicUrlFetchError('BLOCKED_TARGET', 'sensitive internal detail'),
      );

    const res = await postProxy({ url: 'https://example.com/harmless' });
    const json = await res.json();

    expect(res.status).toBe(403);
    expect(json).toMatchObject({ errorCode: 'INVALID_URL', error: 'URL target is not allowed' });
    expect(JSON.stringify(json)).not.toContain('sensitive internal detail');
  });

  it('returns TOO_MANY_REDIRECTS after five followed redirects', async () => {
    mocks.fetchPinnedPublicUrl.mockImplementation(async () =>
      pinned(
        new Response(null, {
          status: 302,
          headers: { location: 'https://example.com/loop' },
        }),
      ),
    );

    const res = await postProxy({ url: 'https://example.com/loop' });
    const json = await res.json();

    expect(res.status).toBe(502);
    expect(json).toMatchObject({ errorCode: 'TOO_MANY_REDIRECTS' });
    expect(mocks.fetchPinnedPublicUrl).toHaveBeenCalledTimes(6);
  });

  it('maps an initially blocked URL to a generic 403 without leaking target details', async () => {
    mocks.fetchPinnedPublicUrl.mockRejectedValue(
      new mocks.PublicUrlFetchError('BLOCKED_TARGET', '127.0.0.1 is private'),
    );

    const res = await postProxy({ url: 'http://127.0.0.1/metadata' });
    const json = await res.json();

    expect(res.status).toBe(403);
    expect(json).toMatchObject({ errorCode: 'INVALID_URL', error: 'URL target is not allowed' });
    expect(JSON.stringify(json)).not.toContain('127.0.0.1');
  });

  it('rejects a redirect without Location', async () => {
    mocks.fetchPinnedPublicUrl.mockResolvedValue(pinned(new Response(null, { status: 302 })));

    const res = await postProxy({ url: 'https://example.com/bad-redirect' });
    const json = await res.json();

    expect(res.status).toBe(502);
    expect(json).toMatchObject({ errorCode: 'UPSTREAM_ERROR' });
  });

  it('forwards upstream 404 but collapses upstream 503 to 502', async () => {
    mocks.fetchPinnedPublicUrl.mockResolvedValueOnce(pinned(new Response(null, { status: 404 })));
    const missing = await postProxy({ url: 'https://example.com/not-found' });
    expect(missing.status).toBe(404);

    mocks.fetchPinnedPublicUrl.mockResolvedValueOnce(pinned(new Response(null, { status: 503 })));
    const failed = await postProxy({ url: 'https://example.com/server-error' });
    expect(failed.status).toBe(502);
  });

  it('rejects an upstream Content-Length over 25 MiB', async () => {
    mocks.fetchPinnedPublicUrl.mockResolvedValue(
      pinned(
        new Response(null, {
          status: 200,
          headers: { 'Content-Length': String(26 * 1024 * 1024) },
        }),
      ),
    );

    const res = await postProxy({ url: 'https://example.com/huge-asset' });
    const json = await res.json();

    expect(res.status).toBe(502);
    expect(json).toMatchObject({
      errorCode: 'UPSTREAM_ERROR',
      error: 'Upstream asset is too large',
    });
  });

  it('caps the actual JSON request stream at 8 KiB', async () => {
    const { POST } = await import('@/app/api/proxy-media/route');
    const res = await POST(
      new Request('http://localhost/api/proxy-media', {
        method: 'POST',
        headers: { 'Content-Type': 'application/json' },
        body: JSON.stringify({ url: `https://example.com/${'x'.repeat(9 * 1024)}` }),
      }),
    );

    expect(res.status).toBe(413);
    expect(mocks.fetchPinnedPublicUrl).not.toHaveBeenCalled();
  });

  it('requires JSON and rejects malformed JSON before fetching', async () => {
    const { POST } = await import('@/app/api/proxy-media/route');
    const wrongType = await POST(
      new Request('http://localhost/api/proxy-media', {
        method: 'POST',
        headers: { 'Content-Type': 'text/plain' },
        body: '{}',
      }),
    );
    const malformed = await POST(
      new Request('http://localhost/api/proxy-media', {
        method: 'POST',
        headers: { 'Content-Type': 'application/json' },
        body: '{',
      }),
    );

    expect(wrongType.status).toBe(415);
    expect(malformed.status).toBe(400);
    expect(mocks.fetchPinnedPublicUrl).not.toHaveBeenCalled();
  });
});