diff --git a/middleware.test.ts b/middleware.test.ts index 2a95fe74..da2cba90 100644 --- a/middleware.test.ts +++ b/middleware.test.ts @@ -42,7 +42,7 @@ describe('middleware', () => { expect(response.status).toBe(429); }); - it('returns too many requests error body when rate limit fails', async () => { + it('returns too many requests JSON body for non-streak route when rate limit fails', async () => { vi.mocked(rateLimit).mockResolvedValue({ success: false, limit: 60, @@ -50,12 +50,30 @@ describe('middleware', () => { reset: 123456789, }); - const request = new NextRequest('http://localhost:3000/api/streak?user=octocat'); + const request = new NextRequest('http://localhost:3000/api/github'); const response = await middleware(request); await expect(response.json()).resolves.toEqual({ error: 'Too many requests', }); + + expect(response.headers.get('Content-Type')).toContain('application/json'); + }); + + it('returns SVG body for streak route when rate limit fails', async () => { + vi.mocked(rateLimit).mockResolvedValue({ + success: false, + limit: 60, + remaining: 0, + reset: 123456789, + }); + + const request = new NextRequest('http://localhost:3000/api/streak?user=octocat'); + const response = await middleware(request); + + await expect(response.text()).resolves.toContain(' { diff --git a/middleware.ts b/middleware.ts index bbb09744..cd0e3053 100644 --- a/middleware.ts +++ b/middleware.ts @@ -2,6 +2,19 @@ import { NextResponse } from 'next/server'; import type { NextRequest } from 'next/server'; import { rateLimit } from './lib/rate-limit'; +function generateRateLimitSVG() { + return ` + + + + Rate Limit Exceeded + + + Please try again later + +`; +} + /** * Middleware to enforce rate limiting on specific API routes. * @@ -26,16 +39,30 @@ export async function middleware(request: NextRequest) { // 60 requests per 60,000ms (1 minute) const result = await rateLimit(ip, 60, 60000); + const rateLimitHeaders = { + 'X-RateLimit-Limit': result.limit.toString(), + 'X-RateLimit-Remaining': result.remaining.toString(), + 'X-RateLimit-Reset': result.reset.toString(), + }; + if (!result.success) { + if (request.nextUrl.pathname.startsWith('/api/streak')) { + return new NextResponse(generateRateLimitSVG(), { + status: 429, + headers: { + 'Content-Type': 'image/svg+xml', + ...rateLimitHeaders, + }, + }); + } + return NextResponse.json( { error: 'Too many requests' }, { status: 429, headers: { 'Content-Type': 'application/json', - 'X-RateLimit-Limit': result.limit.toString(), - 'X-RateLimit-Remaining': result.remaining.toString(), - 'X-RateLimit-Reset': result.reset.toString(), + ...rateLimitHeaders, }, } ); @@ -43,9 +70,9 @@ export async function middleware(request: NextRequest) { // Add rate limit headers to the response for successful requests const response = NextResponse.next(); - response.headers.set('X-RateLimit-Limit', result.limit.toString()); - response.headers.set('X-RateLimit-Remaining', result.remaining.toString()); - response.headers.set('X-RateLimit-Reset', result.reset.toString()); + response.headers.set('X-RateLimit-Limit', rateLimitHeaders['X-RateLimit-Limit']); + response.headers.set('X-RateLimit-Remaining', rateLimitHeaders['X-RateLimit-Remaining']); + response.headers.set('X-RateLimit-Reset', rateLimitHeaders['X-RateLimit-Reset']); return response; }