Reject misrouted gateway credentials

Author: lawrencecchen

internal/proxy/api_key_gateway.go

@@ -45,6 +45,66 @@ type apiKeyGatewaySpec struct {
 	blockedAPIKeyPrefixes []string
 }
 
+func (s Server) rejectMisroutedGatewayCredentials(next http.Handler) http.Handler {
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if s.requestUsesConfiguredGatewayCredential(r) {
+			http.Error(w, "gateway credential used outside its gateway route", http.StatusBadRequest)
+			return
+		}
+		next.ServeHTTP(w, r)
+	})
+}
+
+func (s Server) requestUsesConfiguredGatewayCredential(r *http.Request) bool {
+	if r == nil {
+		return false
+	}
+	var configured []string
+	if s.Gemini != nil {
+		configured = append(configured, strings.TrimSpace(s.Gemini.GatewayToken))
+	}
+	if s.AnthropicGateway != nil {
+		configured = append(configured, strings.TrimSpace(s.AnthropicGateway.GatewayToken))
+	}
+	if s.OpenAIGateway != nil {
+		configured = append(configured, strings.TrimSpace(s.OpenAIGateway.GatewayToken))
+	}
+	if s.Bedrock != nil {
+		configured = append(configured, strings.TrimSpace(s.Bedrock.GatewayToken))
+	}
+	var presented []string
+	presented = append(presented, r.Header.Values("X-Goog-Api-Key")...)
+	presented = append(presented, r.Header.Values("X-Api-Key")...)
+	for _, header := range r.Header.Values("Authorization") {
+		if scheme, value, ok := strings.Cut(strings.TrimSpace(header), " "); ok && strings.EqualFold(scheme, "Bearer") {
+			presented = append(presented, strings.TrimSpace(value))
+		}
+	}
+	for _, value := range r.Header.Values("Sec-WebSocket-Protocol") {
+		for _, protocol := range strings.Split(value, ",") {
+			protocol = strings.TrimSpace(protocol)
+			if strings.HasPrefix(protocol, openAIWebSocketCredentialPrefix) {
+				presented = append(presented, strings.TrimPrefix(protocol, openAIWebSocketCredentialPrefix))
+			}
+		}
+	}
+	query := r.URL.Query()
+	presented = append(presented, query["key"]...)
+	presented = append(presented, query["api_key"]...)
+	for _, candidate := range presented {
+		candidate = strings.TrimSpace(candidate)
+		if candidate == "" {
+			continue
+		}
+		for _, token := range configured {
+			if len(candidate) == len(token) && token != "" && subtle.ConstantTimeCompare([]byte(candidate), []byte(token)) == 1 {
+				return true
+			}
+		}
+	}
+	return false
+}
+
 func (s Server) apiKeyGatewayHandler(config *APIKeyGatewayConfig, spec apiKeyGatewaySpec) http.Handler {
 	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		if config == nil || !config.configured() {

internal/proxy/api_key_gateway_test.go

@@ -666,3 +666,50 @@ func TestAPIKeyGatewayRoutesDoNotCaptureRootProxyPaths(t *testing.T) {
 		t.Fatalf("gateway upstream requests = %d", got)
 	}
 }
+
+func TestRootProxyRejectsMisroutedGatewayCredentials(t *testing.T) {
+	t.Parallel()
+
+	var rootRequests atomic.Int32
+	root := httptest.NewServer(http.HandlerFunc(func(http.ResponseWriter, *http.Request) {
+		rootRequests.Add(1)
+	}))
+	defer root.Close()
+	rootURL, err := url.Parse(root.URL)
+	if err != nil {
+		t.Fatal(err)
+	}
+	handler := Server{
+		Upstream:         rootURL,
+		Gemini:           &GeminiConfig{GatewayToken: "gemini-team"},
+		AnthropicGateway: &APIKeyGatewayConfig{GatewayToken: "anthropic-team"},
+		OpenAIGateway:    &APIKeyGatewayConfig{GatewayToken: "openai-team"},
+	}.Handler()
+	tests := []struct {
+		name string
+		path string
+		set  func(http.Header)
+	}{
+		{"gemini header", "/v1/models", func(h http.Header) { h.Set("X-Goog-Api-Key", "gemini-team") }},
+		{"gemini query", "/v1alpha/models?key=gemini-team", func(http.Header) {}},
+		{"anthropic header", "/v1/messages", func(h http.Header) { h.Set("X-Api-Key", "anthropic-team") }},
+		{"openai bearer", "/v1/responses", func(h http.Header) { h.Set("Authorization", "Bearer openai-team") }},
+		{"openai websocket", "/v1/realtime", func(h http.Header) {
+			h.Set("Sec-WebSocket-Protocol", "realtime, openai-insecure-api-key.openai-team")
+		}},
+	}
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			req := httptest.NewRequest(http.MethodPost, test.path, nil)
+			test.set(req.Header)
+			rec := httptest.NewRecorder()
+			handler.ServeHTTP(rec, req)
+			if rec.Code != http.StatusBadRequest {
+				t.Fatalf("status = %d body = %s", rec.Code, rec.Body.String())
+			}
+		})
+	}
+	if got := rootRequests.Load(); got != 0 {
+		t.Fatalf("misrouted credentials reached root upstream: requests = %d", got)
+	}
+}

internal/proxy/proxy.go

@@ -951,7 +951,7 @@ func (s Server) Handler() http.Handler {
 	mux.Handle("/api/v1/", openAIHandler)
 	mux.Handle("/openai", openAIHandler)
 	mux.Handle("/openai/", openAIHandler)
-	mux.Handle("/", s.proxyHandler())
+	mux.Handle("/", s.rejectMisroutedGatewayCredentials(s.proxyHandler()))
 	return mux
 }