search the right resource

server/lib/aiGatewayTrust.ts

@@ -3,6 +3,15 @@ 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";
+
 /**
  * Derive a Traefik-injected trust token from the server secret.
  * Traefik overwrites this header on inference routes so the AI gateway can
@@ -36,3 +45,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/traefik/getTraefikConfig.ts

@@ -26,6 +26,7 @@ import { sanitize, encodePath, validatePathRewriteConfig } from "./utils";
 import regionalCache from "@server/lib/cache";
 import {
     AI_GATEWAY_TRUST_HEADER,
+    AI_GATEWAY_RESOURCE_TYPE_HEADER,
     getAiGatewayTrustToken
 } from "@server/lib/aiGatewayTrust";
 
@@ -752,17 +753,35 @@ export async function getTraefikConfig(
             aiGatewayHost = undefined;
         }
 
-        // The trust header is the same for every inference route on this
+        // The trust token is the same for every inference route on this exit
-        // exit node, so it's defined once here and attached to each router
+        // node, so it's defined once here and attached to each router below
-        // below instead of being duplicated into a per-resource middleware.
+        // instead of being duplicated into a per-resource middleware. Two
-        const aiGatewayTrustMiddlewareName = "ai-gateway-trust-headers";
+        // 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"
                 }
             }
         };
@@ -836,7 +855,7 @@ export async function getTraefikConfig(
                 config.getRawConfig().traefik.additional_middlewares || [];
             const routerMiddlewares = [
                 badgerMiddlewareName,
-                aiGatewayTrustMiddlewareName,
+                aiGatewayTrustMiddlewareNameResource,
                 irHeadersMiddlewareName,
                 ...additionalMiddlewares
             ];
@@ -939,7 +958,7 @@ export async function getTraefikConfig(
             const additionalMiddlewares =
                 config.getRawConfig().traefik.additional_middlewares || [];
             const routerMiddlewares = [
-                aiGatewayTrustMiddlewareName,
+                aiGatewayTrustMiddlewareNameSiteResource,
                 srHeadersMiddlewareName,
                 ...additionalMiddlewares
             ];
@@ -952,7 +971,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 +984,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 } : {})
             };
 

server/private/lib/traefik/getTraefikConfig.ts

@@ -61,6 +61,7 @@ import { build } from "@server/build";
 import regionalCache from "#private/lib/cache";
 import {
     AI_GATEWAY_TRUST_HEADER,
+    AI_GATEWAY_RESOURCE_TYPE_HEADER,
     getAiGatewayTrustToken
 } from "@server/lib/aiGatewayTrust";
 
@@ -1569,14 +1570,32 @@ export async function getTraefikConfig(
         const aiGatewayOverride =
             config.getRawConfig().server.ai_gateway_override;
 
-        // The trust header is the same for every inference route on this
+        // The trust token is the same for every inference route on this exit
-        // exit node, so it's defined once here and attached to each router
+        // node, so it's defined once here and attached to each router below
-        // below instead of being duplicated into a per-resource middleware.
+        // instead of being duplicated into a per-resource middleware. Two
-        const aiGatewayTrustMiddlewareName = "ai-gateway-trust-headers";
+        // variants exist (public resource vs. siteResource) so the resource
-        config_output.http.middlewares[aiGatewayTrustMiddlewareName] = {
+        // 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";
+        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"
                 }
             }
         };
@@ -1650,7 +1669,7 @@ export async function getTraefikConfig(
                 config.getRawConfig().traefik.additional_middlewares || [];
             const routerMiddlewares = [
                 badgerMiddlewareName,
-                aiGatewayTrustMiddlewareName
+                aiGatewayTrustMiddlewareNameResource
             ];
 
             if (aiGatewayOverride) {
@@ -1750,7 +1769,9 @@ export async function getTraefikConfig(
 
             const additionalMiddlewares =
                 config.getRawConfig().traefik.additional_middlewares || [];
-            const routerMiddlewares: string[] = [aiGatewayTrustMiddlewareName];
+            const routerMiddlewares: string[] = [
+                aiGatewayTrustMiddlewareNameSiteResource
+            ];
 
             if (aiGatewayOverride) {
                 const srHeadersMiddlewareName = `${srKey}-headers-middleware`;
@@ -1775,7 +1796,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.
                 };
             }
 
@@ -1788,7 +1809,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 } : {})
             };
 

server/routers/aiGateway/pipeline.ts

@@ -41,7 +41,10 @@ import { isIpInCidr } from "@server/lib/ip";
 import { localCache } from "@server/lib/cache";
 import {
     AI_GATEWAY_TRUST_HEADER,
-    isAiGatewayTrustHeaderValid
+    AI_GATEWAY_RESOURCE_TYPE_HEADER,
+    isAiGatewayTrustHeaderValid,
+    getAiGatewayResourceType,
+    type AiGatewayResourceType
 } from "@server/lib/aiGatewayTrust";
 import logger from "@server/logger";
 import HttpCode from "@server/types/HttpCode";
@@ -232,7 +235,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).
@@ -241,10 +245,7 @@ async function resolveRequestUser(
         const virtualApiKeyId =
             getRequestHeader(req, remoteHeaders.virtual_api_key_id) || null;
         const userId = getRequestHeader(req, remoteHeaders.user_id);
-        logger.debug("+++++++AI gateway request identity from trust header", {
+
-            virtualApiKeyId,
-            userId
-        });
         if (userId) {
             const username =
                 getRequestHeader(req, remoteHeaders.user) || userId;
@@ -294,6 +295,14 @@ async function resolveRequestUser(
 
     // TODO: MAKE SURE THIS CAN NOT BE SPOOFED AND CAN BE TRUSTED AS AN INTERNAL ADDRESS FROM A NODE
 
+    // 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 };
+    }
+
     const ip = req.ip;
     if (!ip) {
         return { user: null, virtualApiKeyId: null };
@@ -326,39 +335,46 @@ function getRequestHeader(req: Request, name: string): string | undefined {
     return raw;
 }
 
-async function resolveTarget(host: string): Promise<ResolvedTarget | null> {
+// Which table a host is looked up in depends on which Traefik router the
-    const [[resourceRow], [siteResourceRow]] = await Promise.all([
+// request came through, per the trust middleware's resource-type header -
-        db
+// falls back to checking both (public preferred on overlap) only when that
-            .select({
+// header is absent, e.g. a request that reached the gateway outside Traefik.
-                resourceId: resources.resourceId,
+async function resolveTarget(
-                orgId: resources.orgId
+    host: string,
-            })
+    resourceType: AiGatewayResourceType | null
-            .from(resources)
+): Promise<ResolvedTarget | null> {
-            .where(
+    if (resourceType === "resource") {
-                and(
+        return resolveResourceTarget(host);
-                    eq(resources.fullDomain, host),
+    }
-                    eq(resources.mode, "inference"),
+    if (resourceType === "site-resource") {
-                    eq(resources.enabled, true)
+        return resolveSiteResourceTarget(host);
-                )
+    }
-            )
+
-            .limit(1),
+    const [resourceTarget, siteResourceTarget] = await Promise.all([
-        db
+        resolveResourceTarget(host),
-            .select({
+        resolveSiteResourceTarget(host)
-                siteResourceId: siteResources.siteResourceId,
+    ]);
-                orgId: siteResources.orgId
+    return resourceTarget ?? siteResourceTarget;
-            })
+}
-            .from(siteResources)
+
-            .where(
+async function resolveResourceTarget(
-                and(
+    host: string
-                    eq(siteResources.fullDomain, host),
+): Promise<ResolvedTarget | null> {
-                    eq(siteResources.mode, "inference"),
+    const [resourceRow] = await db
-                    eq(siteResources.enabled, true)
+        .select({
-                )
+            resourceId: resources.resourceId,
-            )
+            orgId: resources.orgId
-            .limit(1)
+        })
-    ]);
+        .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
@@ -408,6 +424,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
@@ -753,7 +790,10 @@ export async function handleAiGatewayProxy(
             headers: req.headers
         });
 
-        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).json({
                 error: {
@@ -808,7 +848,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,
@@ -924,7 +964,8 @@ 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()
         ]);
 
         const headers: Record<string, string> = {};

server/routers/aiGateway/targetRouting.ts

@@ -33,7 +33,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
@@ -64,7 +67,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 = {