diff --git a/CHANGELOG.md b/CHANGELOG.md index 5ce1511..d1f6ffb 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -15,9 +15,8 @@ has a matching GitHub release whose body comes from its section below. is set. A non-streaming answer carries `X-Berth-Fallback: `; a streaming one, whose SSE preamble was already sent, gets the fallback's stream after one `: berth: fallback ` comment. The gateway never starts, leases or drains the fallback, and the primary's start is never - cancelled, so later requests land on it. The request reaches the fallback without its - `Authorization` header, so the node's API key never leaves the gateway. One hop: `berth config check` - refuses chains. + cancelled, so later requests land on it. The fallback never receives the node's API key, only its + own key from its entry when it has one. One hop: `berth config check` refuses chains. - Cgroup attribution closes the one known stowaway window (chaos scenario 3: a child whose parent exits within one 2 s tick). On Linux with a systemd user manager, `berth run` execs the engine in its own transient cgroup scope (`systemd-run --user @@ -52,6 +51,13 @@ has a matching GitHub release whose body comes from its section below. `docs/`: install, engines, leasing, models, gateway, the harbor, interfaces, benchmarks, roadmap and FAQ. The README's old Tested hardware, Swap numbers and Roadmap headings link to their new pages. +### Security + +- The gateway no longer forwards the node's API key to engines. A proxied request, fallback and + ComfyUI routes included, carries the engine's own key when its entry has one (`api_key_env` or + `api_key_file`) and no `Authorization` when it has none. If an engine of yours checked the node's + key, give it its own key in its `berth.toml` entry. + ### Fixed - The wake page and start errors show an engine's last log line without terminal color codes diff --git a/docs/gateway.md b/docs/gateway.md index e14b505..db052b8 100644 --- a/docs/gateway.md +++ b/docs/gateway.md @@ -21,6 +21,11 @@ - **API key.** Bearer token on every route; `berth key` prints it (created on first use, mode 600). Loopback only; browser origins and non-loopback Host names are refused. +The node key authenticates the client to the gateway and stays there. A proxied request carries the +engine's own key when its entry has one (`api_key_env` or `api_key_file`), and no `Authorization` +when it has none. A credential that will not resolve (an unset variable, an unreadable file) sends +no `Authorization` and logs one WARN line that names the engine and the error, never the key. + ## Switch engines by naming a model Two engines that each want the whole card can't run together, and most harnesses can't start an engine @@ -75,8 +80,9 @@ request keeps waiting for its own engine. The primary's start is never cancelled requests that come later land on it. `fallback_model` names the model the fallback lists the request under when the two disagree, a non-streaming answer carries `X-Berth-Fallback: `, and a streaming answer, whose headers went out with the keep-alive preamble, carries one -`: berth: fallback ` comment before its data. The fallback gets the request without its -`Authorization` header: the node's API key stays with the gateway. One hop, per engine, opt-in: `berth config check` +`: berth: fallback ` comment before its data. The fallback gets the request without the +node's `Authorization` header; when the fallback's entry has its own key (`api_key_env` or `api_key_file`), +that is the one the request carries. One hop, per engine, opt-in: `berth config check` refuses chains and a fallback that is not an engine in the file. ```toml diff --git a/internal/gateway/enginekey_test.go b/internal/gateway/enginekey_test.go new file mode 100644 index 0000000..2fa6717 --- /dev/null +++ b/internal/gateway/enginekey_test.go @@ -0,0 +1,256 @@ +// Copyright 2026 Trevin Tindall (T-Crypt) +// SPDX-License-Identifier: Apache-2.0 + +//go:build unix + +package gateway + +import ( + "fmt" + "io" + "log/slog" + "net/http" + "net/http/httptest" + "os" + "strings" + "sync" + "testing" + "time" + + "github.com/T-Crypt/berth/internal/broker" + "github.com/T-Crypt/berth/internal/config" +) + +// keyEngine is an OpenAI-style engine that is up at once and records the +// Authorization of the chat request it receives, so a test can say what the +// gateway sent it. +func keyEngine(t *testing.T, model string) (*httptest.Server, func() string) { + t.Helper() + var mu sync.Mutex + auth := "" + mux := http.NewServeMux() + mux.HandleFunc("/v1/models", func(w http.ResponseWriter, _ *http.Request) { + fmt.Fprintf(w, `{"object":"list","data":[{"id":%q}]}`, model) + }) + mux.HandleFunc("/v1/chat/completions", func(w http.ResponseWriter, r *http.Request) { + mu.Lock() + auth = r.Header.Get("Authorization") + mu.Unlock() + w.Header().Set("Content-Type", "text/event-stream") + fmt.Fprint(w, "data: {\"choices\":[{\"delta\":{\"content\":\"key-marker\"}}]}\n\n") + fmt.Fprint(w, "data: [DONE]\n\n") + }) + srv := httptest.NewServer(mux) + t.Cleanup(srv.Close) + return srv, func() string { mu.Lock(); defer mu.Unlock(); return auth } +} + +// keyGateway pairs a configured engine at url with the test daemon and the +// gateway in front of it. +func keyGateway(t *testing.T, e config.Engine) (*httptest.Server, *broker.Broker) { + t.Helper() + sock, b := daemon(t) + g := &Gateway{Socket: sock, Ready: 5 * time.Second, Key: testKey, Engines: []config.Engine{e}} + gw := httptest.NewServer(g.Handler()) + t.Cleanup(gw.Close) + return gw, b +} + +// A client's node key authenticates it to the gateway and goes no further: +// an engine without its own credential receives no Authorization at all. +func TestProxySendsNoKeyToACredentiallessEngine(t *testing.T) { + srv, saw := keyEngine(t, "llama3") + gw, b := keyGateway(t, config.Engine{Name: "ollama", Type: "ollama", URL: srv.URL, Serves: []string{"llama3"}}) + resp := call(t, "POST", gw.URL+"/v1/chat/completions", `{"model":"llama3"}`) + body, _ := io.ReadAll(resp.Body) + resp.Body.Close() + if resp.StatusCode != http.StatusOK || !strings.Contains(string(body), "key-marker") { + t.Fatalf("%d %s", resp.StatusCode, body) + } + if got := saw(); got != "" { + t.Fatalf("the engine received Authorization %q, want none", got) + } + if n := len(b.Status().Leases); n != 0 { + t.Fatalf("%d leases for an engine that was already up", n) + } +} + +// An engine with its own credential receives that key on every proxied +// request, not the node's. +func TestProxySendsTheEngineOwnKey(t *testing.T) { + const env = "GW_TEST_ENGINE_KEY" + t.Setenv(env, "engine-key-9876543210") + srv, saw := keyEngine(t, "llama3") + gw, _ := keyGateway(t, config.Engine{Name: "ollama", Type: "ollama", URL: srv.URL, + Serves: []string{"llama3"}, APIKeyEnv: env}) + resp := call(t, "POST", gw.URL+"/v1/chat/completions", `{"model":"llama3"}`) + body, _ := io.ReadAll(resp.Body) + resp.Body.Close() + if resp.StatusCode != http.StatusOK || !strings.Contains(string(body), "key-marker") { + t.Fatalf("%d %s", resp.StatusCode, body) + } + if got := saw(); got != "Bearer engine-key-9876543210" { + t.Fatalf("the engine received %q, want its own bearer", got) + } +} + +// A fallback that switched in for the slow primary carries the fallback's +// own key when it has one, instead of the plain drop the node key used to +// get. +func TestFallbackSendsItsOwnKey(t *testing.T) { + const env = "GW_TEST_FALLBACK_KEY" + t.Setenv(env, "fallback-key-123456") + var mu sync.Mutex + auth := "" + mux := http.NewServeMux() + mux.HandleFunc("/v1/models", func(w http.ResponseWriter, _ *http.Request) { + fmt.Fprint(w, `{"object":"list","data":[{"id":"cpu"}]}`) + }) + mux.HandleFunc("/v1/chat/completions", func(w http.ResponseWriter, r *http.Request) { + mu.Lock() + auth = r.Header.Get("Authorization") + mu.Unlock() + w.Header().Set("Content-Type", "text/event-stream") + fmt.Fprint(w, "data: {\"choices\":[{\"delta\":{\"content\":\"from the cpu fallback\"}}]}\n\n") + fmt.Fprint(w, "data: [DONE]\n\n") + }) + fb := httptest.NewServer(mux) + t.Cleanup(fb.Close) + + sock, b := daemon(t) + primary := freeAddr(t) + t.Setenv("BERTH_GW_ENGINE", primary+",big,3s") + g := &Gateway{Socket: sock, Ready: 10 * time.Second, Key: testKey, + Log: slog.New(slog.NewTextHandler(io.Discard, nil)), + Engines: []config.Engine{ + {Name: "strata", Type: "llamacpp", URL: "http://" + primary, Start: []string{os.Args[0]}, + Serves: []string{"big"}, Priority: "interactive", + Fallback: "strata-cpu", FallbackAfter: config.Duration{Duration: 300 * time.Millisecond}, FallbackModel: "cpu"}, + {Name: "strata-cpu", Type: "llamacpp", URL: fb.URL, Serves: []string{"cpu"}, APIKeyEnv: env}, + }} + gw := httptest.NewServer(g.Handler()) + t.Cleanup(gw.Close) + + resp := call(t, "POST", gw.URL+"/v1/chat/completions", `{"model":"big"}`) + body, _ := io.ReadAll(resp.Body) + resp.Body.Close() + if resp.StatusCode != http.StatusOK || resp.Header.Get("X-Berth-Fallback") != "strata-cpu" { + t.Fatalf("status %d %q: the request did not go to the fallback", resp.StatusCode, body) + } + mu.Lock() + got := auth + mu.Unlock() + if got != "Bearer fallback-key-123456" { + t.Fatalf("the fallback received %q, want its own bearer", got) + } + if n := len(b.Status().Leases); n == 0 { + t.Fatal("the primary's start never took its lease") + } +} + +// A ComfyUI route the gateway proxies (/prompt with the engine up, the +// side routes) reaches the engine with no Authorization: the node key is +// the client's credential for the gateway, and a ComfyUI with no key of +// its own wants none. +func TestComfyRoutesSendNoKey(t *testing.T) { + var mu sync.Mutex + auth := "" + see := func(w http.ResponseWriter, r *http.Request) { + mu.Lock() + if a := r.Header.Get("Authorization"); a != "" { + auth = a + } + mu.Unlock() + fmt.Fprint(w, "comfy-marker") + } + mux := http.NewServeMux() + mux.HandleFunc("/prompt", func(w http.ResponseWriter, r *http.Request) { + see(w, r) + fmt.Fprint(w, `{"prompt_id":"gw-test-1"}`) + }) + for _, p := range []string{"/queue", "/history/", "/view", "/upload/"} { + mux.HandleFunc(p, see) + } + srv := httptest.NewServer(mux) + t.Cleanup(srv.Close) + + sock, b := daemon(t) + g := &Gateway{Socket: sock, Ready: 5 * time.Second, Key: testKey, + Engines: []config.Engine{{Name: "comfyui", Type: "comfyui", URL: srv.URL}}} + gw := httptest.NewServer(g.Handler()) + t.Cleanup(gw.Close) + + for _, c := range []struct{ method, path, body, want string }{ + {"POST", "/prompt", `{"prompt":{"3":{"class_type":"KSampler"}}}`, `"prompt_id":"gw-test-1"`}, + {"GET", "/queue", "", "comfy-marker"}, + {"GET", "/history/gw-test-1", "", "comfy-marker"}, + {"GET", "/view?filename=out.png", "", "comfy-marker"}, + {"POST", "/upload/image", "bytes", "comfy-marker"}, + } { + resp := call(t, c.method, gw.URL+c.path, c.body) + body, _ := io.ReadAll(resp.Body) + resp.Body.Close() + if resp.StatusCode != http.StatusOK || !strings.Contains(string(body), c.want) { + t.Fatalf("%s %s: %d %s", c.method, c.path, resp.StatusCode, body) + } + } + mu.Lock() + got := auth + mu.Unlock() + if got != "" { + t.Fatalf("ComfyUI received Authorization %q, want none", got) + } + if n := len(b.Status().Leases); n != 0 { + t.Fatalf("%d leases for a ComfyUI that was already up", n) + } +} + +// A credential that will not resolve (an unset variable) sends no +// Authorization and logs one WARN line naming the engine and the error, +// never the key: the request goes through, and the engine refuses in its +// own words if it needs one. +func TestUnresolvableKeySendsNoneAndWarnsOnce(t *testing.T) { + const env = "GW_TEST_MISSING_KEY" + if os.Getenv(env) != "" { + t.Skipf("%s is set in the test environment", env) + } + srv, saw := keyEngine(t, "llama3") + sock, b := daemon(t) + log := &syncBuf{} + g := &Gateway{Socket: sock, Ready: 5 * time.Second, Key: testKey, + Log: slog.New(slog.NewTextHandler(log, nil)), + Engines: []config.Engine{{Name: "ollama", Type: "ollama", URL: srv.URL, + Serves: []string{"llama3"}, APIKeyEnv: env}}} + gw := httptest.NewServer(g.Handler()) + t.Cleanup(gw.Close) + + resp := call(t, "POST", gw.URL+"/v1/chat/completions", `{"model":"llama3"}`) + body, _ := io.ReadAll(resp.Body) + resp.Body.Close() + if resp.StatusCode != http.StatusOK || !strings.Contains(string(body), "key-marker") { + t.Fatalf("%d %s", resp.StatusCode, body) + } + if got := saw(); got != "" { + t.Fatalf("the engine received Authorization %q, want none", got) + } + lines := strings.Split(strings.TrimSpace(log.String()), "\n") + warns := 0 + for _, l := range lines { + if !strings.Contains(l, "engine key not resolved") { + continue + } + warns++ + if !strings.Contains(l, `engine=ollama`) || !strings.Contains(l, env) { + t.Fatalf("the WARN line does not name the engine and the cause: %s", l) + } + if strings.Contains(l, testKey) { + t.Fatalf("the WARN line carries a key: %s", l) + } + } + if warns != 1 { + t.Fatalf("%d WARN lines for one request, want one:\n%s", warns, log.String()) + } + if n := len(b.Status().Leases); n != 0 { + t.Fatalf("%d leases for an engine that was already up", n) + } +} diff --git a/internal/gateway/gateway.go b/internal/gateway/gateway.go index 3ba872d..bebbc2a 100644 --- a/internal/gateway/gateway.go +++ b/internal/gateway/gateway.go @@ -270,6 +270,44 @@ type target struct { live bool // an engine already answering lists the model } +// engineEntry is the configured entry behind a target: the one the +// gateway may start it when that is set, otherwise the entry whose url +// is the target's (an already-running configured engine). A discovered +// engine has no entry. +func (g *Gateway) engineEntry(t target) *config.Engine { + if t.engine != nil { + return t.engine + } + want := strings.TrimRight(t.url, "/") + for i := range g.Engines { + if strings.TrimRight(g.Engines[i].URL, "/") == want { + return &g.Engines[i] + } + } + return nil +} + +// engineAuth is the Authorization a proxied request to the engine +// carries: the engine's own key when it has one, none otherwise. The +// client's node key is a credential for this gateway and goes no +// further: the engine, its logs and any proxy in front of it would see +// it. A credential that will not resolve sends nothing and logs one +// WARN line naming the engine and the error, never the key. +func (g *Gateway) engineAuth(entry *config.Engine, r *http.Request) { + r.Header.Del("Authorization") + if entry == nil { + return + } + key, err := entry.APIKey() + if err != nil { + g.log().Warn("engine key not resolved", "engine", entry.Name, "err", err) + return + } + if key != "" { + r.Header.Set("Authorization", "Bearer "+key) + } +} + // live is every engine the node reads that is up, configured or discovered. func (g *Gateway) live() []proto.Engine { if st := g.nodeStatus(); st != nil { @@ -533,11 +571,15 @@ func (g *Gateway) forward(w http.ResponseWriter, r *http.Request) { } } base := t.url + // The request ends at the fallback when it switched to it, else the + // engine that serves the model. Either way it carries that engine's + // own key when it has one: the client's node key is a credential for + // this gateway and goes no further, a fallback included. if fb != nil { base, body = g.tookFallback(w, preamble, t.engine, fb, waited, body) - // The node key is the client's credential for this gateway; a - // fallback may run on another box and never gets it. - r.Header.Del("Authorization") + g.engineAuth(fb, r) + } else { + g.engineAuth(g.engineEntry(t), r) } u, err := url.Parse(base) if err != nil { @@ -680,7 +722,7 @@ func (g *Gateway) comfyPrompt(w http.ResponseWriter, r *http.Request) { return } } - g.comfyProxyTo(t.url, w, r) + g.comfyProxyTo(t, w, r) } // comfyProxy passes the rest of the ComfyUI surface (/ws, /history, /view, @@ -697,17 +739,18 @@ func (g *Gateway) comfyProxy(w http.ResponseWriter, r *http.Request) { apiError(w, http.StatusServiceUnavailable, "ComfyUI is not running; POST /prompt first to start it") return } - g.comfyProxyTo(t.url, w, r) + g.comfyProxyTo(t, w, r) } -func (g *Gateway) comfyProxyTo(base string, w http.ResponseWriter, r *http.Request) { - u, err := url.Parse(base) +func (g *Gateway) comfyProxyTo(t target, w http.ResponseWriter, r *http.Request) { + u, err := url.Parse(t.url) if err != nil { apiError(w, http.StatusInternalServerError, "engine url: "+err.Error()) return } p := httputil.NewSingleHostReverseProxy(u) p.FlushInterval = -1 // stream progress as ComfyUI sends it + g.engineAuth(g.engineEntry(t), r) r.Host = u.Host p.ServeHTTP(w, r) }