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

server/routers/aiGateway/pipeline.ts

@@ -45,7 +45,11 @@ import { isIpInCidr } from "@server/lib/ip";
 import { localCache } from "@server/lib/cache";
 import {
     AI_GATEWAY_TRUST_HEADER,
-    isAiGatewayTrustHeaderValid
+    AI_GATEWAY_RESOURCE_TYPE_HEADER,
+    AI_GATEWAY_CLIENT_IP_HEADER,
+    isAiGatewayTrustHeaderValid,
+    getAiGatewayResourceType,
+    type AiGatewayResourceType
 } from "@server/lib/aiGatewayTrust";
 import logger from "@server/logger";
 import HttpCode from "@server/types/HttpCode";
@@ -241,7 +245,8 @@ async function buildRequestUser(
 async function resolveRequestUser(
     req: Request,
     resourceId: number | null,
-    orgId: string | null
+    orgId: string | null,
+    resourceType: AiGatewayResourceType | null
 ): Promise<RequestIdentity> {
     // Public inference: identity comes from Badger via Remote-* only when the
     // Traefik trust header proves the request passed verify-session (VAK).
@@ -250,6 +255,7 @@ async function resolveRequestUser(
         const virtualApiKeyId =
             getRequestHeader(req, remoteHeaders.virtual_api_key_id) || null;
         const userId = getRequestHeader(req, remoteHeaders.user_id);
+
         if (userId) {
             const username =
                 getRequestHeader(req, remoteHeaders.user) || userId;
@@ -299,7 +305,20 @@ async function resolveRequestUser(
 
     // TODO: MAKE SURE THIS CAN NOT BE SPOOFED AND CAN BE TRUSTED AS AN INTERNAL ADDRESS FROM A NODE
 
-    const ip = req.ip;
+    // Only siteResources are reached over a client's exit-node tunnel, so
+    // only they can be attributed to a user by IP. Public resources must be
+    // authenticated via Badger above; the trust middleware stamps this
+    // header per-router so we don't need to re-derive it from resourceId.
+    if (resourceType !== "site-resource") {
+        return { user: null, virtualApiKeyId: null };
+    }
+
+    // Prefer the dedicated header Badger stamps at the Traefik hop (opt-in
+    // via server.enable_ai_gateway_client_ip_header) over req.ip, since an
+    // intermediary proxy between Traefik and this gateway may overwrite
+    // X-Forwarded-For/X-Real-Ip rather than appending to them, corrupting
+    // what req.ip resolves to.
+    const ip = getRequestHeader(req, AI_GATEWAY_CLIENT_IP_HEADER) || req.ip;
     if (!ip) {
         return { user: null, virtualApiKeyId: null };
     }
@@ -331,39 +350,46 @@ function getRequestHeader(req: Request, name: string): string | undefined {
     return raw;
 }
 
-async function resolveTarget(host: string): Promise<ResolvedTarget | null> {
-    const [[resourceRow], [siteResourceRow]] = await Promise.all([
-        db
-            .select({
-                resourceId: resources.resourceId,
-                orgId: resources.orgId
-            })
-            .from(resources)
-            .where(
-                and(
-                    eq(resources.fullDomain, host),
-                    eq(resources.mode, "inference"),
-                    eq(resources.enabled, true)
-                )
-            )
-            .limit(1),
-        db
-            .select({
-                siteResourceId: siteResources.siteResourceId,
-                orgId: siteResources.orgId
-            })
-            .from(siteResources)
-            .where(
-                and(
-                    eq(siteResources.fullDomain, host),
-                    eq(siteResources.mode, "inference"),
-                    eq(siteResources.enabled, true)
-                )
-            )
-            .limit(1)
-    ]);
+// Which table a host is looked up in depends on which Traefik router the
+// request came through, per the trust middleware's resource-type header -
+// falls back to checking both (public preferred on overlap) only when that
+// header is absent, e.g. a request that reached the gateway outside Traefik.
+async function resolveTarget(
+    host: string,
+    resourceType: AiGatewayResourceType | null
+): Promise<ResolvedTarget | null> {
+    if (resourceType === "resource") {
+        return resolveResourceTarget(host);
+    }
+    if (resourceType === "site-resource") {
+        return resolveSiteResourceTarget(host);
+    }
+
+    const [resourceTarget, siteResourceTarget] = await Promise.all([
+        resolveResourceTarget(host),
+        resolveSiteResourceTarget(host)
+    ]);
+    return resourceTarget ?? siteResourceTarget;
+}
+
+async function resolveResourceTarget(
+    host: string
+): Promise<ResolvedTarget | null> {
+    const [resourceRow] = await db
+        .select({
+            resourceId: resources.resourceId,
+            orgId: resources.orgId
+        })
+        .from(resources)
+        .where(
+            and(
+                eq(resources.fullDomain, host),
+                eq(resources.mode, "inference"),
+                eq(resources.enabled, true)
+            )
+        )
+        .limit(1);
 
-    // Prefer public inference resources when both match the same host.
     if (resourceRow) {
         const [attachmentRows, resourcePatterns] = await Promise.all([
             db
@@ -413,6 +439,27 @@ async function resolveTarget(host: string): Promise<ResolvedTarget | null> {
         };
     }
 
+    return null;
+}
+
+async function resolveSiteResourceTarget(
+    host: string
+): Promise<ResolvedTarget | null> {
+    const [siteResourceRow] = await db
+        .select({
+            siteResourceId: siteResources.siteResourceId,
+            orgId: siteResources.orgId
+        })
+        .from(siteResources)
+        .where(
+            and(
+                eq(siteResources.fullDomain, host),
+                eq(siteResources.mode, "inference"),
+                eq(siteResources.enabled, true)
+            )
+        )
+        .limit(1);
+
     if (siteResourceRow) {
         const [attachmentRows, resourcePatterns] = await Promise.all([
             db
@@ -764,8 +811,10 @@ export async function handleAiGatewayProxy(
                 );
         }
 
-        logger.debug(`AI gateway ${capability} request for host: ${host}`);
-        const target = await resolveTarget(host);
+        const resourceType = getAiGatewayResourceType(
+            req.headers as Record<string, string>
+        );
+        const target = await resolveTarget(host, resourceType);
         if (!target) {
             return res
                 .status(HttpCode.NOT_FOUND)
@@ -839,7 +888,7 @@ export async function handleAiGatewayProxy(
         const requestedModel = def.extractModel(req);
 
         const [identity, selection] = await Promise.all([
-            resolveRequestUser(req, resourceId, orgId),
+            resolveRequestUser(req, resourceId, orgId, resourceType),
             selectProvider(
                 capableAttachments,
                 resourceListsByProvider,
@@ -977,7 +1026,9 @@ export async function handleAiGatewayProxy(
             "upgrade",
             "content-length",
             "accept-encoding",
-            AI_GATEWAY_TRUST_HEADER.toLowerCase()
+            AI_GATEWAY_TRUST_HEADER.toLowerCase(),
+            AI_GATEWAY_RESOURCE_TYPE_HEADER.toLowerCase(),
+            AI_GATEWAY_CLIENT_IP_HEADER.toLowerCase()
         ]);
 
         const headers: Record<string, string> = {};

server/routers/aiGateway/targetRouting.ts

@@ -34,7 +34,10 @@ import {
     type RequestUser
 } from "@server/routers/aiGateway/pipeline";
 import { streamAiGatewayResponse } from "@server/routers/aiGateway/streamAiGatewayResponse";
-import { AI_GATEWAY_TRUST_HEADER } from "@server/lib/aiGatewayTrust";
+import {
+    AI_GATEWAY_TRUST_HEADER,
+    AI_GATEWAY_RESOURCE_TYPE_HEADER
+} from "@server/lib/aiGatewayTrust";
 
 // Short TTL: long enough to spare the DB on a burst of requests, short
 // enough that target/site changes (added, removed, exit node moved) show up
@@ -65,7 +68,8 @@ const SKIP_HEADERS = new Set([
     "upgrade",
     "content-length",
     "accept-encoding",
-    AI_GATEWAY_TRUST_HEADER.toLowerCase()
+    AI_GATEWAY_TRUST_HEADER.toLowerCase(),
+    AI_GATEWAY_RESOURCE_TYPE_HEADER.toLowerCase()
 ]);
 
 type ResolvedProviderTarget = {