341 lines
11 KiB
Go
341 lines
11 KiB
Go
package server
|
|
|
|
import (
|
|
"ai-router/internal/protocol"
|
|
"encoding/json"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"strings"
|
|
"testing"
|
|
"testing/fstest"
|
|
"time"
|
|
)
|
|
|
|
func testServer(t *testing.T) *Server {
|
|
t.Helper()
|
|
s, e := New(":memory:", "test-password", fstest.MapFS{"index.html": &fstest.MapFile{Data: []byte("SPA")}})
|
|
if e != nil {
|
|
t.Fatal(e)
|
|
}
|
|
t.Cleanup(func() { s.Close() })
|
|
return s
|
|
}
|
|
func call(s *Server, m, p, b string, c *http.Cookie) *httptest.ResponseRecorder {
|
|
r := httptest.NewRequest(m, p, strings.NewReader(b))
|
|
r.Header.Set("Origin", "http://example.com")
|
|
r.Header.Set("X-Requested-With", "ai-router")
|
|
if c != nil {
|
|
r.AddCookie(c)
|
|
}
|
|
w := httptest.NewRecorder()
|
|
s.ServeHTTP(w, r)
|
|
return w
|
|
}
|
|
func loginCookie(t *testing.T, s *Server) *http.Cookie {
|
|
t.Helper()
|
|
w := call(s, "POST", "/admin/api/login", `{"password":"test-password"}`, nil)
|
|
if w.Code != 200 {
|
|
t.Fatal(w.Code, w.Body.String())
|
|
}
|
|
return w.Result().Cookies()[0]
|
|
}
|
|
func TestAuthentication(t *testing.T) {
|
|
s := testServer(t)
|
|
for _, origin := range []string{"", "http://evil.test", "https://example.com"} {
|
|
r := httptest.NewRequest("POST", "/admin/api/login", strings.NewReader(`{"password":"test-password"}`))
|
|
r.Header.Set("Origin", origin)
|
|
r.Header.Set("X-Requested-With", "ai-router")
|
|
w := httptest.NewRecorder()
|
|
s.ServeHTTP(w, r)
|
|
if w.Code != 403 {
|
|
t.Fatal(origin, w.Code)
|
|
}
|
|
}
|
|
c := loginCookie(t, s)
|
|
if !c.HttpOnly || c.SameSite != http.SameSiteStrictMode {
|
|
t.Fatal(c)
|
|
}
|
|
if w := call(s, "GET", "/admin/api/channels", "", nil); w.Code != 401 {
|
|
t.Fatal(w.Code)
|
|
}
|
|
if w := call(s, "POST", "/admin/api/logout", "", c); w.Code != 200 {
|
|
t.Fatal(w.Code)
|
|
}
|
|
if w := call(s, "GET", "/admin/api/channels", "", c); w.Code != 401 {
|
|
t.Fatal(w.Code)
|
|
}
|
|
if w := call(s, "GET", "/admin/api/session", "", nil); w.Code != 200 || w.Header().Get("Content-Security-Policy") == "" || w.Header().Get("X-Frame-Options") != "DENY" || w.Header().Get("Referrer-Policy") != "no-referrer" {
|
|
t.Fatal("security headers missing", w.Code, w.Header())
|
|
}
|
|
for i := 0; i < 10; i++ {
|
|
call(s, "POST", "/admin/api/login", `{"password":"wrong"}`, nil)
|
|
}
|
|
if w := call(s, "POST", "/admin/api/login", `{"password":"wrong"}`, nil); w.Code != 429 {
|
|
t.Fatal(w.Code)
|
|
}
|
|
// While locked, a correct password gets the same 429: the lockout must
|
|
// not become an oracle that confirms candidate passwords.
|
|
if w := call(s, "POST", "/admin/api/login", `{"password":"test-password"}`, nil); w.Code != 429 {
|
|
t.Fatal(w.Code)
|
|
}
|
|
// The lockout is per IP; a fresh address still logs in.
|
|
r := httptest.NewRequest("POST", "/admin/api/login", strings.NewReader(`{"password":"test-password"}`))
|
|
r.Header.Set("Origin", "http://example.com")
|
|
r.Header.Set("X-Requested-With", "ai-router")
|
|
r.RemoteAddr = "192.0.2.9:1"
|
|
w := httptest.NewRecorder()
|
|
s.ServeHTTP(w, r)
|
|
if w.Code != 200 {
|
|
t.Fatal("fresh IP rejected:", w.Code)
|
|
}
|
|
}
|
|
func TestMappingPermissions(t *testing.T) {
|
|
cs := []channel{{ID: "stable", Name: "renamed", Enabled: true, Models: map[string]string{"alias": "vendor/real", "vendor/real": "wrong"}}}
|
|
k := apiKey{Channels: []string{"stable"}}
|
|
_, actual, e := resolve(cs, k, "renamed/alias")
|
|
if e != nil || actual != "vendor/real" {
|
|
t.Fatal(actual, e)
|
|
}
|
|
if _, _, e = resolve(cs, apiKey{Channels: []string{"renamed"}}, "renamed/alias"); e == nil {
|
|
t.Fatal("name granted access")
|
|
}
|
|
if _, _, e = resolve(cs, k, "renamed/vendor/unknown"); e == nil {
|
|
t.Fatal("unknown allowed")
|
|
}
|
|
cs[0].Enabled = false
|
|
if _, _, e = resolve(cs, k, "renamed/alias"); e == nil {
|
|
t.Fatal("disabled allowed")
|
|
}
|
|
}
|
|
func TestLegacyPassthroughIgnored(t *testing.T) {
|
|
s := testServer(t)
|
|
cookie := loginCookie(t, s)
|
|
up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
t.Error("unmapped model reached upstream")
|
|
w.WriteHeader(http.StatusBadGateway)
|
|
}))
|
|
defer up.Close()
|
|
body := map[string]any{
|
|
"name": "legacy", "base_url": up.URL, "protocol": "openai-chat",
|
|
"enabled": true, "passthrough": true, "models": map[string]string{"alias": "vendor/real"},
|
|
}
|
|
save := func() channel {
|
|
t.Helper()
|
|
b, err := json.Marshal(body)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
w := call(s, "POST", "/admin/api/channels", string(b), cookie)
|
|
if w.Code != 200 || strings.Contains(w.Body.String(), `"passthrough"`) {
|
|
t.Fatal(w.Code, w.Body.String())
|
|
}
|
|
var c channel
|
|
if err := json.Unmarshal(w.Body.Bytes(), &c); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
return c
|
|
}
|
|
c := save()
|
|
var stored bool
|
|
if err := s.db.QueryRow("SELECT passthrough FROM channels WHERE id=?", c.ID).Scan(&stored); err != nil || stored {
|
|
t.Fatal("new channel enabled passthrough", stored, err)
|
|
}
|
|
body["id"] = c.ID
|
|
save()
|
|
if err := s.db.QueryRow("SELECT passthrough FROM channels WHERE id=?", c.ID).Scan(&stored); err != nil || stored {
|
|
t.Fatal("update enabled passthrough", stored, err)
|
|
}
|
|
if _, err := s.db.Exec("UPDATE channels SET passthrough=1 WHERE id=?", c.ID); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
cs, err := s.loadChannels()
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if _, actual, err := resolve(cs, apiKey{AllChannels: true}, "legacy/alias"); err != nil || actual != "vendor/real" {
|
|
t.Fatal(actual, err)
|
|
}
|
|
w := call(s, "GET", "/admin/api/channels", "", cookie)
|
|
if w.Code != 200 || strings.Contains(w.Body.String(), `"passthrough"`) {
|
|
t.Fatal(w.Code, w.Body.String())
|
|
}
|
|
w = call(s, "POST", "/admin/api/keys", `{"name":"client","all_channels":true}`, cookie)
|
|
var key map[string]string
|
|
if err := json.Unmarshal(w.Body.Bytes(), &key); err != nil || w.Code != 200 {
|
|
t.Fatal(w.Code, w.Body.String(), err)
|
|
}
|
|
for _, path := range []string{"/v1/chat/completions", "/v1/responses", "/v1/messages"} {
|
|
for _, model := range []string{"legacy/unknown", "legacy/vendor/real"} {
|
|
t.Run(path+"/"+model, func(t *testing.T) {
|
|
r := httptest.NewRequest("POST", path, strings.NewReader(`{"model":"`+model+`","messages":[]}`))
|
|
r.Header.Set("Authorization", "Bearer "+key["key"])
|
|
w := httptest.NewRecorder()
|
|
s.ServeHTTP(w, r)
|
|
if w.Code != http.StatusForbidden {
|
|
t.Fatal(w.Code, w.Body.String())
|
|
}
|
|
})
|
|
}
|
|
}
|
|
save()
|
|
cs, err = s.loadChannels()
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if _, _, err := resolve(cs, apiKey{AllChannels: true}, "legacy/unknown"); err == nil {
|
|
t.Fatal("updated legacy channel allowed unknown model")
|
|
}
|
|
}
|
|
|
|
func TestKeyRevocationAndProxy(t *testing.T) {
|
|
s := testServer(t)
|
|
c := loginCookie(t, s)
|
|
up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
if r.URL.Path != "/v1/chat/completions" || r.Header.Get("Authorization") != "Bearer upstream-secret" {
|
|
t.Error(r.URL.Path, r.Header)
|
|
}
|
|
var b map[string]any
|
|
json.NewDecoder(r.Body).Decode(&b)
|
|
if b["model"] != "vendor/real" || b["custom"] != "keep" {
|
|
t.Error(b)
|
|
}
|
|
writeJSON(w, map[string]any{"choices": []any{}, "usage": map[string]int{"prompt_tokens": 7, "completion_tokens": 3}})
|
|
}))
|
|
defer up.Close()
|
|
b, _ := json.Marshal(channel{Name: "demo", BaseURL: up.URL + "/v1", Protocol: "openai-chat", Enabled: true, Models: map[string]string{"alias": "vendor/real"}, APIKey: "upstream-secret"})
|
|
w := call(s, "POST", "/admin/api/channels", string(b), c)
|
|
if w.Code != 200 {
|
|
t.Fatal(w.Body.String())
|
|
}
|
|
var ch channel
|
|
json.Unmarshal(w.Body.Bytes(), &ch)
|
|
if strings.Contains(w.Body.String(), "upstream-secret") {
|
|
t.Fatal("secret leak")
|
|
}
|
|
w = call(s, "POST", "/admin/api/keys", `{"name":"client","channels":["`+ch.ID+`"]}`, c)
|
|
if w.Code != 200 {
|
|
t.Fatal(w.Body.String())
|
|
}
|
|
var key map[string]string
|
|
json.Unmarshal(w.Body.Bytes(), &key)
|
|
r := httptest.NewRequest("POST", "/v1/chat/completions", strings.NewReader(`{"model":"demo/alias","messages":[],"custom":"keep"}`))
|
|
r.Header.Set("x-api-key", key["key"])
|
|
w = httptest.NewRecorder()
|
|
s.ServeHTTP(w, r)
|
|
if w.Code != 200 {
|
|
t.Fatal(w.Code, w.Body.String())
|
|
}
|
|
var id, stored string
|
|
if e := s.db.QueryRow("SELECT id,hash FROM keys").Scan(&id, &stored); e != nil || stored == key["key"] {
|
|
t.Fatal(e, stored)
|
|
}
|
|
w = call(s, "DELETE", "/admin/api/keys/"+id, "", c)
|
|
if w.Code != 200 {
|
|
t.Fatal(w.Body.String())
|
|
}
|
|
w = call(s, "PATCH", "/admin/api/keys/"+id, `{"enabled":true}`, c)
|
|
if w.Code != 404 {
|
|
t.Fatal("revived revoked key", w.Code)
|
|
}
|
|
r = httptest.NewRequest("GET", "/v1/models", nil)
|
|
r.Header.Set("Authorization", "Bearer "+key["key"])
|
|
w = httptest.NewRecorder()
|
|
s.ServeHTTP(w, r)
|
|
if w.Code != 401 {
|
|
t.Fatal(w.Code)
|
|
}
|
|
var n int
|
|
if err := s.flushLogs(); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
s.db.QueryRow("SELECT count(*) FROM logs").Scan(&n)
|
|
if n != 1 {
|
|
t.Fatal(n)
|
|
}
|
|
}
|
|
func TestStatsUnknownAndFilters(t *testing.T) {
|
|
s := testServer(t)
|
|
c := loginCookie(t, s)
|
|
now := time.Now().UTC().Format(time.RFC3339Nano)
|
|
s.record(requestLog{CreatedAt: now, KeyID: "a", KeyName: "A", Status: 200, Usage: protocol.Usage{Input: 8, Output: 2, Known: true}})
|
|
s.record(requestLog{CreatedAt: now, KeyID: "a", KeyName: "A", Status: 400})
|
|
s.record(requestLog{CreatedAt: now, KeyID: "b", Status: 200})
|
|
if err := s.flushLogs(); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
w := call(s, "GET", "/admin/api/stats?key_id=a", "", c)
|
|
if w.Code != 200 {
|
|
t.Fatal(w.Body.String())
|
|
}
|
|
var out struct {
|
|
Summary summary `json:"summary"`
|
|
}
|
|
if e := json.Unmarshal(w.Body.Bytes(), &out); e != nil {
|
|
t.Fatal(e)
|
|
}
|
|
if out.Summary.Requests != 2 || out.Summary.Success != 1 || out.Summary.Unknown != 1 || out.Summary.Input == nil || *out.Summary.Input != 8 {
|
|
t.Fatal(w.Body.String())
|
|
}
|
|
w = call(s, "GET", "/admin/api/logs?key_id=b", "", c)
|
|
if !strings.Contains(w.Body.String(), `"input_tokens":null`) {
|
|
t.Fatal(w.Body.String())
|
|
}
|
|
w = call(s, "GET", "/admin/api/stats?from=invalid", "", c)
|
|
if w.Code != 400 {
|
|
t.Fatal(w.Code)
|
|
}
|
|
w = call(s, "GET", "/admin/api/unknown", "", c)
|
|
if w.Code != 404 {
|
|
t.Fatal(w.Code)
|
|
}
|
|
w = call(s, "GET", "/some/spa/path", "", nil)
|
|
if w.Code != 200 || w.Body.String() != "SPA" {
|
|
t.Fatal(w.Code, w.Body.String())
|
|
}
|
|
}
|
|
|
|
func TestPublicOriginCaseInsensitive(t *testing.T) {
|
|
s := testServer(t)
|
|
if e := s.SetPublicOrigin("https://Router.Example.com"); e != nil {
|
|
t.Fatal(e)
|
|
}
|
|
r := httptest.NewRequest("POST", "/admin/api/login", strings.NewReader(`{"password":"test-password"}`))
|
|
r.Header.Set("Origin", "https://router.example.com")
|
|
r.Header.Set("X-Requested-With", "ai-router")
|
|
w := httptest.NewRecorder()
|
|
s.ServeHTTP(w, r)
|
|
if w.Code != 200 || !w.Result().Cookies()[0].Secure {
|
|
t.Fatal(w.Code, w.Header())
|
|
}
|
|
}
|
|
|
|
func TestPublicOriginAndRedirect(t *testing.T) {
|
|
s := testServer(t)
|
|
if e := s.SetPublicOrigin("https://router.example"); e != nil {
|
|
t.Fatal(e)
|
|
}
|
|
r := httptest.NewRequest("POST", "/admin/api/login", strings.NewReader(`{"password":"test-password"}`))
|
|
r.Header.Set("Origin", "https://router.example")
|
|
r.Header.Set("X-Requested-With", "ai-router")
|
|
w := httptest.NewRecorder()
|
|
s.ServeHTTP(w, r)
|
|
if w.Code != 200 || !w.Result().Cookies()[0].Secure {
|
|
t.Fatal(w.Code, w.Header())
|
|
}
|
|
reached := false
|
|
target := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { reached = true }))
|
|
defer target.Close()
|
|
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { http.Redirect(w, r, target.URL, 307) }))
|
|
defer upstream.Close()
|
|
req, _ := http.NewRequest("POST", upstream.URL, strings.NewReader("body"))
|
|
req.Header.Set("Authorization", "Bearer private")
|
|
res, e := s.client.Do(req)
|
|
if e != nil {
|
|
t.Fatal(e)
|
|
}
|
|
res.Body.Close()
|
|
if reached || res.StatusCode != 307 {
|
|
t.Fatal("redirect followed")
|
|
}
|
|
}
|