262 lines
11 KiB
Go
262 lines
11 KiB
Go
package server
|
||
|
||
import (
|
||
"io"
|
||
"log/slog"
|
||
"net/http"
|
||
"net/http/httptest"
|
||
"strings"
|
||
"sync/atomic"
|
||
"testing"
|
||
"testing/fstest"
|
||
"unicode/utf8"
|
||
|
||
"code.cubecraftlabs.com/Maniacs_Mansion/Meal-Quest/internal/household"
|
||
)
|
||
|
||
const (
|
||
aiPath = "/api/v1/settings/ai"
|
||
aiKeyPath = "/api/v1/settings/ai/api-key"
|
||
aiVerifyPath = "/api/v1/settings/ai/verify"
|
||
aiModelsPath = "/api/v1/settings/ai/models"
|
||
aiResetPromptPath = "/api/v1/settings/ai/reset-prompt"
|
||
testAIKey = "sk-test-secret-ai-key-7d21"
|
||
testWrongAIKey = "sk-test-wrong-ai-key-0b9e"
|
||
)
|
||
|
||
// newAIRouter returns a router with a fresh, seeded AI settings service.
|
||
func newAIRouter(t *testing.T) http.Handler {
|
||
t.Helper()
|
||
logger := slog.New(slog.NewTextHandler(io.Discard, nil))
|
||
r, err := NewRouter(logger, "dev", nil, stubAuth{}, stubSessions{}, newTestSettings(t), newTestSearch(t), newTestAI(t), newTestChat(t), newTestDinners(t), newTestSuggest(t),
|
||
newTestPlan(t), newTestShopping(t), newTestGoogle(t), newTestPush(t), newTestIdeas(t), newTestHome(t), newTestRunner(t), newTestPhotos(t), fstest.MapFS{"index.html": {Data: []byte(fixtureIndex)}})
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
return r
|
||
}
|
||
|
||
// newModelsServer is an OpenAI-compatible server whose GET /models accepts
|
||
// only testAIKey. While echo is set it answers 500 with a body quoting the
|
||
// Authorization header it received, so a response that passed the vendor
|
||
// body through unredacted would leak the key. Its model list holds a model
|
||
// whose ID and name quote the key.
|
||
func newModelsServer(t *testing.T, echo *atomic.Bool) *httptest.Server {
|
||
t.Helper()
|
||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||
auth := r.Header.Get("Authorization")
|
||
switch {
|
||
case r.URL.Path != "/v1/models":
|
||
http.NotFound(w, r)
|
||
case echo.Load():
|
||
w.WriteHeader(http.StatusInternalServerError)
|
||
io.WriteString(w, `{"error":{"message":"upstream failed for `+auth+`"}}`)
|
||
case auth != "Bearer "+testAIKey:
|
||
w.WriteHeader(http.StatusUnauthorized)
|
||
io.WriteString(w, `{"error":{"message":"invalid key `+auth+`"}}`)
|
||
default:
|
||
io.WriteString(w, `{"data":[{"id":"gpt-test"},{"id":"gpt-test-mini","name":"GPT Test Mini"},`+
|
||
`{"id":"gpt-`+testAIKey+`","name":"Leaky `+testAIKey+`"}]}`)
|
||
}
|
||
}))
|
||
t.Cleanup(srv.Close)
|
||
return srv
|
||
}
|
||
|
||
func errorCode(body map[string]any) any {
|
||
e, _ := body["error"].(map[string]any)
|
||
return e["code"]
|
||
}
|
||
|
||
// TestAISettingsAPI drives every /settings/ai route and scans every
|
||
// response body for the API keys.
|
||
func TestAISettingsAPI(t *testing.T) {
|
||
t.Parallel()
|
||
router := newAIRouter(t)
|
||
var echo atomic.Bool
|
||
vendor := newModelsServer(t, &echo)
|
||
var bodies []string
|
||
call := func(method, target, body string) (*httptest.ResponseRecorder, map[string]any) {
|
||
t.Helper()
|
||
var r io.Reader
|
||
if body != "" {
|
||
r = strings.NewReader(body)
|
||
}
|
||
rec, decoded := serve(t, router, signedIn(method, target, r))
|
||
bodies = append(bodies, rec.Body.String())
|
||
return rec, decoded
|
||
}
|
||
|
||
// The seed: anthropic, the default persona, no key.
|
||
rec, body := call(http.MethodGet, aiPath, "")
|
||
if rec.Code != http.StatusOK || body["provider"] != "anthropic" || body["system_prompt"] != household.DefaultPersonaPrompt ||
|
||
body["temperature"] != 0.7 || body["api_key_last4"] != nil || body["key_verified_at"] != nil {
|
||
t.Errorf("GET seed = %d %v", rec.Code, body)
|
||
}
|
||
for _, field := range []string{"provider", "base_url", "model", "temperature", "system_prompt", "api_key_last4", "key_verified_at"} {
|
||
if _, ok := body[field]; !ok {
|
||
t.Errorf("GET has no %s field: %v", field, body)
|
||
}
|
||
}
|
||
for _, path := range []string{aiVerifyPath, aiModelsPath} {
|
||
method := http.MethodPost
|
||
if path == aiModelsPath {
|
||
method = http.MethodGet
|
||
}
|
||
if rec, body := call(method, path, ""); rec.Code != http.StatusConflict || errorCode(body) != "not_configured" {
|
||
t.Errorf("%s %s without a key = %d %v, want 409 not_configured", method, path, rec.Code, body)
|
||
}
|
||
}
|
||
|
||
for _, bad := range []string{
|
||
`{}`,
|
||
`{"provider":"google"}`,
|
||
`{"provider":"openai_compatible","base_url":"not a url"}`,
|
||
`{"base_url":"https://evil.example"}`,
|
||
`{"temperature":1.1}`,
|
||
`{"temperature":-0.1}`,
|
||
`{"temperature":0.25}`,
|
||
`{"temperature":"0.5"}`,
|
||
`{"model":""}`,
|
||
`{"model":null}`,
|
||
`{"system_prompt":"` + strings.Repeat("é", 8001) + `"}`,
|
||
`{"system_prompt":"` + strings.Repeat("€", 8001) + `"}`,
|
||
`{"system_prompt":"` + strings.Repeat("a", maxAISettingsBody) + `"}`,
|
||
`{"api_key":"` + testAIKey + `"}`,
|
||
} {
|
||
if rec, body := call(http.MethodPatch, aiPath, bad); rec.Code != http.StatusBadRequest {
|
||
t.Errorf("PATCH %.80s = %d %v, want 400", bad, rec.Code, body)
|
||
}
|
||
}
|
||
|
||
// The longest prompt fits the body limit however its 8,000 characters
|
||
// are written: three-byte, four-byte, or as \u-escaped surrogate pairs.
|
||
for _, char := range []string{"€", "🍲", `\ud83c\udf72`} {
|
||
rec, body := call(http.MethodPatch, aiPath, `{"system_prompt":"`+strings.Repeat(char, 8000)+`"}`)
|
||
if prompt, _ := body["system_prompt"].(string); rec.Code != http.StatusOK || utf8.RuneCountInString(prompt) != 8000 {
|
||
t.Errorf("PATCH 8,000 × %s = %d, %d characters", char, rec.Code, utf8.RuneCountInString(prompt))
|
||
}
|
||
}
|
||
|
||
rec, body = call(http.MethodPatch, aiPath, `{"provider":"openai_compatible","base_url":"`+vendor.URL+`/v1","model":"gpt-test","temperature":0.3,"system_prompt":"`+strings.Repeat("é", 8000)+`"}`)
|
||
if rec.Code != http.StatusOK || body["provider"] != "openai_compatible" || body["base_url"] != vendor.URL+"/v1" ||
|
||
body["model"] != "gpt-test" || body["temperature"] != 0.3 {
|
||
t.Errorf("PATCH openai_compatible = %d %v", rec.Code, body)
|
||
}
|
||
|
||
for _, bad := range []string{`{}`, `{"api_key":"short"}`, `{"api_key":"sk-line\nbreak"}`, `{"api_key":null}`, `{"key":"` + testAIKey + `"}`} {
|
||
if rec, body := call(http.MethodPut, aiKeyPath, bad); rec.Code != http.StatusBadRequest {
|
||
t.Errorf("PUT api-key %s = %d %v, want 400", bad, rec.Code, body)
|
||
}
|
||
}
|
||
|
||
// A rejected key is a typed auth failure and records nothing.
|
||
rec, body = call(http.MethodPut, aiKeyPath, `{"api_key":"`+testWrongAIKey+`"}`)
|
||
if rec.Code != http.StatusOK || body["api_key_last4"] != "0b9e" || body["key_verified_at"] != nil {
|
||
t.Errorf("PUT wrong key = %d %v", rec.Code, body)
|
||
}
|
||
if rec, body := call(http.MethodPost, aiVerifyPath, ""); rec.Code != http.StatusBadGateway || errorCode(body) != "auth_failed" {
|
||
t.Errorf("POST verify with the wrong key = %d %v, want 502 auth_failed", rec.Code, body)
|
||
}
|
||
if rec, body := call(http.MethodGet, aiModelsPath, ""); rec.Code != http.StatusBadGateway || errorCode(body) != "auth_failed" {
|
||
t.Errorf("GET models with the wrong key = %d %v, want 502 auth_failed", rec.Code, body)
|
||
}
|
||
|
||
rec, body = call(http.MethodPut, aiKeyPath, `{"api_key":"`+testAIKey+`"}`)
|
||
if rec.Code != http.StatusOK || body["api_key_last4"] != "7d21" {
|
||
t.Errorf("PUT key = %d %v", rec.Code, body)
|
||
}
|
||
if _, ok := body["api_key"]; ok {
|
||
t.Errorf("PUT response has an api_key field: %v", body)
|
||
}
|
||
|
||
// The vendor quoting the key in an error body: the body excerpt
|
||
// reaches the response redacted.
|
||
echo.Store(true)
|
||
for _, c := range []struct{ method, path string }{{http.MethodPost, aiVerifyPath}, {http.MethodGet, aiModelsPath}} {
|
||
rec, body := call(c.method, c.path, "")
|
||
if rec.Code != http.StatusBadGateway || errorCode(body) != "provider_error" || !strings.Contains(rec.Body.String(), "upstream failed for Bearer [REDACTED]") {
|
||
t.Errorf("%s %s against an echoing vendor = %d %v, want 502 provider_error", c.method, c.path, rec.Code, body)
|
||
}
|
||
}
|
||
echo.Store(false)
|
||
|
||
rec, body = call(http.MethodPost, aiVerifyPath, "")
|
||
if rec.Code != http.StatusOK || body["key_verified_at"] == nil || body["api_key_last4"] != "7d21" {
|
||
t.Errorf("POST verify = %d %v", rec.Code, body)
|
||
}
|
||
rec, body = call(http.MethodGet, aiPath, "")
|
||
if rec.Code != http.StatusOK || body["key_verified_at"] == nil || body["model"] != "gpt-test" || body["temperature"] != 0.3 {
|
||
t.Errorf("GET after verify = %d %v", rec.Code, body)
|
||
}
|
||
|
||
rec, body = call(http.MethodGet, aiModelsPath, "")
|
||
models, _ := body["models"].([]any)
|
||
if rec.Code != http.StatusOK || len(models) != 3 {
|
||
t.Fatalf("GET models = %d %v", rec.Code, body)
|
||
}
|
||
if m := models[1].(map[string]any); m["id"] != "gpt-test-mini" || m["display_name"] != "GPT Test Mini" {
|
||
t.Errorf("GET models[1] = %v", m)
|
||
}
|
||
if m := models[2].(map[string]any); m["id"] != "gpt-[REDACTED]" || m["display_name"] != "Leaky [REDACTED]" {
|
||
t.Errorf("GET models[2] = %v, want the key redacted", m)
|
||
}
|
||
|
||
rec, body = call(http.MethodPatch, aiPath, `{"system_prompt":"Be brief."}`)
|
||
if rec.Code != http.StatusOK || body["system_prompt"] != "Be brief." || body["key_verified_at"] == nil {
|
||
t.Errorf("PATCH system_prompt = %d %v", rec.Code, body)
|
||
}
|
||
rec, body = call(http.MethodPost, aiResetPromptPath, "")
|
||
if rec.Code != http.StatusOK || body["system_prompt"] != household.DefaultPersonaPrompt {
|
||
t.Errorf("POST reset-prompt = %d %v", rec.Code, body)
|
||
}
|
||
|
||
// The stored key is never pointed at another server or provider.
|
||
for _, patch := range []string{`{"base_url":"https://attacker.example/v1"}`, `{"provider":"openai","base_url":""}`} {
|
||
if rec, body := call(http.MethodPatch, aiPath, patch); rec.Code != http.StatusBadRequest || errorCode(body) != "invalid_settings" {
|
||
t.Errorf("PATCH %s with a key stored = %d %v, want 400 invalid_settings", patch, rec.Code, body)
|
||
}
|
||
}
|
||
if rec, body := call(http.MethodPut, aiKeyPath, `{"api_key":""}`); rec.Code != http.StatusOK || body["api_key_last4"] != nil || body["key_verified_at"] != nil {
|
||
t.Errorf("PUT empty key = %d %v", rec.Code, body)
|
||
}
|
||
rec, body = call(http.MethodPatch, aiPath, `{"base_url":"https://llm.example/v1"}`)
|
||
if rec.Code != http.StatusOK || body["base_url"] != "https://llm.example/v1" {
|
||
t.Errorf("PATCH base_url without a key = %d %v", rec.Code, body)
|
||
}
|
||
|
||
if len(bodies) < 30 {
|
||
t.Fatalf("scanned only %d response bodies", len(bodies))
|
||
}
|
||
for _, b := range bodies {
|
||
for _, key := range []string{testAIKey, testWrongAIKey} {
|
||
if strings.Contains(b, key) {
|
||
t.Errorf("response holds an API key: %s", b)
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
func TestAISettingsWritesNeedCSRF(t *testing.T) {
|
||
t.Parallel()
|
||
router := newAIRouter(t)
|
||
for _, c := range []struct{ method, path, body string }{
|
||
{http.MethodPatch, aiPath, `{"temperature":0.5}`},
|
||
{http.MethodPut, aiKeyPath, `{"api_key":"` + testAIKey + `"}`},
|
||
{http.MethodPost, aiVerifyPath, ""},
|
||
{http.MethodPost, aiResetPromptPath, ""},
|
||
} {
|
||
req := httptest.NewRequest(c.method, c.path, strings.NewReader(c.body))
|
||
req.AddCookie(validSessionCookie)
|
||
rec := httptest.NewRecorder()
|
||
router.ServeHTTP(rec, req)
|
||
if rec.Code != http.StatusForbidden {
|
||
t.Errorf("%s %s without the CSRF header = %d, want 403", c.method, c.path, rec.Code)
|
||
}
|
||
}
|
||
rec, body := serve(t, router, signedIn(http.MethodGet, aiPath, nil))
|
||
if rec.Code != http.StatusOK || body["api_key_last4"] != nil || body["temperature"] != 0.7 {
|
||
t.Errorf("GET after refused writes = %d %v, want the seed", rec.Code, body)
|
||
}
|
||
}
|