diff --git a/agent/user_agent_test.go b/agent/user_agent_test.go index 8699be5..0614ea9 100644 --- a/agent/user_agent_test.go +++ b/agent/user_agent_test.go @@ -52,6 +52,29 @@ func TestValidateUserAgentRejectsNewline(t *testing.T) { } } +// The rule is "no control character at all" (r < 0x20 or r == 0x7f), so the +// non-newline controls are rejected too, TAB included: a header value is one +// printable token, not something to lay out with tabs. +func TestValidateUserAgentRejectsControlCharacters(t *testing.T) { + for name, in := range map[string]string{ + "NUL": "good/1.0\x00", + "US": "good/1.0\x1f", + "DEL": "good/1.0\x7f", + "TAB": "good/1.0\tX", + } { + if _, err := ValidateUserAgent(in); err == nil { + t.Errorf("%s: expected a user agent containing %q to be rejected", name, in) + } + } +} + +func TestValidateUserAgentAcceptsPrintableToken(t *testing.T) { + const want = "MyTool/1.0 (contact-free)" + if ua, err := ValidateUserAgent(" " + want + " "); err != nil || ua != want { + t.Fatalf("printable value: got (%q, %v), want (%q, nil)", ua, err, want) + } +} + func TestValidateUserAgentRejectsTooLong(t *testing.T) { if _, err := ValidateUserAgent(strings.Repeat("a", MaxUserAgentLen+1)); err == nil { t.Fatalf("expected a user agent longer than %d to be rejected", MaxUserAgentLen) diff --git a/server/llm_key_endpoint_test.go b/server/llm_key_endpoint_test.go index e4af8c4..6e59826 100644 --- a/server/llm_key_endpoint_test.go +++ b/server/llm_key_endpoint_test.go @@ -81,6 +81,44 @@ func TestLLMTestGlobalKeyOnlyForGlobalEndpoint(t *testing.T) { } } +// The connection test applies the same User-Agent rule as saving a profile: a +// value with a control character is a 400 before any request leaves, and a +// printable value reaches the wire exactly as given. +func TestLLMTestValidatesUserAgent(t *testing.T) { + var hits atomic.Int32 + var gotUA atomic.Value + fake := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + hits.Add(1) + gotUA.Store(r.Header.Get("User-Agent")) + 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}}`)) + })) + t.Cleanup(fake.Close) + s := &Server{} + post := func(ua string) *httptest.ResponseRecorder { + b, _ := json.Marshal(map[string]any{ + "provider": "anthropic", "model": "m", "base_url": fake.URL, "api_key": "k", "streaming": false, "user_agent": ua, + }) + req := httptest.NewRequest(http.MethodPost, "/api/llm/test", bytes.NewReader(b)) + req.Header.Set("Content-Type", "application/json") + rec := httptest.NewRecorder() + s.testLLM(rec, req) + return rec + } + + if rec := post("good/1.0\r\nX-Injected: 1"); rec.Code != http.StatusBadRequest || hits.Load() != 0 { + t.Fatalf("control character in user_agent: status %d hits=%d body=%s", rec.Code, hits.Load(), rec.Body.String()) + } + + const want = "MyTool/1.0 (contact-free)" + if rec := post(" " + want + " "); rec.Code != http.StatusOK { + t.Fatalf("printable user_agent: status %d body=%s", rec.Code, rec.Body.String()) + } + if ua, _ := gotUA.Load().(string); hits.Load() != 1 || ua != want { + t.Fatalf("test request User-Agent = %q (hits=%d), want %q", ua, hits.Load(), want) + } +} + // 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) { diff --git a/server/server.go b/server/server.go index 322dd80..fbc2325 100644 --- a/server/server.go +++ b/server/server.go @@ -1488,6 +1488,14 @@ func (s *Server) testLLM(w http.ResponseWriter, r *http.Request) { if stored != nil && cfg.UserAgent == "" { cfg.UserAgent = strings.TrimSpace(stored.UserAgent) } + // Same rule as saving a profile: a value with control characters or over the + // length cap never reaches the wire, not even from the test button. + ua, err := agent.ValidateUserAgent(cfg.UserAgent) + if err != nil { + writeErr(w, 400, err.Error()) + return + } + cfg.UserAgent = ua // Keyless formats (Ollama) skip the key resolution entirely. if cfg.NeedsAPIKey() { // API key resolution precedence: form input > the named profile's stored key > the global config's key.