diff --git a/src/invocation-plane-services/llm-api-gateway/README.md b/src/invocation-plane-services/llm-api-gateway/README.md index 100f98ca7..d07385e42 100644 --- a/src/invocation-plane-services/llm-api-gateway/README.md +++ b/src/invocation-plane-services/llm-api-gateway/README.md @@ -52,6 +52,9 @@ Each request is normalized into a function-scoped request context. and is forwarded to NVCF gRPC auth when that adapter is configured. - `X-Request-ID` is accepted if present, otherwise the gateway generates one. - `X-NVCF-Target-Region` is forwarded into the request context. +- `X-Priority` is reserved for the gateway and derived from the caller's + resolved priority; a client-supplied `X-Priority` on the LLM endpoints is + rejected with 400. Configured functions control the downstream `model`, service tier, routing method, and per-function rate limits. Prompt rendering and exact prompt diff --git a/src/invocation-plane-services/llm-api-gateway/api/BUILD.bazel b/src/invocation-plane-services/llm-api-gateway/api/BUILD.bazel index 3f47bef11..6863f3355 100644 --- a/src/invocation-plane-services/llm-api-gateway/api/BUILD.bazel +++ b/src/invocation-plane-services/llm-api-gateway/api/BUILD.bazel @@ -28,6 +28,7 @@ go_library( "openai_endpoint_handlers.go", "openai_remaining_handlers.go", "openai_surface.go", + "priority_guard.go", "provider_error.go", "proxy_handler.go", "ratelimit.go", @@ -84,6 +85,7 @@ go_test( "middleware_test.go", "model_ids_test.go", "openai_routes_test.go", + "priority_guard_test.go", "rate_limit_accounting_test.go", "responses_handler_test.go", "session_affinity_test.go", diff --git a/src/invocation-plane-services/llm-api-gateway/api/priority_guard.go b/src/invocation-plane-services/llm-api-gateway/api/priority_guard.go new file mode 100644 index 000000000..372161194 --- /dev/null +++ b/src/invocation-plane-services/llm-api-gateway/api/priority_guard.go @@ -0,0 +1,44 @@ +/* +SPDX-FileCopyrightText: Copyright (c) NVIDIA CORPORATION & AFFILIATES. All rights reserved. +SPDX-License-Identifier: Apache-2.0 + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +package api + +import ( + "fmt" + "net/http" + + echo "github.com/labstack/echo/v4" +) + +const headerPriority = "X-Priority" + +// rejectClientSuppliedPriority rejects requests that arrive with a +// client-supplied X-Priority header. +// Registered on the LLM route group, not globally. The check fires on header presence +// rather than value so a present-but-empty X-Priority (which Header.Get cannot +// distinguish from absent) is rejected too. +func rejectClientSuppliedPriority(next echo.HandlerFunc) echo.HandlerFunc { + return func(ec echo.Context) error { + if len(ec.Request().Header.Values(headerPriority)) > 0 { + return echo.NewHTTPError( + http.StatusBadRequest, + fmt.Sprintf("%s is a reserved header and cannot be set by the client", headerPriority), + ) + } + return next(ec) + } +} diff --git a/src/invocation-plane-services/llm-api-gateway/api/priority_guard_test.go b/src/invocation-plane-services/llm-api-gateway/api/priority_guard_test.go new file mode 100644 index 000000000..cf69d2fce --- /dev/null +++ b/src/invocation-plane-services/llm-api-gateway/api/priority_guard_test.go @@ -0,0 +1,179 @@ +/* +SPDX-FileCopyrightText: Copyright (c) NVIDIA CORPORATION & AFFILIATES. All rights reserved. +SPDX-License-Identifier: Apache-2.0 + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +package api + +import ( + "net/http" + "net/http/httptest" + "strings" + "testing" + + echo "github.com/labstack/echo/v4" + + "github.com/NVIDIA/nvcf/src/invocation-plane-services/llm-gateway/config" + "github.com/NVIDIA/nvcf/src/invocation-plane-services/llm-gateway/nvcf" +) + +// newPriorityGuardTestServer wires the production route registration so these +// tests fail if the guard is ever dropped from the LLM route group. +func newPriorityGuardTestServer(authClient InvocationAuthClient) *echo.Echo { + cfg := config.Default() + e := echo.New() + e.Use(NewContextMiddleware(cfg)) + e.Use(NewNVCFAuthMiddleware(authClient)) + RegisterRoutes(e, NewHandlers(cfg, nil, nil)) + return e +} + +func TestLLMRoutesRejectClientSuppliedPriority(t *testing.T) { + t.Parallel() + + paths := []string{"/v1/chat/completions", "/v1/responses", "/v1/embeddings"} + values := []struct { + name string + priority string + }{ + {name: "explicit value", priority: "0"}, + {name: "empty value", priority: ""}, + } + + for _, path := range paths { + for _, value := range values { + t.Run(path+" "+value.name, func(t *testing.T) { + t.Parallel() + + authClient := &stubInvocationAuthClient{ + authResponse: &nvcf.InvocationAuthResponse{ + RoutingKey: "fn-chat", + ClientAuthID: "subject-123", + AuthContext: map[string]string{"ncaId": "nca-456"}, + RateLimitKey: "nca-456", + }, + } + + e := newPriorityGuardTestServer(authClient) + + req := httptest.NewRequest( + http.MethodPost, + path, + strings.NewReader(`{"model":"fn-chat/company-name/model-name","messages":[{"role":"user","content":"hello"}]}`), + ) + req.Header.Set(echo.HeaderContentType, echo.MIMEApplicationJSON) + req.Header.Set(echo.HeaderAuthorization, "Bearer sk-live") + req.Header.Set(headerPriority, value.priority) + rec := httptest.NewRecorder() + + e.ServeHTTP(rec, req) + + if rec.Code != http.StatusBadRequest { + t.Fatalf("status = %d, want %d: %s", rec.Code, http.StatusBadRequest, rec.Body.String()) + } + if !strings.Contains(rec.Body.String(), "reserved header") { + t.Fatalf("body = %q, want the reserved-header rejection message", rec.Body.String()) + } + }) + } + } + + // The guard is route middleware on the LLM group, so it runs after the + // global auth middleware: the reject does not skip authentication. + t.Run("auth runs before the reject", func(t *testing.T) { + t.Parallel() + + authClient := &stubInvocationAuthClient{ + authResponse: &nvcf.InvocationAuthResponse{ + RoutingKey: "fn-chat", + ClientAuthID: "subject-123", + AuthContext: map[string]string{"ncaId": "nca-456"}, + RateLimitKey: "nca-456", + }, + } + + e := newPriorityGuardTestServer(authClient) + + req := httptest.NewRequest( + http.MethodPost, + "/v1/chat/completions", + strings.NewReader(`{"model":"fn-chat/company-name/model-name","messages":[{"role":"user","content":"hello"}]}`), + ) + req.Header.Set(echo.HeaderContentType, echo.MIMEApplicationJSON) + req.Header.Set(echo.HeaderAuthorization, "Bearer sk-live") + req.Header.Set(headerPriority, "0") + rec := httptest.NewRecorder() + + e.ServeHTTP(rec, req) + + if rec.Code != http.StatusBadRequest { + t.Fatalf("status = %d, want %d: %s", rec.Code, http.StatusBadRequest, rec.Body.String()) + } + if !strings.Contains(rec.Body.String(), "reserved header") { + t.Fatalf("body = %q, want the reserved-header rejection message", rec.Body.String()) + } + if authClient.authorizeCalls != 1 { + t.Fatalf("authorize calls = %d, want 1", authClient.authorizeCalls) + } + }) + + // The guard is scoped to the LLM route group; non-proxied routes such as + // the health endpoints accept requests regardless of X-Priority. + t.Run("healthz not rejected", func(t *testing.T) { + t.Parallel() + + e := newPriorityGuardTestServer(&stubInvocationAuthClient{}) + + req := httptest.NewRequest(http.MethodGet, "/healthz", nil) + req.Header.Set(headerPriority, "1") + rec := httptest.NewRecorder() + + e.ServeHTTP(rec, req) + + if rec.Code != http.StatusOK { + t.Fatalf("status = %d, want %d: %s", rec.Code, http.StatusOK, rec.Body.String()) + } + }) +} + +// A request without X-Priority must pass the guard untouched and reach the next +// handler. Positive control: without this, a guard that rejected unconditionally +// would still satisfy every rejection case above while breaking all legitimate +// LLM traffic. The guard runs after authentication in the wired server, so its +// only remaining job on a clean request is to forward it, which this asserts +// directly rather than driving the full provider stack behind the LLM routes. +func TestRejectClientSuppliedPriorityForwardsRequestWithoutHeader(t *testing.T) { + t.Parallel() + + forwarded := false + next := func(c echo.Context) error { + forwarded = true + return c.NoContent(http.StatusNoContent) + } + + req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil) + rec := httptest.NewRecorder() + c := echo.New().NewContext(req, rec) + + if err := rejectClientSuppliedPriority(next)(c); err != nil { + t.Fatalf("guard returned error for a request without X-Priority: %v", err) + } + if !forwarded { + t.Fatal("guard did not forward a request without X-Priority to the next handler") + } + if rec.Code != http.StatusNoContent { + t.Fatalf("status = %d, want %d", rec.Code, http.StatusNoContent) + } +} diff --git a/src/invocation-plane-services/llm-api-gateway/api/routes.go b/src/invocation-plane-services/llm-api-gateway/api/routes.go index a35c5fa96..e016e00e4 100644 --- a/src/invocation-plane-services/llm-api-gateway/api/routes.go +++ b/src/invocation-plane-services/llm-api-gateway/api/routes.go @@ -31,7 +31,7 @@ func RegisterRoutes(e *echo.Echo, handlers *Handlers) { return c.NoContent(http.StatusOK) }) - group := e.Group("") + group := e.Group("", rejectClientSuppliedPriority) handlers.AsOpenAIChatHandlers().RegisterRoutes(group) handlers.AsResponsesHandlers().RegisterRoutes(group) // proxy handlers only contain embedding route for now, but could be extended to other routes in the future