fix(llm): validate custom User-Agent in the connection-test path

testLLM set cfg.UserAgent from the request (or the stored-profile
fallback) with no validation, while pgSaveProfile runs
agent.ValidateUserAgent and returns 400 on a bad value. A UA with a
control character could therefore reach the wire from the test button
though it is rejected on save. Call ValidateUserAgent on the final value
(after the stored fallback) and return 400 with the same message style,
mirroring pgSaveProfile; the fallback itself is unchanged.

ValidateUserAgent rejects every control character (r < 0x20 or r ==
0x7f), TAB included, so the header value stays a single printable token.
Extend the agent tests to cover NUL, US, DEL and TAB, and that a
printable slash-token value with spaces survives trimming unchanged. Add
a Postgres-free server test driving testLLM through an httptest endpoint:
a control-character UA is a 400 before any request leaves, a printable UA
reaches the wire verbatim.
This commit is contained in:
BaaRai
2026-11-08 18:15:33 +01:00
parent 13dde2da15
commit 4acdf3bb26
3 changed files with 69 additions and 0 deletions
+23
View File
@@ -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)
+38
View File
@@ -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) {
+8
View File
@@ -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.