From b879b21975bba867b09bdd3cf8172d5676e38514 Mon Sep 17 00:00:00 2001 From: Lukasz Raczylo Date: Thu, 20 Aug 2026 12:02:06 +0100 Subject: [PATCH] feat: MCP/A2A registry and authed streaming proxy MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add GET /v1/mcp/servers and GET /v1/agents: group-filtered, name-sorted registry listings of configured MCP servers and A2A agents, with their gateway-relative proxy (and, for agents, card) URLs. Add handleTargetProxy for /mcp/{name}/{rest...} and /a2a/{name}/{rest...}: reverse-proxies to the target's configured URL. Checks run unknown name (404), group deny via allowsMCP/ allowsAgent (403), request-rate limits (429/503), then Upgrade (501) and path traversal (400) around the proxy call itself. No upstream credential is injected — MCP servers and A2A agents are in-cluster targets, unlike native provider passthrough. Extract proxyUpstream from routes_passthrough.go's handlePassthrough into a shared reverse-proxy core (header stripping, SSE-safe flushWriter streaming, dead-upstream 502, canceled-client log-only, optional JSON usage teeing) used by both native passthrough and the new target proxy. Behavior-preserving for Task 13: its tests pass unmodified. Validate every configured MCP-server and agent target URL at construction (url.Parse, http/https scheme only) — a malformed or non-HTTP target is a constructor error, not a per-request 502. --- llmgateway.go | 44 +++- mcp_a2a.go | 214 ++++++++++++++++ mcp_a2a_test.go | 578 ++++++++++++++++++++++++++++++++++++++++++ routes_passthrough.go | 137 +++++++--- 4 files changed, 931 insertions(+), 42 deletions(-) create mode 100644 mcp_a2a.go create mode 100644 mcp_a2a_test.go diff --git a/llmgateway.go b/llmgateway.go index a05aa15..b9fa238 100644 --- a/llmgateway.go +++ b/llmgateway.go @@ -107,7 +107,12 @@ type Gateway struct { limiter *limiter registry *modelRegistry adapters map[string]providerAdapter - name string + // targetClient is the shared, connection-pooled *http.Client the + // MCP/A2A target proxy (mcp_a2a.go) issues every upstream request + // through — built once via newAdapterHTTPClient, the same constructor + // each provider adapter uses for its own client. + targetClient *http.Client + name string } // New creates the middleware. NOTE: no tail call — Yaegi zeroes @@ -124,6 +129,9 @@ func newGateway(ctx context.Context, next http.Handler, config *Config, name str if config == nil || len(config.Providers) == 0 { return nil, errors.New("llmgateway: at least one provider must be configured") } + if err := validateTargetURLs(config); err != nil { + return nil, err + } auth, err := newAuthStore(config) if err != nil { return nil, err @@ -148,6 +156,7 @@ func newGateway(ctx context.Context, next http.Handler, config *Config, name str return nil, err } g.adapters = adapters + g.targetClient = newAdapterHTTPClient() registry, err := newModelRegistry(adapters, config, g.errorf) if err != nil { @@ -268,6 +277,39 @@ func (g *Gateway) ServeHTTP(w http.ResponseWriter, r *http.Request) { return } + if r.Method == http.MethodGet && (r.URL.Path == "/v1/mcp/servers" || r.URL.Path == "/v1/agents") { + _, grp, ok := g.auth.identify(r) + if !ok { + writeOAIError(sw, http.StatusUnauthorized, "authentication_error", "invalid or missing API key") + return + } + if r.URL.Path == "/v1/mcp/servers" { + g.handleMCPServers(sw, grp) + } else { + g.handleAgents(sw, grp) + } + return + } + + if name, rest, ok := targetRoute(r.URL.EscapedPath(), targetKindMCP); ok { + u, grp, authOK := g.auth.identify(r) + if !authOK { + writeOAIError(sw, http.StatusUnauthorized, "authentication_error", "invalid or missing API key") + return + } + g.handleTargetProxy(sw, r, u, grp, targetKindMCP, name, rest) + return + } + if name, rest, ok := targetRoute(r.URL.EscapedPath(), targetKindAgent); ok { + u, grp, authOK := g.auth.identify(r) + if !authOK { + writeOAIError(sw, http.StatusUnauthorized, "authentication_error", "invalid or missing API key") + return + } + g.handleTargetProxy(sw, r, u, grp, targetKindAgent, name, rest) + return + } + if providerName, rest, ok := passthroughRoute(r.URL.EscapedPath()); ok { if _, known := g.adapters[providerName]; known { u, grp, ok := g.auth.identify(r) diff --git a/mcp_a2a.go b/mcp_a2a.go new file mode 100644 index 0000000..d4cee68 --- /dev/null +++ b/mcp_a2a.go @@ -0,0 +1,214 @@ +package traefikllmgateway + +import ( + "encoding/json" + "fmt" + "net/http" + "net/url" + "sort" +) + +// Target kinds handleTargetProxy and resolveTarget dispatch on, selecting +// which of Config.MCPServers/Config.Agents to consult and which of +// group.allowsMCP/group.allowsAgent authorizes the request. +const ( + targetKindMCP = "mcp" + targetKindAgent = "a2a" +) + +// defaultAgentCardPath is the well-known A2A agent-card path used for a +// listed agent whose AgentConfig.Card is empty. +const defaultAgentCardPath = "/.well-known/agent-card.json" + +// mcpServerListing is one entry of GET /v1/mcp/servers: the configured +// name and the gateway-relative URL a client proxies requests to it +// through. +type mcpServerListing struct { + Name string `json:"name"` + URL string `json:"url"` +} + +// agentListing is one entry of GET /v1/agents: the configured name, the +// gateway-relative proxy URL, and the gateway-relative agent-card URL. +type agentListing struct { + Name string `json:"name"` + URL string `json:"url"` + Card string `json:"card"` +} + +// handleMCPServers implements GET /v1/mcp/servers: grp's visible MCP +// servers, filtered by group.allowsMCP and sorted by name for a +// deterministic response. +func (g *Gateway) handleMCPServers(w http.ResponseWriter, grp *group) { + names := make([]string, 0, len(g.cfg.MCPServers)) + for name := range g.cfg.MCPServers { + if grp.allowsMCP(name) { + names = append(names, name) + } + } + sort.Strings(names) + + servers := make([]mcpServerListing, len(names)) + for i, name := range names { + servers[i] = mcpServerListing{Name: name, URL: "/mcp/" + name} + } + + w.Header().Set("Content-Type", "application/json") + _ = json.NewEncoder(w).Encode(map[string]any{"servers": servers}) // headers already committed; nothing useful to do on encode failure +} + +// handleAgents implements GET /v1/agents: grp's visible A2A agents, +// filtered by group.allowsAgent and sorted by name for a deterministic +// response. Each entry's card URL is the agent's own gateway-relative +// proxy URL plus its configured card path, defaulting to +// defaultAgentCardPath when AgentConfig.Card is empty. +func (g *Gateway) handleAgents(w http.ResponseWriter, grp *group) { + names := make([]string, 0, len(g.cfg.Agents)) + for name := range g.cfg.Agents { + if grp.allowsAgent(name) { + names = append(names, name) + } + } + sort.Strings(names) + + agents := make([]agentListing, len(names)) + for i, name := range names { + cardPath := g.cfg.Agents[name].Card + if cardPath == "" { + cardPath = defaultAgentCardPath + } + agents[i] = agentListing{Name: name, URL: "/a2a/" + name, Card: "/a2a/" + name + cardPath} + } + + w.Header().Set("Content-Type", "application/json") + _ = json.NewEncoder(w).Encode(map[string]any{"agents": agents}) // headers already committed; nothing useful to do on encode failure +} + +// targetRoute splits path into a target name and its rest path for the +// "/{prefix}/{name}/{rest...}" MCP/A2A convention (prefix is +// targetKindMCP or targetKindAgent), reusing passthroughRoute's own +// EscapedPath-based, traversal-safe splitting twice: once to peel prefix +// off path, once more on what is left to peel name off rest. ok is false +// when path's first segment is not exactly prefix, or when no name +// segment follows it — a bare "/mcp" or "/mcp/" carries no target name to +// route to. +func targetRoute(path, prefix string) (name, rest string, ok bool) { + first, remainder, matched := passthroughRoute(path) + if !matched || first != prefix { + return "", "", false + } + return passthroughRoute("/" + remainder) +} + +// resolveTarget looks up name in the Gateway's MCP-server or agent +// registry, selected by kind (targetKindMCP or targetKindAgent). known is +// false when no such name is configured at all; allowed reports whether +// grp's glob patterns permit it, meaningful only when known is true. +func (g *Gateway) resolveTarget(kind, name string, grp *group) (targetURL string, allowed, known bool) { + switch kind { + case targetKindMCP: + tc, ok := g.cfg.MCPServers[name] + if !ok { + return "", false, false + } + return tc.URL, grp.allowsMCP(name), true + case targetKindAgent: + ac, ok := g.cfg.Agents[name] + if !ok { + return "", false, false + } + return ac.URL, grp.allowsAgent(name), true + default: + return "", false, false + } +} + +// handleTargetProxy implements the MCP/A2A target proxy routes: +// "/mcp/{name}/{rest...}" and "/a2a/{name}/{rest...}" reverse-proxy rest, +// verbatim, to name's configured target URL (Config.MCPServers[name].URL +// or Config.Agents[name].URL, per kind). It shares its proxy mechanics +// with handlePassthrough via proxyUpstream — hop-by-hop header stripping, +// the client's own gateway credential stripped, SSE-safe incremental +// flushing, a dead upstream mapped to 502, a canceled client context +// logged rather than surfaced. Unlike handlePassthrough, no upstream +// credential is injected (proxyUpstream's injectAuth argument is nil): +// MCP servers and A2A agents are in-cluster targets that trust the +// gateway's network position, not a per-provider API key the gateway +// holds on the caller's behalf. Accounting is likewise passthrough-only — +// a target request is counted by checkAndCount below, but its response +// is never teed off for token/cost usage extraction. +// +// Checks run: unknown name (404) before group authorization (403) before +// rate limits (429/503) before the capability/path checks (Upgrade→501, +// traversal→400) that gate the proxy call itself — a caller addressing a +// target that does not exist, or that they may not use, or that their +// group is already over budget on, learns that before any detail of how +// they tried to use it. +func (g *Gateway) handleTargetProxy(w http.ResponseWriter, r *http.Request, u *user, grp *group, kind, name, rest string) { + targetURL, allowed, known := g.resolveTarget(kind, name, grp) + if !known { + writeOAIError(w, http.StatusNotFound, "invalid_request_error", "unknown "+kind+" target") + return + } + if !allowed { + writeOAIError(w, http.StatusForbidden, "invalid_request_error", kind+" access denied") + return + } + + scopes := buildLimitScopes(u, grp) + if violation := g.limiter.checkAndCount(scopes); violation != nil { + writeLimitViolation(w, violation) + return + } + + if r.Header.Get("Upgrade") != "" { + writeOAIError(w, http.StatusNotImplemented, "invalid_request_error", "websocket/upgrade passthrough not supported") + return + } + if hasTraversalSegment(rest) { + writeOAIError(w, http.StatusBadRequest, "invalid_request_error", "invalid path") + return + } + + upstreamURL := targetURL + if rest != "" { + upstreamURL += "/" + rest + } + if r.URL.RawQuery != "" { + upstreamURL += "?" + r.URL.RawQuery + } + + g.proxyUpstream(w, r, upstreamURL, g.targetClient, nil, false, kind+" target (name "+name+")") +} + +// validateTargetURLs checks that every configured MCP-server and agent +// target URL parses (url.Parse) and uses an http or https scheme. A +// malformed or non-HTTP target is rejected once, here, at construction — +// a constructor error, not a 502 the first time some caller happens to +// address it. +func validateTargetURLs(cfg *Config) error { + for name, tc := range cfg.MCPServers { + if err := validateTargetURL(tc.URL); err != nil { + return fmt.Errorf("llmgateway: mcpServers %q: %w", name, err) + } + } + for name, ac := range cfg.Agents { + if err := validateTargetURL(ac.URL); err != nil { + return fmt.Errorf("llmgateway: agents %q: %w", name, err) + } + } + return nil +} + +// validateTargetURL reports an error unless raw parses as a URL with an +// http or https scheme. +func validateTargetURL(raw string) error { + u, err := url.Parse(raw) + if err != nil { + return fmt.Errorf("invalid url %q: %w", raw, err) + } + if u.Scheme != "http" && u.Scheme != "https" { + return fmt.Errorf("url %q must use http or https scheme, got %q", raw, u.Scheme) + } + return nil +} diff --git a/mcp_a2a_test.go b/mcp_a2a_test.go new file mode 100644 index 0000000..b50f905 --- /dev/null +++ b/mcp_a2a_test.go @@ -0,0 +1,578 @@ +package traefikllmgateway + +import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "strings" + "testing" + "time" +) + +// --- unit tests: targetRoute --- + +func TestTargetRoute(t *testing.T) { + tests := []struct { + path string + prefix string + wantName string + wantRest string + wantOK bool + }{ + {"/mcp/myserver/tools/list", "mcp", "myserver", "tools/list", true}, + {"/mcp/myserver", "mcp", "myserver", "", true}, + {"/mcp/myserver/", "mcp", "myserver", "", true}, + {"/mcp", "mcp", "", "", false}, + {"/mcp/", "mcp", "", "", false}, + {"/a2a/agent1/.well-known/agent-card.json", "a2a", "agent1", ".well-known/agent-card.json", true}, + {"/a2a/agent1", "a2a", "agent1", "", true}, + {"/notmcp/foo", "mcp", "", "", false}, + {"/", "mcp", "", "", false}, + {"", "mcp", "", "", false}, + } + for _, tt := range tests { + t.Run(tt.prefix+tt.path, func(t *testing.T) { + gotName, gotRest, gotOK := targetRoute(tt.path, tt.prefix) + if gotOK != tt.wantOK || gotName != tt.wantName || gotRest != tt.wantRest { + t.Errorf("targetRoute(%q, %q) = (%q, %q, %v), want (%q, %q, %v)", + tt.path, tt.prefix, gotName, gotRest, gotOK, tt.wantName, tt.wantRest, tt.wantOK) + } + }) + } +} + +// --- registry filtering: GET /v1/mcp/servers, GET /v1/agents --- + +func TestHandleMCPServers_GroupFiltered_SortedByName(t *testing.T) { + cfg := CreateConfig() + cfg.Providers = map[string]*ProviderConfig{"openai": {Type: "openai", APIKey: "k"}} + cfg.MCPServers = map[string]*TargetConfig{ + "charlie": {URL: "http://mcp-charlie.internal"}, + "alpha": {URL: "http://mcp-alpha.internal"}, + "bravo": {URL: "http://mcp-bravo.internal"}, + } + cfg.Groups = map[string]*GroupConfig{"limited": {MCPServers: []string{"alpha", "charlie"}}} + cfg.Users = &UsersConfig{Inline: []*UserConfig{{Name: "alice", Group: "limited", APIKey: "sk-alice"}}} + next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}) + h, err := New(context.Background(), next, cfg, "llmgw") + if err != nil { + t.Fatalf("New: %v", err) + } + + req := httptest.NewRequest(http.MethodGet, "/v1/mcp/servers", nil) + req.Header.Set("Authorization", "Bearer sk-alice") + rec := httptest.NewRecorder() + h.ServeHTTP(rec, req) + + if rec.Code != http.StatusOK { + t.Fatalf("status = %d, want 200, body=%s", rec.Code, rec.Body.String()) + } + + var got struct { + Servers []struct { + Name string `json:"name"` + URL string `json:"url"` + } `json:"servers"` + } + if err := json.Unmarshal(rec.Body.Bytes(), &got); err != nil { + t.Fatalf("decode response: %v", err) + } + if len(got.Servers) != 2 { + t.Fatalf("servers = %+v, want 2 entries (alpha, charlie; bravo denied by group)", got.Servers) + } + if got.Servers[0].Name != "alpha" || got.Servers[0].URL != "/mcp/alpha" { + t.Errorf("servers[0] = %+v, want name=alpha url=/mcp/alpha", got.Servers[0]) + } + if got.Servers[1].Name != "charlie" || got.Servers[1].URL != "/mcp/charlie" { + t.Errorf("servers[1] = %+v, want name=charlie url=/mcp/charlie", got.Servers[1]) + } +} + +func TestHandleAgents_GroupFiltered_SortedByName_CardPath(t *testing.T) { + cfg := CreateConfig() + cfg.Providers = map[string]*ProviderConfig{"openai": {Type: "openai", APIKey: "k"}} + cfg.Agents = map[string]*AgentConfig{ + "zeta": {URL: "http://agent-zeta.internal"}, + "delta": {URL: "http://agent-delta.internal", Card: "/custom/card.json"}, + } + cfg.Groups = map[string]*GroupConfig{"default": {}} + cfg.Users = &UsersConfig{Inline: []*UserConfig{{Name: "alice", Group: "default", APIKey: "sk-alice"}}} + next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}) + h, err := New(context.Background(), next, cfg, "llmgw") + if err != nil { + t.Fatalf("New: %v", err) + } + + req := httptest.NewRequest(http.MethodGet, "/v1/agents", nil) + req.Header.Set("Authorization", "Bearer sk-alice") + rec := httptest.NewRecorder() + h.ServeHTTP(rec, req) + + if rec.Code != http.StatusOK { + t.Fatalf("status = %d, want 200, body=%s", rec.Code, rec.Body.String()) + } + + var got struct { + Agents []struct { + Name string `json:"name"` + URL string `json:"url"` + Card string `json:"card"` + } `json:"agents"` + } + if err := json.Unmarshal(rec.Body.Bytes(), &got); err != nil { + t.Fatalf("decode response: %v", err) + } + if len(got.Agents) != 2 { + t.Fatalf("agents = %+v, want 2 entries", got.Agents) + } + // sorted: delta before zeta + if got.Agents[0].Name != "delta" || got.Agents[0].Card != "/a2a/delta/custom/card.json" { + t.Errorf("agents[0] = %+v, want name=delta card=/a2a/delta/custom/card.json", got.Agents[0]) + } + if got.Agents[1].Name != "zeta" || got.Agents[1].Card != "/a2a/zeta/.well-known/agent-card.json" { + t.Errorf("agents[1] = %+v, want name=zeta card=/a2a/zeta/.well-known/agent-card.json (default cardPath)", got.Agents[1]) + } +} + +func TestHandleMCPServers_Unauthenticated_Returns401(t *testing.T) { + cfg := CreateConfig() + cfg.Providers = map[string]*ProviderConfig{"openai": {Type: "openai", APIKey: "k"}} + cfg.MCPServers = map[string]*TargetConfig{"alpha": {URL: "http://mcp-alpha.internal"}} + next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}) + h, err := New(context.Background(), next, cfg, "llmgw") + if err != nil { + t.Fatalf("New: %v", err) + } + + for _, path := range []string{"/v1/mcp/servers", "/v1/agents"} { + t.Run(path, func(t *testing.T) { + req := httptest.NewRequest(http.MethodGet, path, nil) + rec := httptest.NewRecorder() + h.ServeHTTP(rec, req) + if rec.Code != http.StatusUnauthorized { + t.Fatalf("status = %d, want 401, body=%s", rec.Code, rec.Body.String()) + } + }) + } +} + +// --- target proxy: strips gateway auth, injects nothing --- + +func TestHandleTargetProxy_StripsGatewayAuth_NoUpstreamInjection(t *testing.T) { + var gotAuth, gotAPIKeyHeader, gotConnection string + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + gotAuth = r.Header.Get("Authorization") + gotAPIKeyHeader = r.Header.Get("x-api-key") + gotConnection = r.Header.Get("Connection") + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte("mcp-ok")) + })) + defer srv.Close() + + cfg := CreateConfig() + cfg.Providers = map[string]*ProviderConfig{"openai": {Type: "openai", APIKey: "k"}} + cfg.MCPServers = map[string]*TargetConfig{"alpha": {URL: srv.URL}} + cfg.Groups = map[string]*GroupConfig{"default": {}} + cfg.Users = &UsersConfig{Inline: []*UserConfig{{Name: "alice", Group: "default", APIKey: "sk-alice"}}} + next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}) + h, err := New(context.Background(), next, cfg, "llmgw") + if err != nil { + t.Fatalf("New: %v", err) + } + + req := httptest.NewRequest(http.MethodGet, "/mcp/alpha/tools/list", nil) + req.Header.Set("Authorization", "Bearer sk-alice") + req.Header.Set("Connection", "keep-alive") + rec := httptest.NewRecorder() + h.ServeHTTP(rec, req) + + if rec.Code != http.StatusOK { + t.Fatalf("status = %d, want 200, body=%s", rec.Code, rec.Body.String()) + } + if rec.Body.String() != "mcp-ok" { + t.Errorf("body = %q, want verbatim upstream body", rec.Body.String()) + } + if gotAuth != "" { + t.Errorf("upstream saw Authorization %q, want stripped and nothing injected (in-cluster target)", gotAuth) + } + if gotAPIKeyHeader != "" { + t.Errorf("upstream saw x-api-key %q, want stripped", gotAPIKeyHeader) + } + if gotConnection != "" { + t.Errorf("upstream saw hop-by-hop Connection header %q, want stripped", gotConnection) + } +} + +// --- SSE-safe streaming through the target proxy --- + +func TestHandleTargetProxy_SSE_FlushesIncrementally(t *testing.T) { + continueCh := make(chan struct{}) + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + w.WriteHeader(http.StatusOK) + fl, ok := w.(http.Flusher) + if !ok { + t.Fatal("fake upstream ResponseWriter does not support Flush") + } + _, _ = w.Write([]byte("data: chunk-1\n\n")) + fl.Flush() + <-continueCh + _, _ = w.Write([]byte("data: chunk-2\n\n")) + fl.Flush() + })) + defer srv.Close() + + cfg := CreateConfig() + cfg.Providers = map[string]*ProviderConfig{"openai": {Type: "openai", APIKey: "k"}} + cfg.MCPServers = map[string]*TargetConfig{"alpha": {URL: srv.URL}} + cfg.Groups = map[string]*GroupConfig{"default": {}} + cfg.Users = &UsersConfig{Inline: []*UserConfig{{Name: "alice", Group: "default", APIKey: "sk-alice"}}} + next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}) + h, err := New(context.Background(), next, cfg, "llmgw") + if err != nil { + t.Fatalf("New: %v", err) + } + + req := httptest.NewRequest(http.MethodGet, "/mcp/alpha/stream", nil) + req.Header.Set("Authorization", "Bearer sk-alice") + sw := newSyncFlushWriter() + done := make(chan struct{}) + go func() { + h.ServeHTTP(sw, req) + close(done) + }() + + select { + case <-sw.writeCh: + case <-time.After(5 * time.Second): + t.Fatal("timed out waiting for the first chunk to reach the client") + } + _, writesAfterFirst, _, bodyAfterFirst := sw.snapshot() + if writesAfterFirst != 1 { + t.Fatalf("writes after first chunk = %d, want exactly 1 (chunks must not coalesce)", writesAfterFirst) + } + if !strings.Contains(bodyAfterFirst, "chunk-1") || strings.Contains(bodyAfterFirst, "chunk-2") { + t.Fatalf("body after first chunk = %q, want chunk-1 only", bodyAfterFirst) + } + + close(continueCh) + select { + case <-done: + case <-time.After(5 * time.Second): + t.Fatal("timed out waiting for ServeHTTP to finish") + } + + status, writes, flushes, body := sw.snapshot() + if status != http.StatusOK { + t.Fatalf("status = %d, want 200", status) + } + if writes < 2 { + t.Fatalf("want >=2 writes (one per upstream chunk), got %d", writes) + } + if flushes < writes { + t.Fatalf("want a flush per write (incremental delivery), got %d writes and %d flushes", writes, flushes) + } + if !strings.Contains(body, "chunk-1") || !strings.Contains(body, "chunk-2") { + t.Errorf("body missing streamed chunks, got %q", body) + } +} + +// --- authz and validation checks --- + +func TestHandleTargetProxy_UnknownName_Returns404(t *testing.T) { + cfg := CreateConfig() + cfg.Providers = map[string]*ProviderConfig{"openai": {Type: "openai", APIKey: "k"}} + cfg.MCPServers = map[string]*TargetConfig{"alpha": {URL: "http://mcp-alpha.internal"}} + cfg.Agents = map[string]*AgentConfig{"agent1": {URL: "http://agent1.internal"}} + cfg.Groups = map[string]*GroupConfig{"default": {}} + cfg.Users = &UsersConfig{Inline: []*UserConfig{{Name: "alice", Group: "default", APIKey: "sk-alice"}}} + next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}) + h, err := New(context.Background(), next, cfg, "llmgw") + if err != nil { + t.Fatalf("New: %v", err) + } + + for _, path := range []string{"/mcp/nosuchserver/tools/list", "/a2a/nosuchagent/tasks/send"} { + t.Run(path, func(t *testing.T) { + req := httptest.NewRequest(http.MethodGet, path, nil) + req.Header.Set("Authorization", "Bearer sk-alice") + rec := httptest.NewRecorder() + h.ServeHTTP(rec, req) + if rec.Code != http.StatusNotFound { + t.Fatalf("status = %d, want 404, body=%s", rec.Code, rec.Body.String()) + } + }) + } +} + +func TestHandleTargetProxy_Unauthenticated_Returns401(t *testing.T) { + cfg := CreateConfig() + cfg.Providers = map[string]*ProviderConfig{"openai": {Type: "openai", APIKey: "k"}} + cfg.MCPServers = map[string]*TargetConfig{"alpha": {URL: "http://mcp-alpha.internal"}} + next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}) + h, err := New(context.Background(), next, cfg, "llmgw") + if err != nil { + t.Fatalf("New: %v", err) + } + + // Even a name that does not exist must be rejected as 401 before 404 — + // auth runs before the unknown-target check. + req := httptest.NewRequest(http.MethodGet, "/mcp/nosuchserver/tools/list", nil) + rec := httptest.NewRecorder() + h.ServeHTTP(rec, req) + if rec.Code != http.StatusUnauthorized { + t.Fatalf("status = %d, want 401, body=%s", rec.Code, rec.Body.String()) + } +} + +func TestHandleTargetProxy_GroupDenies_Returns403(t *testing.T) { + cfg := CreateConfig() + cfg.Providers = map[string]*ProviderConfig{"openai": {Type: "openai", APIKey: "k"}} + cfg.MCPServers = map[string]*TargetConfig{ + "alpha": {URL: "http://mcp-alpha.internal"}, + "beta": {URL: "http://mcp-beta.internal"}, + } + cfg.Groups = map[string]*GroupConfig{"limited": {MCPServers: []string{"alpha"}}} + cfg.Users = &UsersConfig{Inline: []*UserConfig{{Name: "alice", Group: "limited", APIKey: "sk-alice"}}} + next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}) + h, err := New(context.Background(), next, cfg, "llmgw") + if err != nil { + t.Fatalf("New: %v", err) + } + + req := httptest.NewRequest(http.MethodGet, "/mcp/beta/tools/list", nil) + req.Header.Set("Authorization", "Bearer sk-alice") + rec := httptest.NewRecorder() + h.ServeHTTP(rec, req) + + if rec.Code != http.StatusForbidden { + t.Fatalf("status = %d, want 403, body=%s", rec.Code, rec.Body.String()) + } +} + +func TestHandleTargetProxy_TraversalPath_Returns400(t *testing.T) { + upstreamCalled := false + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + upstreamCalled = true + w.WriteHeader(http.StatusOK) + })) + defer srv.Close() + + cfg := CreateConfig() + cfg.Providers = map[string]*ProviderConfig{"openai": {Type: "openai", APIKey: "k"}} + cfg.MCPServers = map[string]*TargetConfig{"alpha": {URL: srv.URL}} + cfg.Groups = map[string]*GroupConfig{"default": {}} + cfg.Users = &UsersConfig{Inline: []*UserConfig{{Name: "alice", Group: "default", APIKey: "sk-alice"}}} + next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}) + h, err := New(context.Background(), next, cfg, "llmgw") + if err != nil { + t.Fatalf("New: %v", err) + } + + for _, path := range []string{ + "/mcp/alpha/../secret", + "/mcp/alpha/..%2f..%2fsecret", + } { + t.Run(path, func(t *testing.T) { + req := httptest.NewRequest(http.MethodGet, path, nil) + req.Header.Set("Authorization", "Bearer sk-alice") + rec := httptest.NewRecorder() + h.ServeHTTP(rec, req) + if rec.Code != http.StatusBadRequest { + t.Fatalf("status = %d, want 400, body=%s", rec.Code, rec.Body.String()) + } + }) + } + if upstreamCalled { + t.Error("upstream must never be called for a rejected traversal path") + } +} + +func TestHandleTargetProxy_UpgradeHeader_Returns501(t *testing.T) { + cfg := CreateConfig() + cfg.Providers = map[string]*ProviderConfig{"openai": {Type: "openai", APIKey: "k"}} + cfg.MCPServers = map[string]*TargetConfig{"alpha": {URL: "http://mcp-alpha.internal"}} + cfg.Groups = map[string]*GroupConfig{"default": {}} + cfg.Users = &UsersConfig{Inline: []*UserConfig{{Name: "alice", Group: "default", APIKey: "sk-alice"}}} + next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}) + h, err := New(context.Background(), next, cfg, "llmgw") + if err != nil { + t.Fatalf("New: %v", err) + } + + req := httptest.NewRequest(http.MethodGet, "/mcp/alpha/realtime", nil) + req.Header.Set("Authorization", "Bearer sk-alice") + req.Header.Set("Upgrade", "websocket") + req.Header.Set("Connection", "Upgrade") + rec := httptest.NewRecorder() + h.ServeHTTP(rec, req) + + if rec.Code != http.StatusNotImplemented { + t.Fatalf("status = %d, want 501, body=%s", rec.Code, rec.Body.String()) + } +} + +// --- empty rest proxies to target root --- + +func TestHandleTargetProxy_EmptyRest_ProxiesToRoot(t *testing.T) { + var gotPath string + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + gotPath = r.URL.Path + w.WriteHeader(http.StatusOK) + })) + defer srv.Close() + + cfg := CreateConfig() + cfg.Providers = map[string]*ProviderConfig{"openai": {Type: "openai", APIKey: "k"}} + cfg.MCPServers = map[string]*TargetConfig{"alpha": {URL: srv.URL}} + cfg.Groups = map[string]*GroupConfig{"default": {}} + cfg.Users = &UsersConfig{Inline: []*UserConfig{{Name: "alice", Group: "default", APIKey: "sk-alice"}}} + next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}) + h, err := New(context.Background(), next, cfg, "llmgw") + if err != nil { + t.Fatalf("New: %v", err) + } + + req := httptest.NewRequest(http.MethodGet, "/mcp/alpha", nil) + req.Header.Set("Authorization", "Bearer sk-alice") + rec := httptest.NewRecorder() + h.ServeHTTP(rec, req) + + if rec.Code != http.StatusOK { + t.Fatalf("status = %d, want 200, body=%s", rec.Code, rec.Body.String()) + } + if gotPath != "/" { + t.Errorf("upstream path = %q, want %q (empty rest proxies to target root)", gotPath, "/") + } +} + +// --- dead upstream maps to 502, same as native passthrough --- + +func TestHandleTargetProxy_DeadUpstream_Returns502(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {})) + deadURL := srv.URL + srv.Close() + + cfg := CreateConfig() + cfg.Providers = map[string]*ProviderConfig{"openai": {Type: "openai", APIKey: "k"}} + cfg.MCPServers = map[string]*TargetConfig{"alpha": {URL: deadURL}} + cfg.Groups = map[string]*GroupConfig{"default": {}} + cfg.Users = &UsersConfig{Inline: []*UserConfig{{Name: "alice", Group: "default", APIKey: "sk-alice"}}} + next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}) + h, err := New(context.Background(), next, cfg, "llmgw") + if err != nil { + t.Fatalf("New: %v", err) + } + + req := httptest.NewRequest(http.MethodGet, "/mcp/alpha/tools/list", nil) + req.Header.Set("Authorization", "Bearer sk-alice") + rec := httptest.NewRecorder() + h.ServeHTTP(rec, req) + + if rec.Code != http.StatusBadGateway { + t.Fatalf("status = %d, want 502, body=%s", rec.Code, rec.Body.String()) + } +} + +// --- a2a proxies distinctly from mcp --- + +func TestHandleTargetProxy_A2A_ProxiesToAgent(t *testing.T) { + var gotPath string + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + gotPath = r.URL.Path + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte("a2a-ok")) + })) + defer srv.Close() + + cfg := CreateConfig() + cfg.Providers = map[string]*ProviderConfig{"openai": {Type: "openai", APIKey: "k"}} + cfg.Agents = map[string]*AgentConfig{"agent1": {URL: srv.URL}} + cfg.Groups = map[string]*GroupConfig{"default": {}} + cfg.Users = &UsersConfig{Inline: []*UserConfig{{Name: "alice", Group: "default", APIKey: "sk-alice"}}} + next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}) + h, err := New(context.Background(), next, cfg, "llmgw") + if err != nil { + t.Fatalf("New: %v", err) + } + + req := httptest.NewRequest(http.MethodPost, "/a2a/agent1/tasks/send", nil) + req.Header.Set("Authorization", "Bearer sk-alice") + rec := httptest.NewRecorder() + h.ServeHTTP(rec, req) + + if rec.Code != http.StatusOK { + t.Fatalf("status = %d, want 200, body=%s", rec.Code, rec.Body.String()) + } + if rec.Body.String() != "a2a-ok" { + t.Errorf("body = %q, want verbatim upstream body", rec.Body.String()) + } + if gotPath != "/tasks/send" { + t.Errorf("upstream path = %q, want /tasks/send", gotPath) + } +} + +// --- config validation at construction --- + +func TestNewGateway_TargetURLValidation(t *testing.T) { + tests := []struct { + mcpServers map[string]*TargetConfig + agents map[string]*AgentConfig + name string + wantErr bool + }{ + { + name: "valid mcp server url", + mcpServers: map[string]*TargetConfig{"alpha": {URL: "http://mcp-alpha.internal:8080"}}, + wantErr: false, + }, + { + name: "valid https agent url", + agents: map[string]*AgentConfig{"agent1": {URL: "https://agent1.internal"}}, + wantErr: false, + }, + { + name: "malformed mcp url (control character)", + mcpServers: map[string]*TargetConfig{"alpha": {URL: "http://exa\x7fmple.com"}}, + wantErr: true, + }, + { + name: "disallowed scheme for mcp server", + mcpServers: map[string]*TargetConfig{"alpha": {URL: "ftp://mcp-alpha.internal"}}, + wantErr: true, + }, + { + name: "disallowed scheme for agent", + agents: map[string]*AgentConfig{"agent1": {URL: "ws://agent1.internal"}}, + wantErr: true, + }, + { + name: "malformed agent url (invalid port)", + agents: map[string]*AgentConfig{"agent1": {URL: "http://agent1.internal:notaport"}}, + wantErr: true, + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + cfg := CreateConfig() + cfg.Providers = map[string]*ProviderConfig{"openai": {Type: "openai", APIKey: "k"}} + cfg.MCPServers = tt.mcpServers + cfg.Agents = tt.agents + next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}) + h, err := New(context.Background(), next, cfg, "llmgw") + if tt.wantErr { + if err == nil { + t.Fatal("want error, got nil") + } + if h != nil { + t.Fatal("want nil handler on error") + } + return + } + if err != nil { + t.Fatalf("New: %v", err) + } + }) + } +} diff --git a/routes_passthrough.go b/routes_passthrough.go index eb60a8d..bf941a0 100644 --- a/routes_passthrough.go +++ b/routes_passthrough.go @@ -281,27 +281,95 @@ func (g *Gateway) handlePassthrough(w http.ResponseWriter, r *http.Request, u *u upstreamURL += "?" + r.URL.RawQuery } + result, ok := g.proxyUpstream(w, r, upstreamURL, adapter.httpClient(), adapter.injectAuth, true, "passthrough (provider "+providerName+")") + if !ok || !result.isJSON { + // A build/connection/copy failure already wrote its own response + // (or, for a canceled client context, wrote nothing at all — see + // proxyUpstream); a non-JSON response has nothing more to account + // than the request itself, already counted by checkAndCount above. + return + } + if result.tee.truncated { + g.logf("passthrough: response body exceeded %d bytes; skipping usage accounting (provider %q)", maxAccountingTeeBytes, providerName) + return + } + + respUsage, model, unmarshalErr := extractPassthroughUsage(adapter.typeName(), providerName, result.tee.buf.Bytes()) + if unmarshalErr != nil { + g.logf("passthrough: response body did not decode as JSON for usage accounting (provider %q): %v", providerName, unmarshalErr) + } + cost := unifiedCostMicros(providerName+"/"+model, model, respUsage, g.cfg.Pricing) + g.limiter.account(scopes, respUsage, cost) +} + +// proxyResult is what proxyUpstream reports back to its caller once it has +// streamed a response to the client: enough for a caller that wants +// best-effort JSON usage accounting (handlePassthrough) to run it, without +// proxyUpstream itself knowing anything about usage or pricing. tee is +// nil unless accountJSON was true and the response's Content-Type was +// application/json. +type proxyResult struct { + tee *cappedAccountingBuffer + isJSON bool +} + +// proxyUpstream is the shared reverse-proxy core behind both native +// provider passthrough (handlePassthrough, above) and the MCP/A2A target +// proxy (handleTargetProxy, mcp_a2a.go): it builds an upstream request +// from r's method/body/headers, sends it over client, and streams the +// response back to w incrementally through a flushWriter — so an SSE or +// other chunked upstream body reaches the client as each chunk arrives, +// never buffered until the copy finishes. +// +// Every hop-by-hop header (hopByHopHeaders) and the client's own gateway +// credential (gatewayCredentialHeaders) are stripped from the outgoing +// request, and Accept-Encoding (clientNegotiationHeaders) besides, so +// Transport can negotiate and transparently decompress compression +// itself. injectAuth, when non-nil, is called on the built request before +// it is sent — handlePassthrough passes its adapter's injectAuth to swap +// the client's key for the provider's own; handleTargetProxy passes nil, +// since an MCP server or A2A agent is an in-cluster target that receives +// no injected credential at all. +// +// A build failure or a dead upstream writes a 502 envelope to w and +// returns ok=false; a canceled client context (errors.Is context.Canceled) +// is logged, not surfaced, and also returns ok=false, writing nothing — +// the client is already gone. A response-copy failure after headers were +// already written also returns ok=false, since nothing further can be +// written to w at that point either way. logPrefix labels every +// logf/errorf line this call emits, so passthrough and target-proxy +// failures stay distinguishable in the log. +// +// accountJSON gates whether a non-streaming application/json response +// gets teed off into result.tee for the caller's own best-effort usage +// parse afterward: handlePassthrough passes true; handleTargetProxy +// passes false, since target-proxy accounting never goes past the +// request-count checkAndCount already ran before calling in, and teeing a +// response nobody will ever read back would only cost memory for nothing. +func (g *Gateway) proxyUpstream(w http.ResponseWriter, r *http.Request, upstreamURL string, client *http.Client, injectAuth func(*http.Request), accountJSON bool, logPrefix string) (result proxyResult, ok bool) { bodyReader := io.LimitReader(r.Body, maxPassthroughBytes) // gosec G704 (SSRF via taint analysis) flags upstreamURL as // request-derived: it is, by design — this is a reverse proxy, and its // whole job is to forward a client-supplied path/query onto the - // upstream. The host component is never request-derived: adapter.base() - // is a fixed, operator-configured URL from Config.Providers, and - // upstreamURL is built as base()+"/"+rest — string concatenation, not a - // second url.Parse of rest alone — so no client-supplied value can - // change the scheme or host url.Parse resolves for the final string. - // rest is additionally traversal-checked above (hasTraversalSegment - // rejects any ".." segment in its decoded form before this point is - // reached) and is forwarded here in its original escaped form — see - // passthroughRoute's EscapedPath()-based split in llmgateway.go — so a - // client cannot smuggle an encoded "/" (%2f) past this gateway's own + // upstream. The host component is never request-derived: every caller + // builds upstreamURL as an operator-configured base (a provider's + // adapter.base(), or an MCP-server/agent TargetConfig.URL validated at + // construction — see validateTargetURLs) plus "/"+rest — string + // concatenation, not a second url.Parse of rest alone — so no + // client-supplied value can change the scheme or host url.Parse + // resolves for the final string. rest is additionally + // traversal-checked by every caller before this point is reached + // (hasTraversalSegment rejects any ".." segment in its decoded form) + // and is forwarded here in its original escaped form — see + // passthroughRoute's and targetRoute's EscapedPath()-based splits — so + // a client cannot smuggle an encoded "/" (%2f) past this gateway's own // routing only to have a permissive upstream reinterpret it as a real // separator. upstreamReq, err := http.NewRequestWithContext(r.Context(), r.Method, upstreamURL, bodyReader) //nolint:gosec // operator-fixed host, rest is traversal-checked and forwarded escaped; see comment above if err != nil { - g.errorf("passthrough: build upstream request (provider %q): %v", providerName, err) + g.errorf("%s: build upstream request: %v", logPrefix, err) writeOAIError(w, http.StatusBadGateway, "server_error", "upstream connection error") - return + return proxyResult{}, false } upstreamReq.ContentLength = r.ContentLength if r.ContentLength < 0 || r.ContentLength > maxPassthroughBytes { @@ -311,17 +379,19 @@ func (g *Gateway) handlePassthrough(w http.ResponseWriter, r *http.Request, u *u upstreamReq.ContentLength = -1 } copyHeadersExcept(upstreamReq.Header, r.Header, hopByHopHeaders, gatewayCredentialHeaders, clientNegotiationHeaders) - adapter.injectAuth(upstreamReq) + if injectAuth != nil { + injectAuth(upstreamReq) + } - resp, err := adapter.httpClient().Do(upstreamReq) //nolint:gosec // same upstreamReq built above; operator-fixed host, traversal-checked, see its construction comment + resp, err := client.Do(upstreamReq) //nolint:gosec // same upstreamReq built above; operator-fixed host, traversal-checked, see its construction comment if err != nil { if errors.Is(err, context.Canceled) { - g.logf("passthrough: client canceled request to provider %q: %v", providerName, err) - return + g.logf("%s: client canceled request: %v", logPrefix, err) + return proxyResult{}, false } - g.errorf("passthrough: upstream connection error (provider %q): %v", providerName, err) + g.errorf("%s: upstream connection error: %v", logPrefix, err) writeOAIError(w, http.StatusBadGateway, "server_error", "upstream connection error") - return + return proxyResult{}, false } defer resp.Body.Close() //nolint:errcheck // read-side close; nothing actionable on failure @@ -330,10 +400,11 @@ func (g *Gateway) handlePassthrough(w http.ResponseWriter, r *http.Request, u *u fw := newFlushWriter(w) // Only a non-streaming JSON response gets teed off for best-effort - // accounting; every other response streams straight through fw with no - // extra buffering — a large or slow upstream body must never sit - // waiting for full receipt before the client sees its first byte. - isJSON := strings.Contains(strings.ToLower(resp.Header.Get("Content-Type")), "application/json") + // accounting, and only when the caller wants that at all (accountJSON); + // every other response streams straight through fw with no extra + // buffering — a large or slow upstream body must never sit waiting for + // full receipt before the client sees its first byte. + isJSON := accountJSON && strings.Contains(strings.ToLower(resp.Header.Get("Content-Type")), "application/json") var tee *cappedAccountingBuffer var reader io.Reader = resp.Body if isJSON { @@ -342,24 +413,8 @@ func (g *Gateway) handlePassthrough(w http.ResponseWriter, r *http.Request, u *u } if _, err := io.Copy(fw, reader); err != nil { - g.errorf("passthrough: stream response body (provider %q): %v", providerName, err) - return + g.errorf("%s: stream response body: %v", logPrefix, err) + return proxyResult{isJSON: isJSON, tee: tee}, false } - if !isJSON { - // Streaming or non-JSON: usage is unknowable from here, so only the - // request itself gets accounted — already done by checkAndCount - // above. - return - } - if tee.truncated { - g.logf("passthrough: response body exceeded %d bytes; skipping usage accounting (provider %q)", maxAccountingTeeBytes, providerName) - return - } - - respUsage, model, unmarshalErr := extractPassthroughUsage(adapter.typeName(), providerName, tee.buf.Bytes()) - if unmarshalErr != nil { - g.logf("passthrough: response body did not decode as JSON for usage accounting (provider %q): %v", providerName, unmarshalErr) - } - cost := unifiedCostMicros(providerName+"/"+model, model, respUsage, g.cfg.Pricing) - g.limiter.account(scopes, respUsage, cost) + return proxyResult{isJSON: isJSON, tee: tee}, true }