diff --git a/server/llm_key_endpoint_test.go b/server/llm_key_endpoint_test.go new file mode 100644 index 0000000..e4af8c4 --- /dev/null +++ b/server/llm_key_endpoint_test.go @@ -0,0 +1,127 @@ +package server + +import ( + "bytes" + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "sync/atomic" + "testing" + + "github.com/Autumn-27/artex/agent" + "github.com/Autumn-27/artex/db" +) + +func TestSameLLMEndpoint(t *testing.T) { + for _, tc := range []struct { + provA, baseA, provB, baseB string + want bool + }{ + {"anthropic", "https://gw.example", "anthropic", "https://gw.example/", true}, + {"anthropic", "https://gw.example/v1/messages", "anthropic", "https://gw.example", true}, + {"", "", "anthropic", "", true}, // empty provider defaults to anthropic + {"openai", " https://gw.example/v1/chat/completions ", "openai", "https://gw.example/v1", true}, + {"anthropic", "https://evil.example", "anthropic", "https://gw.example", false}, + {"anthropic", "", "anthropic", "https://gw.example", false}, // canonical host != saved host + {"openai", "", "anthropic", "", false}, // different canonical hosts + {"openai", "https://gw.example/v1", "openai-responses", "https://gw.example/v1", false}, + } { + if got := sameLLMEndpoint(tc.provA, tc.baseA, tc.provB, tc.baseB); got != tc.want { + t.Errorf("sameLLMEndpoint(%q,%q,%q,%q)=%v want %v", tc.provA, tc.baseA, tc.provB, tc.baseB, got, tc.want) + } + } +} + +// keyRecorder is a fake Anthropic endpoint that records the x-api-key it receives. +func keyRecorder(t *testing.T) (*httptest.Server, *atomic.Int32, *atomic.Value) { + var hits atomic.Int32 + var key atomic.Value + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + hits.Add(1) + key.Store(r.Header.Get("x-api-key")) + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"id":"m","type":"message","role":"assistant","model":"m","content":[{"type":"text","text":"OK"}],"stop_reason":"end_turn","usage":{"input_tokens":1,"output_tokens":1},"data":[]}`)) + })) + t.Cleanup(srv.Close) + return srv, &hits, &key +} + +func postJSON(t *testing.T, h http.HandlerFunc, path string, body map[string]any) map[string]any { + t.Helper() + b, _ := json.Marshal(body) + req := httptest.NewRequest(http.MethodPost, path, bytes.NewReader(b)) + req.Header.Set("Content-Type", "application/json") + rec := httptest.NewRecorder() + h(rec, req) + var resp map[string]any + _ = json.Unmarshal(rec.Body.Bytes(), &resp) + return resp +} + +// The global key must not follow a base_url typed into the test form. +func TestLLMTestGlobalKeyOnlyForGlobalEndpoint(t *testing.T) { + saved, savedHits, savedKey := keyRecorder(t) + other, otherHits, _ := keyRecorder(t) + s := &Server{} + s.llmCfg = agent.ConfigFrom("anthropic", "m", saved.URL, "global-secret", "") + + resp := postJSON(t, s.testLLM, "/api/llm/test", map[string]any{ + "provider": "anthropic", "model": "m", "base_url": other.URL, "api_key": "", "streaming": false, + }) + if otherHits.Load() != 0 || resp["ok"] != false { + t.Fatalf("global key went to a form-supplied host: hits=%d resp=%v", otherHits.Load(), resp) + } + + postJSON(t, s.testLLM, "/api/llm/test", map[string]any{ + "provider": "anthropic", "model": "m", "base_url": saved.URL + "/", "api_key": "", "streaming": false, + }) + if k, _ := savedKey.Load().(string); savedHits.Load() != 1 || k != "global-secret" { + t.Fatalf("global key not used for its own endpoint: hits=%d key=%q", savedHits.Load(), k) + } +} + +// A stored profile's key goes only to that profile's base URL, for both the +// connection test and the model listing. +func TestStoredProfileKeyOnlyForStoredEndpoint(t *testing.T) { + m, err := NewManager(t.TempDir(), "") + if err != nil { + t.Skipf("postgres unavailable (%v)", err) + } + t.Cleanup(func() { m.Close() }) + ctx, cancel := context.WithCancel(context.Background()) + cancel() + s := New(ctx, m, t.TempDir(), t.TempDir(), t.TempDir()) + s.cfgMu.Lock() + s.llmCfg = agent.Config{} // no global key to fall back on + s.cfgMu.Unlock() + + saved, savedHits, savedKey := keyRecorder(t) + other, otherHits, _ := keyRecorder(t) + pid, err := m.pg.SaveProfile(&db.LLMProfile{ + Name: "key-endpoint", Format: "anthropic", Model: "m", BaseURL: saved.URL, APIKey: "profile-secret", + }) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _, _ = m.pg.Exec("DELETE FROM llm_profiles WHERE id=$1", pid) }) + + for _, ep := range []struct { + path string + h http.HandlerFunc + }{{"/api/llm/test", s.testLLM}, {"/api/llm/models", s.pgListModels}} { + resp := postJSON(t, ep.h, ep.path, map[string]any{ + "provider": "anthropic", "model": "m", "base_url": other.URL, "api_key": "", "profile_id": pid, "streaming": false, + }) + if otherHits.Load() != 0 || resp["ok"] != false { + t.Fatalf("%s: stored key went to a form-supplied host: hits=%d resp=%v", ep.path, otherHits.Load(), resp) + } + before := savedHits.Load() + postJSON(t, ep.h, ep.path, map[string]any{ + "provider": "anthropic", "model": "m", "base_url": saved.URL, "api_key": "", "profile_id": pid, "streaming": false, + }) + if k, _ := savedKey.Load().(string); savedHits.Load() == before || k != "profile-secret" { + t.Fatalf("%s: stored key not used for its own endpoint: key=%q", ep.path, k) + } + } +} diff --git a/server/server.go b/server/server.go index 65fbc45..8a76124 100644 --- a/server/server.go +++ b/server/server.go @@ -1466,19 +1466,38 @@ func (s *Server) testLLM(w http.ResponseWriter, r *http.Request) { if cfg.NeedsAPIKey() { // API key resolution precedence: form input > the named profile's stored key > the global config's key. // A stored profile's key is never sent back to the browser, so the form is empty when testing a stored profile and the key comes from the DB. + // A saved key is only sent to the endpoint it was saved for: base_url and + // provider come from the form, so without this check the test button sends + // the saved key to any host typed into the form. + withheld := false if stored != nil && cfg.APIKey == "" { - // Trim like ConfigFrom does for form values: a key stored with - // surrounding whitespace runs fine in conversations (the runtime - // path re-trims) but would be sent raw here and fail auth. - cfg.APIKey = strings.TrimSpace(stored.APIKey) + if sameLLMEndpoint(req.Provider, req.BaseURL, stored.Format, stored.BaseURL) { + // Trim like ConfigFrom does for form values: a key stored with + // surrounding whitespace runs fine in conversations (the runtime + // path re-trims) but would be sent raw here and fail auth. + cfg.APIKey = strings.TrimSpace(stored.APIKey) + } else { + withheld = true + } } if cfg.APIKey == "" { s.cfgMu.Lock() - cfg.APIKey = s.llmCfg.APIKey + global := s.llmCfg s.cfgMu.Unlock() + if global.APIKey != "" { + if sameLLMEndpoint(req.Provider, req.BaseURL, global.Provider(), global.BaseURL) { + cfg.APIKey = global.APIKey + } else { + withheld = true + } + } } if cfg.APIKey == "" { - writeJSON(w, 200, map[string]any{"ok": false, "error": "no API key provided"}) + msg := "no API key provided" + if withheld { + msg = "no API key provided: the saved key is only sent to the endpoint it was saved for; enter the key to test a different base URL or provider" + } + writeJSON(w, 200, map[string]any{"ok": false, "error": msg}) return } } @@ -1496,6 +1515,16 @@ func (s *Server) testLLM(w http.ResponseWriter, r *http.Request) { }) } +// sameLLMEndpoint reports whether two provider/base URL pairs resolve to the same +// endpoint after ConfigFrom's normalization (trimmed, endpoint suffix stripped; +// an empty base URL means the provider's canonical host). A saved API key may only +// be sent when the request's endpoint matches the one it was saved for. +func sameLLMEndpoint(provA, baseA, provB, baseB string) bool { + a := agent.ConfigFrom(provA, "", baseA, "", "") + b := agent.ConfigFrom(provB, "", baseB, "", "") + return a.Format == b.Format && a.BaseURL == b.BaseURL +} + // truncateReply clips a connection-test reply for display. A model told to answer // "OK" can still ramble (or think out loud); the UI only needs enough to show it // really said something. Rune-based so multibyte text never splits mid-character. diff --git a/server/server_mgmt.go b/server/server_mgmt.go index 65be6bc..5e07ffe 100644 --- a/server/server_mgmt.go +++ b/server/server_mgmt.go @@ -2053,17 +2053,27 @@ func (s *Server) pgListModels(w http.ResponseWriter, r *http.Request) { writeErr(w, 400, err.Error()) return } - // Resolve API key: form input > profile stored key. + // Resolve API key: form input > profile stored key. The stored key is only + // sent to the endpoint it was saved for, since base_url comes from the form. apiKey := strings.TrimSpace(req.APIKey) + withheld := false if apiKey == "" && req.ProfileID != nil { if p, err := s.m.pg.ProfileByID(*req.ProfileID); err == nil && p != nil { - apiKey = p.APIKey + if sameLLMEndpoint(req.Provider, req.BaseURL, p.Format, p.BaseURL) { + apiKey = p.APIKey + } else { + withheld = p.APIKey != "" + } } } provider := strings.TrimSpace(req.Provider) // Ollama serves its model list keyless on the local server: no key gate. if apiKey == "" && provider != "ollama" { - writeJSON(w, 200, map[string]any{"ok": false, "error": "no API Key provided"}) + msg := "no API Key provided" + if withheld { + msg = "no API Key provided: the saved key is only sent to the endpoint it was saved for; enter the key to list models from a different base URL or provider" + } + writeJSON(w, 200, map[string]any{"ok": false, "error": msg}) return }