diff --git a/apps/web/src/app/api/cloud-agent-next/sessions/stream-ticket/route.test.ts b/apps/web/src/app/api/cloud-agent-next/sessions/stream-ticket/route.test.ts new file mode 100644 index 0000000000..aba6064680 --- /dev/null +++ b/apps/web/src/app/api/cloud-agent-next/sessions/stream-ticket/route.test.ts @@ -0,0 +1,56 @@ +import { checkRateLimit } from '@vercel/firewall'; +import { getUserFromAuth } from '@/lib/user/server'; + +jest.mock('@vercel/firewall'); +jest.mock('@/lib/user/server'); +jest.mock('@/lib/drizzle'); +jest.mock('@/lib/cloud-agent/session-ownership'); +jest.mock('@/lib/cloud-agent/stream-ticket'); +jest.mock('@sentry/nextjs'); + +import { POST } from './route'; + +const mockCheckRateLimit = jest.mocked(checkRateLimit); +const mockGetUserFromAuth = jest.mocked(getUserFromAuth); + +function createRequest(): Request { + return new Request('http://localhost:3000/api/cloud-agent-next/sessions/stream-ticket', { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ cloudAgentSessionId: 'session-1' }), + }); +} + +describe('POST /api/cloud-agent-next/sessions/stream-ticket', () => { + beforeEach(() => { + jest.clearAllMocks(); + mockCheckRateLimit.mockResolvedValue({ rateLimited: false }); + }); + + test('returns 429 before authentication when the client IP is rate limited', async () => { + mockCheckRateLimit.mockResolvedValue({ rateLimited: true }); + const request = createRequest(); + + const response = await POST(request); + + expect(response.status).toBe(429); + expect(response.headers.get('Retry-After')).toBe('60'); + await expect(response.json()).resolves.toEqual({ + error: 'Rate limit exceeded. Please try again later.', + }); + expect(mockCheckRateLimit).toHaveBeenCalledWith('stream-ticket-ip', { request }); + expect(mockGetUserFromAuth).not.toHaveBeenCalled(); + }); + + test('continues to authentication when the client IP is allowed', async () => { + const authFailedResponse = Response.json({ error: 'Unauthorized' }, { status: 401 }); + mockGetUserFromAuth.mockResolvedValue({ user: null, authFailedResponse } as never); + const request = createRequest(); + + const response = await POST(request); + + expect(mockCheckRateLimit).toHaveBeenCalledWith('stream-ticket-ip', { request }); + expect(mockGetUserFromAuth).toHaveBeenCalledWith({ adminOnly: false }); + expect(response).toBe(authFailedResponse); + }); +}); diff --git a/apps/web/src/app/api/cloud-agent-next/sessions/stream-ticket/route.ts b/apps/web/src/app/api/cloud-agent-next/sessions/stream-ticket/route.ts index 4b6a55fa8c..c27d893c51 100644 --- a/apps/web/src/app/api/cloud-agent-next/sessions/stream-ticket/route.ts +++ b/apps/web/src/app/api/cloud-agent-next/sessions/stream-ticket/route.ts @@ -8,8 +8,11 @@ import { import { signStreamTicket } from '@/lib/cloud-agent/stream-ticket'; import { TRPCError } from '@trpc/server'; import { captureException } from '@sentry/nextjs'; +import { checkRateLimit } from '@vercel/firewall'; import * as z from 'zod'; +const STREAM_TICKET_IP_RATE_LIMIT_ID = 'stream-ticket-ip'; + const streamTicketSchema = z.object({ cloudAgentSessionId: z.string().min(1), organizationId: z.string().uuid().optional(), @@ -51,6 +54,14 @@ function handleTRPCError(error: unknown): NextResponse { */ export async function POST(request: Request) { try { + const { rateLimited } = await checkRateLimit(STREAM_TICKET_IP_RATE_LIMIT_ID, { request }); + if (rateLimited) { + return NextResponse.json( + { error: 'Rate limit exceeded. Please try again later.' }, + { status: 429, headers: { 'Retry-After': '60' } } + ); + } + const { user, authFailedResponse } = await getUserFromAuth({ adminOnly: false }); if (authFailedResponse) {