diff --git a/docs/architecture.mdx b/docs/architecture.mdx index f7e7ffed..81c1db9a 100644 --- a/docs/architecture.mdx +++ b/docs/architecture.mdx @@ -364,6 +364,25 @@ without sending the request to an engine. It does not broaden the candidate list or refresh inventory synchronously; a later discovery update makes a newly advertised owner eligible. +The Ollama proxy also supports an opt-in exact-artifact capability gate. A +model-bearing inference request with one valid +`X-MrNiceAI-Expected-Artifact-Sha256` value triggers concurrent, bounded +`/api/tags` reads through each already eligible candidate's normal transport. +Only candidates with one unambiguous model record carrying that exact digest +remain eligible. This request-local intersection preserves selection and +scheduler order, happens before reservations or forwarding, and fails closed +when no exact owner can be proven. The expected value is consumed at the proxy. +Inventory redirects are not followed and process-wide concurrency is capped. +At the response commit point, the proxy checks the candidate's `/api/ps` +scheduler inventory and requires the requested model's loaded digest to equal +the preselected digest. A missing, ambiguous, unavailable, or changed live +record fails over before any response byte is released, or fails closed if the +candidate list is exhausted. Any engine-supplied header or trailer named +`X-MrNiceAI-Served-Artifact-Sha256` is removed and replaced with the immutable +verified digest for the candidate that actually served the response, including +after failover. Reserved request trailers are rejected or stripped. Requests +without the expected header retain ordinary behavior. + A request whose model cannot be parsed keeps the ordinary non-model ordering. Model listings are not routed at all. A `GET` of `/v1/models` or `/api/tags` is diff --git a/services/ollama-proxy/README.md b/services/ollama-proxy/README.md index 35f959b5..2520bcae 100644 --- a/services/ollama-proxy/README.md +++ b/services/ollama-proxy/README.md @@ -43,10 +43,13 @@ The proxy listens on `--port` (default 11435) and forwards incoming HTTP request A response *forwarded from an engine* is different: if the engine set its own `Access-Control-Allow-Origin` (Ollama with `OLLAMA_ORIGINS`), that header is passed through untouched rather than replaced with the wildcard, so a deliberately narrow engine policy is never silently widened and a credentialed response is not broken. An engine that sends no CORS header has expressed no policy to keep, so the proxy supplies its own — and drops any `Access-Control-Allow-Credentials` that arrived without an origin, because a browser rejects that header alongside a wildcard origin and would discard the response the fallback exists to make readable. +**Exact-artifact requests.** A model-bearing inference request may carry exactly one ordinary `X-MrNiceAI-Expected-Artifact-Sha256` header containing 64 lowercase hexadecimal characters; the same field is rejected as a trailer. The proxy consumes the header, reads `/api/tags` from the already model-eligible candidates through their normal local or mTLS transports, and retains only owners whose single matching model record has that digest. Redirects are not followed. A fixed worker pool plus a process-wide gate bound the inventory work; selection, scheduler priority, reservations, and failover keep their existing order within the retained set. Invalid bindings return `400`, an available but non-matching inventory returns `412`, and complete inventory unavailability returns `503`, all before inference. Before committing a candidate response, the proxy reads Ollama's scheduler-owned `/api/ps` inventory and requires the live loaded-model digest to match the selected digest. A missing, ambiguous, unavailable, or changed loaded record fails over before any response byte is released, or fails closed when candidates are exhausted. The proxy removes every engine-supplied `X-MrNiceAI-Served-Artifact-Sha256` header and trailer and writes exactly one header from the verified candidate inventory. The expected header never reaches an engine, a manual pin cannot bypass the digest gate, and requests without the expected header retain ordinary routing behavior while still having reserved engine evidence stripped. This binding currently applies to the Ollama proxy; other engine proxies must implement the same contract independently before they can serve an exact-artifact request. + One limit is outside the proxy's control: current Chromium-based browsers gate a request from a public origin to a local or loopback address behind the user's [Local Network Access](https://chromestatus.com/feature/5152728072060928) permission, which replaced the old server-side opt-in header. No header the proxy sends can grant that. A hosted page needs the permission plus a `fetch(url, { targetAddressSpace: 'loopback' })` annotation; a page served from the local machine is unaffected. Node selection: - **Eligibility**: Before routing model-bearing inference, the proxy keeps only nodes whose current Ollama inventory advertises the requested model. Ollama's implicit `:latest` tag is normalized. An empty or non-matching inventory is excluded until a later discovery update; if no advertised owner is routable, the proxy returns a local `502`. +- **Exact artifact**: When the optional expected-artifact header is present, a bounded request-local `/api/tags` read restricts eligibility to the exact digest and a pre-commit `/api/ps` read verifies the loaded runner. No matching installed and loaded owner means no response is committed. - **Auto**: When no eligible node is explicitly selected, the proxy follows `node/set-priority` (see below), then discovered nodes in stable ID order. - **Priority (scheduler-driven)**: The Job Scheduler ranks the cluster least-loaded-first by pending workload plus smoothed GPU pressure and, via `nvpair-ui-broker`, pushes the ordered node list with those per-node counts to this proxy with `node/set-priority`. Auto routing sends the request to the listed node carrying the least estimated load. See [`nvpair-job-scheduler`](../nvpair-job-scheduler/README.md). - **Manual**: Use the `node/select` JSON-RPC method to pin traffic to a specific node. A manual pin **overrides the priority list only when that node is eligible** for the requested model. diff --git a/services/ollama-proxy/artifact_attestation_test.go b/services/ollama-proxy/artifact_attestation_test.go new file mode 100644 index 00000000..ed6b6967 --- /dev/null +++ b/services/ollama-proxy/artifact_attestation_test.go @@ -0,0 +1,377 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +package main + +import ( + "io" + "net/http" + "net/http/httptest" + "strings" + "sync/atomic" + "testing" + "time" +) + +const ( + testArtifactA = "17052f91a42e97930aa6e28a6c6c06a983e6a58dbb00434885a0cf5313e376f7" + testArtifactB = "27052f91a42e97930aa6e28a6c6c06a983e6a58dbb00434885a0cf5313e376f7" +) + +func attestedRequest(body string, digest string) *http.Request { + request := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(body)) + request.Header.Set(expectedArtifactSHA256Header, digest) + return request +} + +func writeArtifactInventory(w http.ResponseWriter, digest string) { + w.Header().Set("Content-Type", "application/json") + _, _ = io.WriteString(w, `{"models":[{"name":"gpt-oss:20b","model":"gpt-oss:20b","digest":"`+digest+`"}]}`) +} + +func TestHandleHTTP_ExactArtifactFiltersBeforeSelectionAndReplacesEngineEvidence(t *testing.T) { + var mismatchInference atomic.Int32 + mismatch := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path == "/api/tags" { + _, _ = io.WriteString(w, `{"models":[{"name":"gpt-oss:20b","model":"gpt-oss:20b","digest":"`+testArtifactB+`"}]}`) + return + } + mismatchInference.Add(1) + w.WriteHeader(http.StatusOK) + })) + defer mismatch.Close() + + var matchInference atomic.Int32 + match := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path == "/api/tags" || r.URL.Path == "/api/ps" { + if got := headerValues(r.Header, expectedArtifactSHA256Header); len(got) != 0 { + t.Errorf("expected-artifact header leaked to inventory: %v", got) + } + writeArtifactInventory(w, testArtifactA) + return + } + matchInference.Add(1) + if got := headerValues(r.Header, expectedArtifactSHA256Header); len(got) != 0 { + t.Errorf("expected-artifact header leaked to engine: %v", got) + } + w.Header().Add(servedArtifactSHA256Header, testArtifactB) + w.Header().Add(servedArtifactSHA256Header, testArtifactB) + _, _ = io.WriteString(w, `{"ok":true}`) + })) + defer match.Close() + + discovery := NewDiscovery() + discovery.AddManual(nodeForModel(t, "mismatch", mismatch.URL, "gpt-oss:20b")) + discovery.AddManual(nodeForModel(t, "match", match.URL, "gpt-oss:20b")) + proxy := testProxy(discovery, 11434) + proxy.SetSelected("mismatch") + recorder := httptest.NewRecorder() + + proxy.handleHTTP(recorder, attestedRequest(`{"model":"gpt-oss:20b"}`, testArtifactA)) + + if recorder.Code != http.StatusOK { + t.Fatalf("status = %d, want 200; body=%s", recorder.Code, recorder.Body.String()) + } + if got := mismatchInference.Load(); got != 0 { + t.Errorf("digest-mismatched selected node received %d inference requests, want 0", got) + } + if got := matchInference.Load(); got != 1 { + t.Errorf("matching node received %d inference requests, want 1", got) + } + if got := headerValues(recorder.Header(), servedArtifactSHA256Header); len(got) != 1 || got[0] != testArtifactA { + t.Fatalf("served artifact values = %v, want exactly [%s]", got, testArtifactA) + } +} + +func TestHandleHTTP_ExactArtifactFailoverKeepsFinalCandidateEvidence(t *testing.T) { + server := func(status int, engineEvidence string) *httptest.Server { + return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path == "/api/tags" || r.URL.Path == "/api/ps" { + writeArtifactInventory(w, testArtifactA) + return + } + w.Header().Set(servedArtifactSHA256Header, engineEvidence) + w.WriteHeader(status) + _, _ = io.WriteString(w, `{"done":true}`) + })) + } + busy := server(http.StatusServiceUnavailable, testArtifactB) + defer busy.Close() + good := server(http.StatusOK, testArtifactB) + defer good.Close() + + discovery := NewDiscovery() + discovery.AddManual(nodeForModel(t, "busy", busy.URL, "gpt-oss:20b")) + discovery.AddManual(nodeForModel(t, "good", good.URL, "gpt-oss:20b")) + proxy := testProxy(discovery, 11434) + proxy.SetSelected("busy") + recorder := httptest.NewRecorder() + + proxy.handleHTTP(recorder, attestedRequest(`{"model":"gpt-oss:20b"}`, testArtifactA)) + + if recorder.Code != http.StatusOK { + t.Fatalf("status = %d, want 200 after failover; body=%s", recorder.Code, recorder.Body.String()) + } + if got := headerValues(recorder.Header(), servedArtifactSHA256Header); len(got) != 1 || got[0] != testArtifactA { + t.Fatalf("served artifact values = %v, want final candidate digest %s", got, testArtifactA) + } +} + +func TestHandleHTTP_ExactArtifactMismatchFailsClosed(t *testing.T) { + var inference atomic.Int32 + engine := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path == "/api/tags" { + _, _ = io.WriteString(w, `{"models":[{"name":"gpt-oss:20b","digest":"`+testArtifactB+`"}]}`) + return + } + inference.Add(1) + w.WriteHeader(http.StatusOK) + })) + defer engine.Close() + discovery := NewDiscovery() + discovery.AddManual(nodeForModel(t, "engine", engine.URL, "gpt-oss:20b")) + recorder := httptest.NewRecorder() + + testProxy(discovery, 11434).handleHTTP(recorder, attestedRequest(`{"model":"gpt-oss:20b"}`, testArtifactA)) + + if recorder.Code != http.StatusPreconditionFailed { + t.Fatalf("status = %d, want 412; body=%s", recorder.Code, recorder.Body.String()) + } + if got := inference.Load(); got != 0 { + t.Fatalf("mismatched engine received %d inference requests, want 0", got) + } + if got := headerValues(recorder.Header(), servedArtifactSHA256Header); len(got) != 0 { + t.Fatalf("rejection carried served-artifact evidence: %v", got) + } +} + +func TestHandleHTTP_ExactArtifactInventoryUnavailableFailsClosed(t *testing.T) { + engine := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.WriteHeader(http.StatusServiceUnavailable) + })) + defer engine.Close() + discovery := NewDiscovery() + discovery.AddManual(nodeForModel(t, "engine", engine.URL, "gpt-oss:20b")) + recorder := httptest.NewRecorder() + + testProxy(discovery, 11434).handleHTTP(recorder, attestedRequest(`{"model":"gpt-oss:20b"}`, testArtifactA)) + + if recorder.Code != http.StatusServiceUnavailable { + t.Fatalf("status = %d, want 503; body=%s", recorder.Code, recorder.Body.String()) + } +} + +func TestHandleHTTP_InvalidArtifactBindingRejectsLocally(t *testing.T) { + tests := []struct { + name string + path string + body string + values []string + }{ + {name: "malformed", path: "/api/chat", body: `{"model":"gpt-oss:20b"}`, values: []string{"not-a-digest"}}, + {name: "uppercase", path: "/api/chat", body: `{"model":"gpt-oss:20b"}`, values: []string{strings.ToUpper(testArtifactA)}}, + {name: "duplicate", path: "/api/chat", body: `{"model":"gpt-oss:20b"}`, values: []string{testArtifactA, testArtifactA}}, + {name: "non-inference", path: "/api/version", body: `{}`, values: []string{testArtifactA}}, + {name: "missing-model", path: "/api/chat", body: `{}`, values: []string{testArtifactA}}, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + request := httptest.NewRequest(http.MethodPost, test.path, strings.NewReader(test.body)) + request.Header[expectedArtifactSHA256Header] = test.values + recorder := httptest.NewRecorder() + testProxy(NewDiscovery(), 11434).handleHTTP(recorder, request) + if recorder.Code != http.StatusBadRequest { + t.Fatalf("status = %d, want 400; body=%s", recorder.Code, recorder.Body.String()) + } + }) + } +} + +func TestHandleHTTP_ExpectedArtifactTrailerIsRejected(t *testing.T) { + request := httptest.NewRequest(http.MethodPost, "/api/chat", strings.NewReader(`{"model":"gpt-oss:20b"}`)) + request.Trailer = http.Header{expectedArtifactSHA256Header: []string{testArtifactA}} + recorder := httptest.NewRecorder() + + testProxy(NewDiscovery(), 11434).handleHTTP(recorder, request) + + if recorder.Code != http.StatusBadRequest { + t.Fatalf("status = %d, want 400; body=%s", recorder.Code, recorder.Body.String()) + } +} + +func TestHandleHTTP_InventoryRedirectIsNotFollowed(t *testing.T) { + var redirected atomic.Int32 + target := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + redirected.Add(1) + writeArtifactInventory(w, testArtifactA) + })) + defer target.Close() + + var inference atomic.Int32 + engine := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path == "/api/tags" { + http.Redirect(w, r, target.URL+"/api/tags", http.StatusTemporaryRedirect) + return + } + inference.Add(1) + w.WriteHeader(http.StatusOK) + })) + defer engine.Close() + discovery := NewDiscovery() + discovery.AddManual(nodeForModel(t, "engine", engine.URL, "gpt-oss:20b")) + recorder := httptest.NewRecorder() + + testProxy(discovery, 11434).handleHTTP(recorder, attestedRequest(`{"model":"gpt-oss:20b"}`, testArtifactA)) + + if recorder.Code != http.StatusServiceUnavailable { + t.Fatalf("status = %d, want 503; body=%s", recorder.Code, recorder.Body.String()) + } + if got := redirected.Load(); got != 0 { + t.Fatalf("redirect target received %d requests, want 0", got) + } + if got := inference.Load(); got != 0 { + t.Fatalf("redirecting engine received %d inference requests, want 0", got) + } +} + +func TestHandleHTTP_LoadedDigestMismatchFailsOverBeforeCommit(t *testing.T) { + var staleInference atomic.Int32 + stale := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case "/api/tags": + writeArtifactInventory(w, testArtifactA) + case "/api/ps": + writeArtifactInventory(w, testArtifactB) + default: + staleInference.Add(1) + _, _ = io.WriteString(w, `{"engine":"stale"}`) + } + })) + defer stale.Close() + + var currentInference atomic.Int32 + current := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path == "/api/tags" || r.URL.Path == "/api/ps" { + writeArtifactInventory(w, testArtifactA) + return + } + currentInference.Add(1) + _, _ = io.WriteString(w, `{"engine":"current"}`) + })) + defer current.Close() + + discovery := NewDiscovery() + discovery.AddManual(nodeForModel(t, "stale", stale.URL, "gpt-oss:20b")) + discovery.AddManual(nodeForModel(t, "current", current.URL, "gpt-oss:20b")) + proxy := testProxy(discovery, 11434) + proxy.SetSelected("stale") + recorder := httptest.NewRecorder() + + proxy.handleHTTP(recorder, attestedRequest(`{"model":"gpt-oss:20b"}`, testArtifactA)) + + if recorder.Code != http.StatusOK || !strings.Contains(recorder.Body.String(), `"current"`) { + t.Fatalf("final response = %d %s, want current candidate", recorder.Code, recorder.Body.String()) + } + if staleInference.Load() != 1 || currentInference.Load() != 1 { + t.Fatalf("inference counts stale=%d current=%d, want 1 each", staleInference.Load(), currentInference.Load()) + } + if got := recorder.Header().Get(servedArtifactSHA256Header); got != testArtifactA { + t.Fatalf("served artifact = %q, want %s", got, testArtifactA) + } +} + +func TestHandleHTTP_ArtifactInventoryConcurrencyIsBounded(t *testing.T) { + var active atomic.Int32 + var maximum atomic.Int32 + engine := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path == "/api/tags" { + current := active.Add(1) + defer active.Add(-1) + for { + observed := maximum.Load() + if current <= observed || maximum.CompareAndSwap(observed, current) { + break + } + } + time.Sleep(25 * time.Millisecond) + writeArtifactInventory(w, testArtifactA) + return + } + if r.URL.Path == "/api/ps" { + writeArtifactInventory(w, testArtifactA) + return + } + _, _ = io.WriteString(w, `{"ok":true}`) + })) + defer engine.Close() + + discovery := NewDiscovery() + for i := 0; i < maxArtifactInventoryConcurrency+5; i++ { + discovery.AddManual(nodeForModel(t, "engine-"+string(rune('a'+i)), engine.URL, "gpt-oss:20b")) + } + recorder := httptest.NewRecorder() + testProxy(discovery, 11434).handleHTTP(recorder, attestedRequest(`{"model":"gpt-oss:20b"}`, testArtifactA)) + + if recorder.Code != http.StatusOK { + t.Fatalf("status = %d, want 200; body=%s", recorder.Code, recorder.Body.String()) + } + if got := maximum.Load(); got > maxArtifactInventoryConcurrency { + t.Fatalf("maximum concurrent inventory reads = %d, want <= %d", got, maxArtifactInventoryConcurrency) + } +} + +func TestHandleHTTP_UnboundRequestRetainsRoutingButStripsReservedEvidence(t *testing.T) { + engine := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.Header().Add(servedArtifactSHA256Header, testArtifactB) + w.WriteHeader(http.StatusOK) + _, _ = io.WriteString(w, `{"done":true}`) + })) + defer engine.Close() + discovery := NewDiscovery() + discovery.AddManual(nodeForModel(t, "engine", engine.URL, "gpt-oss:20b")) + recorder := httptest.NewRecorder() + + testProxy(discovery, 11434).handleHTTP( + recorder, + httptest.NewRequest(http.MethodPost, "/api/chat", strings.NewReader(`{"model":"gpt-oss:20b"}`)), + ) + + if recorder.Code != http.StatusOK { + t.Fatalf("status = %d, want 200", recorder.Code) + } + if got := headerValues(recorder.Header(), servedArtifactSHA256Header); len(got) != 0 { + t.Fatalf("unbound response retained reserved engine evidence: %v", got) + } +} + +func TestHandleHTTP_StripsReservedResponseTrailer(t *testing.T) { + engine := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.Header().Add("Trailer", servedArtifactSHA256Header+", X-Engine-Trace") + w.WriteHeader(http.StatusOK) + _, _ = io.WriteString(w, `{"done":true}`) + w.Header().Set(servedArtifactSHA256Header, testArtifactB) + w.Header().Set("X-Engine-Trace", "preserved") + })) + defer engine.Close() + discovery := NewDiscovery() + discovery.AddManual(nodeForModel(t, "engine", engine.URL, "gpt-oss:20b")) + recorder := httptest.NewRecorder() + + testProxy(discovery, 11434).handleHTTP( + recorder, + httptest.NewRequest(http.MethodPost, "/api/chat", strings.NewReader(`{"model":"gpt-oss:20b"}`)), + ) + + response := recorder.Result() + defer response.Body.Close() + _, _ = io.ReadAll(response.Body) + if got := headerValues(response.Header, servedArtifactSHA256Header); len(got) != 0 { + t.Fatalf("response header retained reserved evidence: %v", got) + } + if got := headerValues(response.Trailer, servedArtifactSHA256Header); len(got) != 0 { + t.Fatalf("response trailer retained reserved evidence: %v", got) + } + if got := response.Trailer.Get("X-Engine-Trace"); got != "preserved" { + t.Fatalf("unrelated trailer = %q, want preserved", got) + } +} diff --git a/services/ollama-proxy/proxy.go b/services/ollama-proxy/proxy.go index ad4ac7a2..2b150c51 100644 --- a/services/ollama-proxy/proxy.go +++ b/services/ollama-proxy/proxy.go @@ -382,6 +382,12 @@ type Proxy struct { plainTransport *http.Transport peerTransports map[string]*http.Transport + // artifactInventorySlots bounds all exact-artifact control-plane reads + // across concurrent client requests. A fixed worker pool also bounds each + // request independently; this process-wide gate prevents a request burst + // from multiplying the number of simultaneous /api/tags and /api/ps calls. + artifactInventorySlots chan struct{} + // nextRequestID is a monotonic counter for tagging RequestStarted / // RequestEvent pairs. Atomic add returns the new value, so request // IDs start at 1 and never collide within a single proxy lifetime. @@ -403,12 +409,13 @@ type Proxy struct { func NewProxy(codec *Codec, discovery *Discovery, port int) *Proxy { return &Proxy{ - codec: codec, - discovery: discovery, - port: port, - targets: reach.NewChooser(), - runID: newRunID(), - activity: nodeactivity.NewReporter(activityReportInterval), + codec: codec, + discovery: discovery, + port: port, + targets: reach.NewChooser(), + runID: newRunID(), + activity: nodeactivity.NewReporter(activityReportInterval), + artifactInventorySlots: make(chan struct{}, maxArtifactInventoryConcurrency), } } @@ -533,9 +540,13 @@ const ( proxyIdleConnTimeout = 90 * time.Second // Inbound http.Server limits — keep IdleTimeout aligned with client // IdleConnTimeout so idle keep-alives are reaped on both sides. - proxyReadHeaderTimeout = 10 * time.Second - proxyServerIdleTimeout = 90 * time.Second - maxModelListBytes = 16 << 20 + proxyReadHeaderTimeout = 10 * time.Second + proxyServerIdleTimeout = 90 * time.Second + maxModelListBytes = 16 << 20 + maxArtifactInventoryConcurrency = 8 + + expectedArtifactSHA256Header = "X-MrNiceAI-Expected-Artifact-Sha256" + servedArtifactSHA256Header = "X-MrNiceAI-Served-Artifact-Sha256" ) // idleClientWriteTimeout bounds how long a single write of streamed response @@ -851,9 +862,10 @@ func (p *Proxy) emitWorkload(method string, w Workload) { // pinned to that peer's exact server cert. Empty peerUUID means a plain-HTTP // dial — the local backend (self) or an explicit manual node. type candidate struct { - id string - url *url.URL - peerUUID string + id string + url *url.URL + peerUUID string + artifactDigest string } // candidateTransport returns the reverse-proxy / model-list transport for a @@ -952,6 +964,15 @@ type retrySignal struct{} func (retrySignal) Error() string { return "ollama-proxy: retry next candidate" } +// artifactVerificationSignal aborts an upstream response whose live loaded- +// model inventory cannot prove the requested digest. The failover loop may try +// the next prequalified candidate, but the rejected response is never exposed. +type artifactVerificationSignal struct{} + +func (artifactVerificationSignal) Error() string { + return "ollama-proxy: loaded artifact verification failed" +} + type modelListItem struct { key string digest string @@ -973,6 +994,285 @@ func ollamaModelKey(model string) string { return model } +func headerValues(h http.Header, name string) []string { + var values []string + for key, current := range h { + if strings.EqualFold(key, name) { + values = append(values, current...) + } + } + return values +} + +func removeHeader(h http.Header, name string) { + for key := range h { + if strings.EqualFold(key, name) { + delete(h, key) + } + } +} + +func headerPresent(h http.Header, name string) bool { + for key := range h { + if strings.EqualFold(key, name) { + return true + } + } + return false +} + +// removeTrailerDeclaration removes a reserved field from a comma-separated +// Trailer declaration without disturbing unrelated trailer names. +func removeTrailerDeclaration(h http.Header, name string) { + for key, values := range h { + if !strings.EqualFold(key, "Trailer") { + continue + } + kept := make([]string, 0, len(values)) + for _, value := range values { + var names []string + for _, current := range strings.Split(value, ",") { + current = strings.TrimSpace(current) + if current != "" && !strings.EqualFold(current, name) { + names = append(names, current) + } + } + if len(names) > 0 { + kept = append(kept, strings.Join(names, ", ")) + } + } + if len(kept) == 0 { + delete(h, key) + } else { + h[key] = kept + } + } +} + +func removeReservedArtifactFields(h http.Header) { + removeHeader(h, expectedArtifactSHA256Header) + removeHeader(h, servedArtifactSHA256Header) + removeTrailerDeclaration(h, expectedArtifactSHA256Header) + removeTrailerDeclaration(h, servedArtifactSHA256Header) +} + +type artifactTrailerSanitizingBody struct { + io.ReadCloser + trailer http.Header +} + +func (b *artifactTrailerSanitizingBody) sanitize() { + removeReservedArtifactFields(b.trailer) +} + +func (b *artifactTrailerSanitizingBody) Read(p []byte) (int, error) { + n, err := b.ReadCloser.Read(p) + if err == io.EOF { + // net/http populates response trailer values only when the body reaches + // EOF. ReverseProxy copies trailers after that read, so sanitize here as + // well as in ModifyResponse. + b.sanitize() + } + return n, err +} + +func (b *artifactTrailerSanitizingBody) Close() error { + b.sanitize() + return b.ReadCloser.Close() +} + +func validArtifactSHA256(value string) bool { + if len(value) != 64 { + return false + } + for _, char := range value { + if (char < '0' || char > '9') && (char < 'a' || char > 'f') { + return false + } + } + return true +} + +func expectedArtifactSHA256(header, trailer http.Header) (string, bool, error) { + values := headerValues(header, expectedArtifactSHA256Header) + // The trust contract requires an ordinary request header. A trailer arrives + // only after the body is consumed and must never become a second input path. + if headerPresent(trailer, expectedArtifactSHA256Header) { + return "", true, fmt.Errorf("expected artifact digest is not accepted as a trailer") + } + if len(values) == 0 { + return "", false, nil + } + if len(values) != 1 || !validArtifactSHA256(values[0]) { + return "", true, fmt.Errorf("expected artifact digest must be exactly one lowercase SHA-256 value") + } + return values[0], true, nil +} + +type candidateArtifactResult struct { + digest string + found bool + err error +} + +func (p *Proxy) acquireArtifactInventorySlot(ctx context.Context) error { + select { + case p.artifactInventorySlots <- struct{}{}: + return nil + case <-ctx.Done(): + return ctx.Err() + } +} + +func (p *Proxy) releaseArtifactInventorySlot() { + <-p.artifactInventorySlots +} + +func (p *Proxy) candidateInventoryDigest( + ctx context.Context, + cand candidate, + model string, + path string, +) (string, bool, error) { + if err := p.acquireArtifactInventorySlot(ctx); err != nil { + return "", false, err + } + defer p.releaseArtifactInventorySlot() + + target := *cand.url + target.Path = path + target.RawPath = "" + target.RawQuery = "" + request, err := http.NewRequestWithContext(ctx, http.MethodGet, target.String(), nil) + if err != nil { + return "", false, err + } + request.Header.Set("Accept", "application/json") + client := &http.Client{ + Timeout: modelListClient.Timeout, + Transport: p.candidateTransport(cand), + CheckRedirect: func(_ *http.Request, _ []*http.Request) error { + return http.ErrUseLastResponse + }, + } + response, err := client.Do(request) + if err != nil { + p.targets.Forget(cand.id) + return "", false, err + } + defer response.Body.Close() + if response.StatusCode != http.StatusOK { + return "", false, fmt.Errorf("inventory %s returned %s", path, response.Status) + } + body, err := io.ReadAll(io.LimitReader(response.Body, maxModelListBytes+1)) + if err != nil { + return "", false, err + } + if len(body) > maxModelListBytes { + return "", false, fmt.Errorf("inventory %s exceeds %d bytes", path, maxModelListBytes) + } + var envelope struct { + Models *[]json.RawMessage `json:"models"` + } + if err := json.Unmarshal(body, &envelope); err != nil { + return "", false, err + } + if envelope.Models == nil { + return "", false, fmt.Errorf("inventory %s has no model array", path) + } + requested := ollamaModelKey(model) + var digest string + matches := 0 + for _, raw := range *envelope.Models { + var identity struct { + Name string `json:"name"` + Model string `json:"model"` + Digest string `json:"digest"` + } + if err := json.Unmarshal(raw, &identity); err != nil { + return "", false, fmt.Errorf("invalid model record from %s: %w", path, err) + } + if identity.Name != "" && identity.Model != "" && ollamaModelKey(identity.Name) != ollamaModelKey(identity.Model) { + return "", false, fmt.Errorf("model record from %s has conflicting identities", path) + } + identityKey := identity.Name + if identityKey == "" { + identityKey = identity.Model + } + if ollamaModelKey(identityKey) != requested { + continue + } + matches++ + digest = identity.Digest + } + if matches == 0 { + return "", false, nil + } + if matches != 1 || !validArtifactSHA256(digest) { + return "", false, fmt.Errorf("requested model record from %s is ambiguous or has an invalid digest", path) + } + return digest, true, nil +} + +// candidateModelDigest reads the candidate's native Ollama inventory through +// the same transport used for inference. The returned digest is bound only to +// one unambiguous record for the requested model in this request-local read. +func (p *Proxy) candidateModelDigest(ctx context.Context, cand candidate, model string) (string, bool, error) { + return p.candidateInventoryDigest(ctx, cand, model, "/api/tags") +} + +// candidateLoadedModelDigest reads Ollama's scheduler-owned live process +// inventory. Ollama derives this digest from the exact Model object held by the +// loaded runner, so it closes the gap between an installed tag and the model +// that was actually selected for the response now waiting to commit. +func (p *Proxy) candidateLoadedModelDigest(ctx context.Context, cand candidate, model string) (string, bool, error) { + return p.candidateInventoryDigest(ctx, cand, model, "/api/ps") +} + +// filterCandidatesByArtifact intersects the already model-eligible, ordered +// candidate list with a bounded request-local digest inventory. The order is +// retained so manual selection, scheduler priority, and failover semantics do +// not change. Each retained candidate carries the immutable digest value read +// for this request, which is later stamped only if that candidate commits. +func (p *Proxy) filterCandidatesByArtifact(ctx context.Context, candidates []candidate, model, expected string) ([]candidate, bool) { + results := make([]candidateArtifactResult, len(candidates)) + workerCount := min(len(candidates), maxArtifactInventoryConcurrency) + jobs := make(chan int, len(candidates)) + for i := range candidates { + jobs <- i + } + close(jobs) + var wg sync.WaitGroup + for range workerCount { + wg.Add(1) + go func() { + defer wg.Done() + for i := range jobs { + results[i].digest, results[i].found, results[i].err = p.candidateModelDigest(ctx, candidates[i], model) + } + }() + } + wg.Wait() + + available := false + matched := make([]candidate, 0, len(candidates)) + for i, result := range results { + if result.err != nil { + slog.Debug("artifact inventory candidate unavailable", + "node_id", candidates[i].id, "target", candidates[i].url.Host, "err", result.err) + continue + } + available = true + if !result.found || result.digest != expected { + continue + } + cand := candidates[i] + cand.artifactDigest = result.digest + matched = append(matched, cand) + } + return matched, available +} + // serveModelList queries every Ollama candidate concurrently and returns the // native /api/tags or OpenAI /v1/models envelope with duplicate model records // removed. Results are merged in candidate order, not completion order, so @@ -1142,11 +1442,50 @@ func (p *Proxy) handleHTTP(w http.ResponseWriter, r *http.Request) { // routing behavior even when their JSON happens to contain a model field. bodyBytes, model := bufferBodyAndModel(r) isInf := isInferenceRequest(r.Method, r.URL.Path) + expectedArtifact, artifactBound, artifactErr := expectedArtifactSHA256(r.Header, r.Trailer) + removeReservedArtifactFields(r.Header) + removeReservedArtifactFields(r.Trailer) + if artifactBound && (artifactErr != nil || !isInf || model == "") { + cors.Apply(w.Header()) + w.Header().Set("Content-Type", "application/json") + w.Header().Set("X-Content-Type-Options", "nosniff") + w.WriteHeader(http.StatusBadRequest) + _, _ = io.WriteString(w, `{"error":"invalid artifact binding"}`) + p.codec.Notify("proxy/request", RequestEvent{ + ID: reqID, Method: r.Method, Path: r.URL.Path, Status: http.StatusBadRequest, + Duration: time.Since(start).Milliseconds(), Error: "invalid artifact binding", + }) + return + } routingModel := "" if isInf { routingModel = model } candidates := p.resolveCandidates(routingModel) + if artifactBound && len(candidates) > 0 { + var inventoryAvailable bool + candidates, inventoryAvailable = p.filterCandidatesByArtifact(r.Context(), candidates, model, expectedArtifact) + if len(candidates) == 0 { + status := http.StatusPreconditionFailed + errorText := "no candidate matches requested model artifact" + body := `{"error":"requested model artifact is unavailable"}` + if !inventoryAvailable { + status = http.StatusServiceUnavailable + errorText = "artifact inventory unavailable" + body = `{"error":"model artifact inventory is unavailable"}` + } + cors.Apply(w.Header()) + w.Header().Set("Content-Type", "application/json") + w.Header().Set("X-Content-Type-Options", "nosniff") + w.WriteHeader(status) + _, _ = io.WriteString(w, body) + p.codec.Notify("proxy/request", RequestEvent{ + ID: reqID, Method: r.Method, Path: r.URL.Path, Status: status, + Duration: time.Since(start).Milliseconds(), Error: errorText, + }) + return + } + } if isInf && model != "" { candidates = p.reserveCandidate(candidates) } @@ -1322,6 +1661,8 @@ func (p *Proxy) handleHTTP(w http.ResponseWriter, r *http.Request) { r.Body = io.NopCloser(bytes.NewReader(bodyBytes)) } retry := false + artifactFailureStatus := 0 + artifactFailureError := "" sc := &statusCapture{ResponseWriter: w, status: http.StatusOK, idle: idleClientWriteTimeout} proxy := &httputil.ReverseProxy{ @@ -1329,6 +1670,8 @@ func (p *Proxy) handleHTTP(w http.ResponseWriter, r *http.Request) { req.URL.Scheme = cand.url.Scheme req.URL.Host = cand.url.Host req.Host = cand.url.Host + removeReservedArtifactFields(req.Header) + removeReservedArtifactFields(req.Trailer) }, // A remote cluster peer is dialed over mTLS (per-peer pinned config); // self/manual candidates use the plain transport. See candidateTransport. @@ -1337,12 +1680,40 @@ func (p *Proxy) handleHTTP(w http.ResponseWriter, r *http.Request) { // have arrived but before the body streams. That's both the retry // decision point and, on commit, the time-to-first-byte boundary. ModifyResponse: func(resp *http.Response) error { + // This is a reserved proxy trust signal. Never pass through a value + // supplied by an engine or another intermediary. + removeReservedArtifactFields(resp.Header) + removeReservedArtifactFields(resp.Trailer) + resp.Body = &artifactTrailerSanitizingBody{ReadCloser: resp.Body, trailer: resp.Trailer} if !last && shouldRetry(resp.StatusCode) { // Abort before streaming: ReverseProxy closes resp.Body and // calls ErrorHandler with our sentinel, then we try next. retry = true return retrySignal{} } + if artifactBound { + loadedDigest, found, err := p.candidateLoadedModelDigest(r.Context(), cand, model) + switch { + case err != nil: + artifactFailureStatus = http.StatusServiceUnavailable + artifactFailureError = "loaded model artifact inventory is unavailable" + slog.Debug("loaded artifact verification unavailable", + "node_id", cand.id, "target", cand.url.Host, "err", err) + case !found: + artifactFailureStatus = http.StatusPreconditionFailed + artifactFailureError = "requested model artifact is not loaded" + case loadedDigest != cand.artifactDigest: + artifactFailureStatus = http.StatusPreconditionFailed + artifactFailureError = "loaded model artifact does not match the requested digest" + } + if artifactFailureStatus != 0 { + if !last { + retry = true + } + return artifactVerificationSignal{} + } + resp.Header.Set(servedArtifactSHA256Header, cand.artifactDigest) + } // Prefer an engine-declared preflight policy so an exact origin plus // Allow-Credentials can pass a credentialed browser fetch. Engines // that publish no policy retain the proxy's permissive 204 fallback. @@ -1404,6 +1775,24 @@ func (p *Proxy) handleHTTP(w http.ResponseWriter, r *http.Request) { if _, ok := err.(retrySignal); ok { return // retryable status — the loop advances to the next candidate } + if _, ok := err.(artifactVerificationSignal); ok { + if !last { + return // uncommitted artifact response — try the next candidate + } + servedNodeID = cand.id + servedTarget = cand.url.Host + proxyErr = artifactFailureError + body, marshalErr := json.Marshal(map[string]string{"error": artifactFailureError}) + if marshalErr != nil { + body = []byte(`{"error":"model artifact verification failed"}`) + } + cors.Apply(ew.Header()) + ew.Header().Set("Content-Type", "application/json") + ew.Header().Set("X-Content-Type-Options", "nosniff") + ew.WriteHeader(artifactFailureStatus) + _, _ = ew.Write(body) + return + } // Transport/dial error (not a status-based retry): forget this // node's confirmed address so the next request re-confirms and // can fail over to another of its published addresses diff --git a/services/versions.json b/services/versions.json index 29d8c230..b4584183 100644 --- a/services/versions.json +++ b/services/versions.json @@ -3,7 +3,7 @@ "product": "0.91.7", "installer": "0.91.7", "components": { - "ollama-proxy": "0.26.2", + "ollama-proxy": "0.27.0", "lmstudio-proxy": "0.16.2", "nvpair-node-info": "0.13.3", "nvpair-node-scanner": "0.20.3",