Address VPN gate review feedback

Author: jahooma

web/src/app/api/v1/chat/completions/__tests__/completions.test.ts

@@ -69,6 +69,7 @@ describe('/api/v1/chat/completions POST endpoint', () => {
   const allowedFreeModeHeaders = (apiKey: string) => ({
     Authorization: `Bearer ${apiKey}`,
     'cf-ipcountry': 'US',
+    'cf-connecting-ip': '203.0.113.10',
   })
 
   beforeEach(() => {

web/src/app/api/v1/freebuff/session/__tests__/session.test.ts

@@ -27,7 +27,10 @@ function makeReq(
   if (apiKey) headers.set('Authorization', `Bearer ${apiKey}`)
   if (opts.instanceId) headers.set(FREEBUFF_INSTANCE_HEADER, opts.instanceId)
   const cfCountry = opts.cfCountry === null ? null : (opts.cfCountry ?? 'US')
-  if (cfCountry) headers.set('cf-ipcountry', cfCountry)
+  if (cfCountry) {
+    headers.set('cf-ipcountry', cfCountry)
+    headers.set('cf-connecting-ip', '203.0.113.10')
+  }
   if (opts.model) headers.set(FREEBUFF_MODEL_HEADER, opts.model)
   return {
     headers,

web/src/app/api/v1/freebuff/session/_handlers.ts

@@ -25,7 +25,6 @@ import type { NextRequest } from 'next/server'
  *  the 10s error retry cadence. The new CLI parses the 403 body directly. */
 async function countryBlockedResponse(
   req: NextRequest,
-  deps: FreebuffSessionDeps,
 ): Promise<NextResponse | null> {
   const countryAccess = await getFreeModeCountryAccess(req, {
     ipinfoToken: env.IPINFO_TOKEN,
@@ -132,7 +131,7 @@ export async function postFreebuffSession(
   const auth = await resolveUser(req, deps)
   if ('error' in auth) return auth.error
 
-  const blocked = await countryBlockedResponse(req, deps)
+  const blocked = await countryBlockedResponse(req)
   if (blocked) return blocked
 
   const requestedModel = req.headers.get(FREEBUFF_MODEL_HEADER) ?? ''
@@ -176,7 +175,7 @@ export async function getFreebuffSession(
   const auth = await resolveUser(req, deps)
   if ('error' in auth) return auth.error
 
-  const blocked = await countryBlockedResponse(req, deps)
+  const blocked = await countryBlockedResponse(req)
   if (blocked) return blocked
 
   try {

web/src/server/__tests__/free-mode-country.test.ts

@@ -20,7 +20,10 @@ const noAnonymousNetwork = {
 describe('free mode country access', () => {
   test('allows allowlisted Cloudflare countries', async () => {
     const access = await getFreeModeCountryAccess(
-      makeReq({ 'cf-ipcountry': 'us' }),
+      makeReq({
+        'cf-ipcountry': 'us',
+        'cf-connecting-ip': '203.0.113.10',
+      }),
       noAnonymousNetwork,
     )
     expect(access.allowed).toBe(true)
@@ -58,6 +61,30 @@ describe('free mode country access', () => {
     expect(access.blockReason).toBe('missing_client_ip')
   })
 
+  test('blocks allowlisted Cloudflare countries when client IP is missing', async () => {
+    const access = await getFreeModeCountryAccess(
+      makeReq({ 'cf-ipcountry': 'US' }),
+      noAnonymousNetwork,
+    )
+    expect(access.allowed).toBe(false)
+    expect(access.countryCode).toBe(null)
+    expect(access.blockReason).toBe('missing_client_ip')
+    expect(access.cfCountry).toBe('US')
+  })
+
+  test('uses CF-Connecting-IP as a client IP fallback', async () => {
+    const access = await getFreeModeCountryAccess(
+      makeReq({
+        'cf-ipcountry': 'US',
+        'cf-connecting-ip': '203.0.113.10',
+      }),
+      noAnonymousNetwork,
+    )
+    expect(access.allowed).toBe(true)
+    expect(access.countryCode).toBe('US')
+    expect(access.hasClientIp).toBe(true)
+  })
+
   test('blocks allowlisted countries when the client IP is an anonymous network', async () => {
     const access = await getFreeModeCountryAccess(
       makeReq({

web/src/server/free-mode-country.ts

@@ -70,6 +70,7 @@ type ResolvedCountryAccess = Omit<
 }
 
 const IPINFO_PRIVACY_CACHE_TTL_MS = 30 * 60 * 1000
+const IPINFO_PRIVACY_CACHE_MAX_ENTRIES = 5000
 const ipinfoPrivacyCache = new Map<
   string,
   { expiresAt: number; privacy: FreeModeIpPrivacy | null }
@@ -80,7 +81,34 @@ export function extractClientIp(req: NextRequest): string | undefined {
   if (forwardedFor) {
     return forwardedFor.split(',')[0].trim()
   }
-  return req.headers.get('x-real-ip') ?? undefined
+  return (
+    req.headers.get('cf-connecting-ip') ??
+    req.headers.get('x-real-ip') ??
+    undefined
+  )
+}
+
+function setIpinfoPrivacyCache(
+  ip: string,
+  privacy: FreeModeIpPrivacy | null,
+): void {
+  const now = Date.now()
+  for (const [cachedIp, cached] of ipinfoPrivacyCache) {
+    if (cached.expiresAt <= now) {
+      ipinfoPrivacyCache.delete(cachedIp)
+    }
+  }
+
+  while (ipinfoPrivacyCache.size >= IPINFO_PRIVACY_CACHE_MAX_ENTRIES) {
+    const oldestIp = ipinfoPrivacyCache.keys().next().value
+    if (!oldestIp) break
+    ipinfoPrivacyCache.delete(oldestIp)
+  }
+
+  ipinfoPrivacyCache.set(ip, {
+    expiresAt: now + IPINFO_PRIVACY_CACHE_TTL_MS,
+    privacy,
+  })
 }
 
 function privacySignalsFromIpinfo(
@@ -123,10 +151,7 @@ export async function lookupIpinfoPrivacy(params: {
   const privacy = {
     signals,
   }
-  ipinfoPrivacyCache.set(params.ip, {
-    expiresAt: Date.now() + IPINFO_PRIVACY_CACHE_TTL_MS,
-    privacy,
-  })
+  setIpinfoPrivacyCache(params.ip, privacy)
   return privacy
 }
 
@@ -218,6 +243,18 @@ export async function getFreeModeCountryAccess(
     }
   }
 
+  if (!clientIp) {
+    return {
+      allowed: false,
+      countryCode: null,
+      blockReason: 'missing_client_ip',
+      cfCountry,
+      geoipCountry: null,
+      ipPrivacy: null,
+      hasClientIp: false,
+    }
+  }
+
   const ipPrivacy = await getIpPrivacy(clientIp, options)
   if (ipPrivacy?.signals.length) {
     return {