Merge branch 'aig' of https://github.com/fosrl/pangolin into aig

server/lib/aiGatewayTrust.ts

@@ -3,6 +3,24 @@ import config from "@server/lib/config";
 
 export const AI_GATEWAY_TRUST_HEADER = "X-Pangolin-Ai-Gateway-Auth";
 
+// Injected by the same Traefik trust middleware as AI_GATEWAY_TRUST_HEADER,
+// but its value differs per router (public inference resource vs. private
+// siteResource) so the gateway can tell which kind of resource a trusted
+// request arrived on without re-deriving it from resourceId/siteResourceId.
+export const AI_GATEWAY_RESOURCE_TYPE_HEADER =
+    "X-Pangolin-Ai-Gateway-Resource-Type";
+
+export type AiGatewayResourceType = "resource" | "site-resource";
+
+// Opt-in (server.enable_ai_gateway_client_ip_header): carries the client IP
+// that Badger resolved at the Traefik hop, so it survives an intermediary
+// proxy between Traefik and the AI gateway that overwrites
+// X-Forwarded-For/X-Real-Ip instead of appending to them. Set by a
+// disableForwardAuth Badger middleware instance (see getTraefikConfig.ts)
+// on the site-resource inference router only, since that's the sole path
+// that resolves request identity from the client IP.
+export const AI_GATEWAY_CLIENT_IP_HEADER = "X-Pangolin-Client-Ip";
+
 /**
  * Derive a Traefik-injected trust token from the server secret.
  * Traefik overwrites this header on inference routes so the AI gateway can
@@ -36,3 +54,16 @@ export function isAiGatewayTrustHeaderValid(
     const value = Array.isArray(raw) ? raw[0] : raw;
     return typeof value === "string" && value === expected;
 }
+
+export function getAiGatewayResourceType(
+    headers: Record<string, string | string[] | undefined> | undefined
+): AiGatewayResourceType | null {
+    if (!headers) {
+        return null;
+    }
+    const raw =
+        headers[AI_GATEWAY_RESOURCE_TYPE_HEADER] ??
+        headers[AI_GATEWAY_RESOURCE_TYPE_HEADER.toLowerCase()];
+    const value = Array.isArray(raw) ? raw[0] : raw;
+    return value === "resource" || value === "site-resource" ? value : null;
+}

server/lib/readConfigFile.ts

@@ -153,6 +153,25 @@ export const configSchema = z
                     })
                     .optional(),
                 trust_proxy: z.int().gte(0).optional().default(1),
+                // Opt-in: have Traefik/Badger stamp the resolved client IP
+                // into a dedicated header (X-Pangolin-Client-Ip) on the
+                // site-resource AI gateway route, so it survives an
+                // intermediary proxy between Traefik and the gateway that
+                // overwrites X-Forwarded-For/X-Real-Ip instead of appending
+                // to them. Off by default since it requires a Badger
+                // version that supports realIpHeader.
+                enable_ai_gateway_client_ip_header: z
+                    .boolean()
+                    .optional()
+                    .default(false)
+                    .transform((val) =>
+                        process.env.ENABLE_AI_GATEWAY_CLIENT_IP_HEADER !==
+                        undefined
+                            ? process.env
+                                  .ENABLE_AI_GATEWAY_CLIENT_IP_HEADER ===
+                                  "true"
+                            : val
+                    ),
                 secret: z.string().pipe(z.string().min(8)).optional(),
                 maxmind_db_path: z.string().optional(),
                 maxmind_asn_path: z.string().optional()
@@ -183,7 +202,8 @@ export const configSchema = z
                     "resource_session_request_param",
                 dashboard_session_length_hours: 720,
                 resource_session_length_hours: 720,
-                trust_proxy: 1
+                trust_proxy: 1,
+                enable_ai_gateway_client_ip_header: false
             }),
         postgres: z
             .object({

server/lib/traefik/getTraefikConfig.ts

@@ -26,6 +26,8 @@ import { sanitize, encodePath, validatePathRewriteConfig } from "./utils";
 import regionalCache from "@server/lib/cache";
 import {
     AI_GATEWAY_TRUST_HEADER,
+    AI_GATEWAY_RESOURCE_TYPE_HEADER,
+    AI_GATEWAY_CLIENT_IP_HEADER,
     getAiGatewayTrustToken
 } from "@server/lib/aiGatewayTrust";
 
@@ -752,21 +754,61 @@ export async function getTraefikConfig(
             aiGatewayHost = undefined;
         }
 
-        // The trust header is the same for every inference route on this
-        // exit node, so it's defined once here and attached to each router
-        // below instead of being duplicated into a per-resource middleware.
-        const aiGatewayTrustMiddlewareName = "ai-gateway-trust-headers";
+        // The trust token is the same for every inference route on this exit
+        // node, so it's defined once here and attached to each router below
+        // instead of being duplicated into a per-resource middleware. Two
+        // variants exist (public resource vs. siteResource) so the resource
+        // type header lets the gateway know which kind of router the
+        // request came through without re-deriving it from resourceId.
+        const aiGatewayTrustMiddlewareNameResource =
+            "ai-gateway-trust-headers-resource";
+        const aiGatewayTrustMiddlewareNameSiteResource =
+            "ai-gateway-trust-headers-site-resource";
         if (!config_output.http.middlewares) {
             config_output.http.middlewares = {};
         }
-        config_output.http.middlewares[aiGatewayTrustMiddlewareName] = {
+        config_output.http.middlewares[aiGatewayTrustMiddlewareNameResource] =
+            {
+                headers: {
+                    customRequestHeaders: {
+                        [AI_GATEWAY_TRUST_HEADER]: getAiGatewayTrustToken(),
+                        [AI_GATEWAY_RESOURCE_TYPE_HEADER]: "resource"
+                    }
+                }
+            };
+        config_output.http.middlewares[
+            aiGatewayTrustMiddlewareNameSiteResource
+        ] = {
             headers: {
                 customRequestHeaders: {
-                    [AI_GATEWAY_TRUST_HEADER]: getAiGatewayTrustToken()
+                    [AI_GATEWAY_TRUST_HEADER]: getAiGatewayTrustToken(),
+                    [AI_GATEWAY_RESOURCE_TYPE_HEADER]: "site-resource"
                 }
             }
         };
 
+        // Opt-in: a Badger instance with forward auth disabled, used only
+        // to stamp the resolved client IP into a dedicated header before
+        // the request reaches whatever sits between Traefik and the AI
+        // gateway. Only the site-resource router below needs this - it's
+        // the only path that resolves request identity from the client IP
+        // (see resolveRequestUser in aiGateway/pipeline.ts) - and it's the
+        // only inference router that doesn't already run Badger.
+        const aiGatewayClientIpMiddlewareName = "ai-gateway-client-ip";
+        const enableAiGatewayClientIpHeader =
+            config.getRawConfig().server.enable_ai_gateway_client_ip_header;
+        if (enableAiGatewayClientIpHeader) {
+            config_output.http.middlewares[aiGatewayClientIpMiddlewareName] =
+                {
+                    plugin: {
+                        badger: {
+                            disableForwardAuth: true,
+                            realIpHeader: AI_GATEWAY_CLIENT_IP_HEADER
+                        }
+                    }
+                };
+        }
+
         // Public inference resources: same TLS/cert-resolver handling as
         // plain http-mode resources, but the service points at the AI
         // gateway instead of any real backend targets.
@@ -836,7 +878,7 @@ export async function getTraefikConfig(
                 config.getRawConfig().traefik.additional_middlewares || [];
             const routerMiddlewares = [
                 badgerMiddlewareName,
-                aiGatewayTrustMiddlewareName,
+                aiGatewayTrustMiddlewareNameResource,
                 irHeadersMiddlewareName,
                 ...additionalMiddlewares
             ];
@@ -939,7 +981,10 @@ export async function getTraefikConfig(
             const additionalMiddlewares =
                 config.getRawConfig().traefik.additional_middlewares || [];
             const routerMiddlewares = [
-                aiGatewayTrustMiddlewareName,
+                ...(enableAiGatewayClientIpHeader
+                    ? [aiGatewayClientIpMiddlewareName]
+                    : []),
+                aiGatewayTrustMiddlewareNameSiteResource,
                 srHeadersMiddlewareName,
                 ...additionalMiddlewares
             ];
@@ -952,7 +997,7 @@ export async function getTraefikConfig(
                     middlewares: [redirectHttpsMiddlewareName],
                     service: serviceName,
                     rule,
-                    priority: 100
+                    priority: 200 // we want to match on the site resource first because the clientIP rule is more specific than the public inference resource rule, which is just the exit node IP range. so we give it a higher priority to ensure it matches first.
                 };
             }
 
@@ -965,7 +1010,7 @@ export async function getTraefikConfig(
                 middlewares: routerMiddlewares,
                 service: serviceName,
                 rule,
-                priority: 100,
+                priority: 200, // we want to match on the site resource first because the clientIP rule is more specific than the public inference resource rule, which is just the exit node IP range. so we give it a higher priority to ensure it matches first.
                 ...(sr.ssl ? { tls } : {})
             };