chore: initial commit
This commit is contained in:
@@ -0,0 +1,265 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"database/sql"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"strings"
|
||||
"time"
|
||||
)
|
||||
|
||||
type channel struct {
|
||||
ID string `json:"id"`
|
||||
Name string `json:"name"`
|
||||
BaseURL string `json:"base_url"`
|
||||
Protocol string `json:"protocol"`
|
||||
Enabled bool `json:"enabled"`
|
||||
Models map[string]string `json:"models"`
|
||||
APIKey string `json:"api_key,omitempty"`
|
||||
}
|
||||
type apiKey struct {
|
||||
Revoked bool `json:"revoked"`
|
||||
ID string `json:"id"`
|
||||
Name string `json:"name"`
|
||||
Prefix string `json:"prefix"`
|
||||
Enabled bool `json:"enabled"`
|
||||
AllChannels bool `json:"all_channels"`
|
||||
Channels []string `json:"channels"`
|
||||
ExpiresAt string `json:"expires_at"`
|
||||
CreatedAt string `json:"created_at"`
|
||||
}
|
||||
|
||||
func (s *Server) loadChannels() ([]channel, error) {
|
||||
rows, e := s.db.Query("SELECT id,name,base_url,protocol,enabled,models,api_key FROM channels ORDER BY name")
|
||||
if e != nil {
|
||||
return nil, e
|
||||
}
|
||||
defer rows.Close()
|
||||
out := []channel{}
|
||||
for rows.Next() {
|
||||
var c channel
|
||||
var m string
|
||||
if e = rows.Scan(&c.ID, &c.Name, &c.BaseURL, &c.Protocol, &c.Enabled, &m, &c.APIKey); e != nil {
|
||||
return nil, e
|
||||
}
|
||||
if e = json.Unmarshal([]byte(m), &c.Models); e != nil {
|
||||
return nil, e
|
||||
}
|
||||
out = append(out, c)
|
||||
}
|
||||
return out, rows.Err()
|
||||
}
|
||||
func (s *Server) channels(w http.ResponseWriter, r *http.Request) {
|
||||
if r.Method == "GET" {
|
||||
cs, e := s.loadChannels()
|
||||
if e != nil {
|
||||
errJSON(w, 500, "database error")
|
||||
return
|
||||
}
|
||||
for i := range cs {
|
||||
cs[i].APIKey = ""
|
||||
}
|
||||
writeJSON(w, cs)
|
||||
return
|
||||
}
|
||||
if r.Method != "POST" {
|
||||
errJSON(w, 405, "method not allowed")
|
||||
return
|
||||
}
|
||||
var c channel
|
||||
if decode(w, r, &c) != nil {
|
||||
errJSON(w, 400, "invalid channel")
|
||||
return
|
||||
}
|
||||
u, e := url.Parse(c.BaseURL)
|
||||
if e != nil || (u.Scheme != "http" && u.Scheme != "https") || u.Hostname() == "" || u.User != nil || u.RawQuery != "" || u.ForceQuery || u.Fragment != "" || strings.Contains(c.BaseURL, "#") || c.Name == "" || strings.Contains(c.Name, "/") || (c.Protocol != "openai-chat" && c.Protocol != "openai-responses" && c.Protocol != "anthropic") {
|
||||
errJSON(w, 400, "invalid channel")
|
||||
return
|
||||
}
|
||||
if c.Models == nil {
|
||||
c.Models = map[string]string{}
|
||||
}
|
||||
for a, m := range c.Models {
|
||||
if a == "" || m == "" {
|
||||
errJSON(w, 400, "invalid model mapping")
|
||||
return
|
||||
}
|
||||
}
|
||||
c.BaseURL = strings.TrimRight(c.BaseURL, "/")
|
||||
m, _ := json.Marshal(c.Models)
|
||||
if c.ID == "" {
|
||||
c.ID = randomID()
|
||||
_, e = s.db.Exec("INSERT INTO channels VALUES(?,?,?,?,?,?,?,?)", c.ID, c.Name, c.BaseURL, c.Protocol, c.Enabled, false, string(m), c.APIKey)
|
||||
} else {
|
||||
var old string
|
||||
if e = s.db.QueryRow("SELECT api_key FROM channels WHERE id=?", c.ID).Scan(&old); e == sql.ErrNoRows {
|
||||
errJSON(w, 404, "not found")
|
||||
return
|
||||
} else if e != nil {
|
||||
errJSON(w, 500, "database error")
|
||||
return
|
||||
}
|
||||
if c.APIKey == "" {
|
||||
c.APIKey = old
|
||||
}
|
||||
_, e = s.db.Exec("UPDATE channels SET name=?,base_url=?,protocol=?,enabled=?,models=?,api_key=? WHERE id=?", c.Name, c.BaseURL, c.Protocol, c.Enabled, string(m), c.APIKey, c.ID)
|
||||
}
|
||||
if e != nil {
|
||||
errJSON(w, 400, "channel could not be saved; names must be unique")
|
||||
return
|
||||
}
|
||||
c.APIKey = ""
|
||||
writeJSON(w, c)
|
||||
}
|
||||
func (s *Server) channelID(w http.ResponseWriter, r *http.Request) {
|
||||
if r.Method != "DELETE" {
|
||||
errJSON(w, 405, "method not allowed")
|
||||
return
|
||||
}
|
||||
res, e := s.db.Exec("DELETE FROM channels WHERE id=?", strings.TrimPrefix(r.URL.Path, "/admin/api/channels/"))
|
||||
mutation(w, res, e)
|
||||
}
|
||||
func mutation(w http.ResponseWriter, res sql.Result, e error) {
|
||||
if e != nil {
|
||||
errJSON(w, 500, "database error")
|
||||
return
|
||||
}
|
||||
n, e := res.RowsAffected()
|
||||
if e != nil {
|
||||
errJSON(w, 500, "database error")
|
||||
return
|
||||
}
|
||||
if n == 0 {
|
||||
errJSON(w, 404, "not found")
|
||||
return
|
||||
}
|
||||
writeJSON(w, map[string]bool{"ok": true})
|
||||
}
|
||||
func (s *Server) keys(w http.ResponseWriter, r *http.Request) {
|
||||
if r.Method == "GET" {
|
||||
rows, e := s.db.Query("SELECT id,name,prefix,enabled,all_channels,channels,expires_at,created_at,revoked FROM keys ORDER BY created_at DESC")
|
||||
if e != nil {
|
||||
errJSON(w, 500, "database error")
|
||||
return
|
||||
}
|
||||
defer rows.Close()
|
||||
out := []apiKey{}
|
||||
for rows.Next() {
|
||||
var k apiKey
|
||||
var ch string
|
||||
if rows.Scan(&k.ID, &k.Name, &k.Prefix, &k.Enabled, &k.AllChannels, &ch, &k.ExpiresAt, &k.CreatedAt, &k.Revoked) != nil {
|
||||
errJSON(w, 500, "database error")
|
||||
return
|
||||
}
|
||||
if json.Unmarshal([]byte(ch), &k.Channels) != nil {
|
||||
errJSON(w, 500, "database error")
|
||||
return
|
||||
}
|
||||
out = append(out, k)
|
||||
}
|
||||
if rows.Err() != nil {
|
||||
errJSON(w, 500, "database error")
|
||||
return
|
||||
}
|
||||
writeJSON(w, out)
|
||||
return
|
||||
}
|
||||
if r.Method != "POST" {
|
||||
errJSON(w, 405, "method not allowed")
|
||||
return
|
||||
}
|
||||
var k apiKey
|
||||
if decode(w, r, &k) != nil || k.Name == "" {
|
||||
errJSON(w, 400, "invalid key")
|
||||
return
|
||||
}
|
||||
if k.ExpiresAt != "" {
|
||||
t, e := time.Parse(time.RFC3339, k.ExpiresAt)
|
||||
if e != nil || !t.After(time.Now()) {
|
||||
errJSON(w, 400, "invalid expiry")
|
||||
return
|
||||
}
|
||||
k.ExpiresAt = t.UTC().Format(time.RFC3339Nano)
|
||||
}
|
||||
if k.Channels == nil {
|
||||
k.Channels = []string{}
|
||||
}
|
||||
for _, id := range k.Channels {
|
||||
var n int
|
||||
if s.db.QueryRow("SELECT count(*) FROM channels WHERE id=?", id).Scan(&n) != nil || n == 0 {
|
||||
errJSON(w, 400, "invalid channel permission")
|
||||
return
|
||||
}
|
||||
}
|
||||
raw := "sk-" + randomID()
|
||||
ch, _ := json.Marshal(k.Channels)
|
||||
_, e := s.db.Exec("INSERT INTO keys(id,name,prefix,hash,enabled,all_channels,channels,expires_at,created_at) VALUES(?,?,?,?,?,?,?,?,?)", randomID(), k.Name, raw[:11], hash(raw), true, k.AllChannels, string(ch), k.ExpiresAt, time.Now().UTC().Format(time.RFC3339Nano))
|
||||
if e != nil {
|
||||
errJSON(w, 500, "database error")
|
||||
return
|
||||
}
|
||||
writeJSON(w, map[string]string{"key": raw})
|
||||
}
|
||||
func (s *Server) keyID(w http.ResponseWriter, r *http.Request) {
|
||||
id := strings.TrimPrefix(r.URL.Path, "/admin/api/keys/")
|
||||
if r.Method == "DELETE" {
|
||||
res, e := s.db.Exec("UPDATE keys SET enabled=0,revoked=1 WHERE id=?", id)
|
||||
mutation(w, res, e)
|
||||
return
|
||||
}
|
||||
if r.Method == "PATCH" {
|
||||
var x struct {
|
||||
Enabled *bool `json:"enabled"`
|
||||
}
|
||||
if decode(w, r, &x) != nil || x.Enabled == nil {
|
||||
errJSON(w, 400, "invalid request")
|
||||
return
|
||||
}
|
||||
res, e := s.db.Exec("UPDATE keys SET enabled=? WHERE id=? AND revoked=0", *x.Enabled, id)
|
||||
mutation(w, res, e)
|
||||
return
|
||||
}
|
||||
errJSON(w, 405, "method not allowed")
|
||||
}
|
||||
func (s *Server) authenticate(r *http.Request) (apiKey, error) {
|
||||
var k apiKey
|
||||
raw := r.Header.Get("x-api-key")
|
||||
if a := r.Header.Get("Authorization"); a != "" {
|
||||
parts := strings.Fields(a)
|
||||
if len(parts) != 2 || !strings.EqualFold(parts[0], "Bearer") {
|
||||
return k, sql.ErrNoRows
|
||||
}
|
||||
if raw != "" && raw != parts[1] {
|
||||
return k, sql.ErrNoRows
|
||||
}
|
||||
raw = parts[1]
|
||||
}
|
||||
if raw == "" {
|
||||
return k, sql.ErrNoRows
|
||||
}
|
||||
var ch string
|
||||
e := s.db.QueryRow("SELECT id,name,prefix,enabled,all_channels,channels,expires_at,created_at FROM keys WHERE hash=? AND enabled=1 AND revoked=0", hash(raw)).Scan(&k.ID, &k.Name, &k.Prefix, &k.Enabled, &k.AllChannels, &ch, &k.ExpiresAt, &k.CreatedAt)
|
||||
if e != nil {
|
||||
return k, e
|
||||
}
|
||||
if k.ExpiresAt != "" {
|
||||
t, e := time.Parse(time.RFC3339Nano, k.ExpiresAt)
|
||||
if e != nil || !time.Now().Before(t) {
|
||||
return apiKey{}, sql.ErrNoRows
|
||||
}
|
||||
}
|
||||
e = json.Unmarshal([]byte(ch), &k.Channels)
|
||||
return k, e
|
||||
}
|
||||
func (k apiKey) allows(id string) bool {
|
||||
if k.AllChannels {
|
||||
return true
|
||||
}
|
||||
for _, v := range k.Channels {
|
||||
if v == id {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
@@ -0,0 +1,136 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"sort"
|
||||
"strings"
|
||||
"time"
|
||||
)
|
||||
|
||||
const (
|
||||
discoveryTimeout = 30 * time.Second
|
||||
discoveryMaxBody = 1 << 20
|
||||
discoveryMaxPages = 20
|
||||
)
|
||||
|
||||
func (s *Server) discoverModels(w http.ResponseWriter, r *http.Request) {
|
||||
if r.Method != http.MethodPost {
|
||||
errJSON(w, 405, "method not allowed")
|
||||
return
|
||||
}
|
||||
var c channel
|
||||
if decode(w, r, &c) != nil {
|
||||
errJSON(w, 400, "invalid channel")
|
||||
return
|
||||
}
|
||||
u, e := url.Parse(c.BaseURL)
|
||||
if e != nil || (u.Scheme != "http" && u.Scheme != "https") || u.Hostname() == "" || u.User != nil || u.RawQuery != "" || u.ForceQuery || strings.Contains(c.BaseURL, "#") || (c.Protocol != "openai-chat" && c.Protocol != "openai-responses" && c.Protocol != "anthropic") {
|
||||
errJSON(w, 400, "invalid channel")
|
||||
return
|
||||
}
|
||||
if c.ID != "" {
|
||||
var old string
|
||||
e = s.db.QueryRowContext(r.Context(), "SELECT api_key FROM channels WHERE id=?", c.ID).Scan(&old)
|
||||
if e == sql.ErrNoRows {
|
||||
errJSON(w, 404, "not found")
|
||||
return
|
||||
}
|
||||
if e != nil {
|
||||
errJSON(w, 500, "database error")
|
||||
return
|
||||
}
|
||||
if c.APIKey == "" {
|
||||
c.APIKey = old
|
||||
}
|
||||
}
|
||||
ctx, cancel := context.WithTimeout(r.Context(), discoveryTimeout)
|
||||
defer cancel()
|
||||
models, e := s.fetchModels(ctx, c)
|
||||
if e != nil {
|
||||
if errors.Is(ctx.Err(), context.DeadlineExceeded) {
|
||||
errJSON(w, 504, "model discovery timed out")
|
||||
} else {
|
||||
errJSON(w, 502, "model discovery failed")
|
||||
}
|
||||
return
|
||||
}
|
||||
writeJSON(w, map[string][]string{"models": models})
|
||||
}
|
||||
|
||||
func (s *Server) fetchModels(ctx context.Context, c channel) ([]string, error) {
|
||||
// Share the bounded transport, but never forward credentials across redirects.
|
||||
client := &http.Client{Transport: s.client.Transport, CheckRedirect: func(*http.Request, []*http.Request) error { return http.ErrUseLastResponse }}
|
||||
endpoint := strings.TrimRight(c.BaseURL, "/") + "/models"
|
||||
if c.Protocol == "anthropic" {
|
||||
endpoint = strings.TrimRight(c.BaseURL, "/") + "/v1/models"
|
||||
}
|
||||
seen := map[string]bool{}
|
||||
cursors := map[string]bool{}
|
||||
after := ""
|
||||
for page := 0; page < discoveryMaxPages; page++ {
|
||||
u := endpoint
|
||||
if after != "" {
|
||||
u += "?" + url.Values{"after_id": {after}}.Encode()
|
||||
}
|
||||
req, e := http.NewRequestWithContext(ctx, http.MethodGet, u, nil)
|
||||
if e != nil {
|
||||
return nil, e
|
||||
}
|
||||
req.Header.Set("Accept", "application/json")
|
||||
if c.Protocol == "anthropic" {
|
||||
req.Header.Set("x-api-key", c.APIKey)
|
||||
req.Header.Set("anthropic-version", "2023-06-01")
|
||||
} else {
|
||||
req.Header.Set("Authorization", "Bearer "+c.APIKey)
|
||||
}
|
||||
res, e := client.Do(req)
|
||||
if e != nil {
|
||||
return nil, e
|
||||
}
|
||||
if res.StatusCode != http.StatusOK {
|
||||
res.Body.Close()
|
||||
return nil, errors.New("upstream status")
|
||||
}
|
||||
body, e := io.ReadAll(io.LimitReader(res.Body, discoveryMaxBody+1))
|
||||
res.Body.Close()
|
||||
if e != nil || len(body) > discoveryMaxBody {
|
||||
return nil, errors.New("upstream body")
|
||||
}
|
||||
var result struct {
|
||||
Data []struct {
|
||||
ID string `json:"id"`
|
||||
} `json:"data"`
|
||||
HasMore bool `json:"has_more"`
|
||||
LastID string `json:"last_id"`
|
||||
}
|
||||
if json.Unmarshal(body, &result) != nil || result.Data == nil {
|
||||
return nil, errors.New("upstream models")
|
||||
}
|
||||
for _, model := range result.Data {
|
||||
if strings.TrimSpace(model.ID) == "" {
|
||||
return nil, errors.New("upstream model id")
|
||||
}
|
||||
seen[model.ID] = true
|
||||
}
|
||||
if c.Protocol != "anthropic" || !result.HasMore {
|
||||
models := make([]string, 0, len(seen))
|
||||
for id := range seen {
|
||||
models = append(models, id)
|
||||
}
|
||||
sort.Strings(models)
|
||||
return models, nil
|
||||
}
|
||||
if result.LastID == "" || cursors[result.LastID] {
|
||||
return nil, errors.New("upstream pagination")
|
||||
}
|
||||
cursors[result.LastID] = true
|
||||
after = result.LastID
|
||||
}
|
||||
return nil, errors.New("upstream page limit")
|
||||
}
|
||||
@@ -0,0 +1,273 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"reflect"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
const discoveryPath = "/admin/api/channels/discover-models"
|
||||
|
||||
func discover(t *testing.T, s *Server, cookie *http.Cookie, c channel) *httptest.ResponseRecorder {
|
||||
t.Helper()
|
||||
b, e := json.Marshal(c)
|
||||
if e != nil {
|
||||
t.Fatal(e)
|
||||
}
|
||||
return call(s, "POST", discoveryPath, string(b), cookie)
|
||||
}
|
||||
|
||||
func TestDiscoveryProtocolsAndReadOnly(t *testing.T) {
|
||||
for _, proto := range []string{"openai-chat", "openai-responses", "anthropic"} {
|
||||
t.Run(proto, func(t *testing.T) {
|
||||
s := testServer(t)
|
||||
cookie := loginCookie(t, s)
|
||||
calls := 0
|
||||
up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
calls++
|
||||
if r.Method != "GET" || r.URL.Path != "/v1/models" {
|
||||
t.Error(r.Method, r.URL)
|
||||
}
|
||||
if proto == "anthropic" {
|
||||
if r.Header.Get("x-api-key") != "secret" || r.Header.Get("anthropic-version") != "2023-06-01" || r.Header.Get("Authorization") != "" {
|
||||
t.Error("wrong anthropic headers")
|
||||
}
|
||||
if r.URL.Query().Get("after_id") == "" {
|
||||
fmt.Fprint(w, `{"data":[{"id":"z"}],"has_more":true,"last_id":"cursor/&?"}`)
|
||||
return
|
||||
}
|
||||
if r.URL.Query().Get("after_id") != "cursor/&?" {
|
||||
t.Error(r.URL)
|
||||
}
|
||||
} else if r.Header.Get("Authorization") != "Bearer secret" || r.Header.Get("x-api-key") != "" {
|
||||
t.Error("wrong openai headers")
|
||||
}
|
||||
fmt.Fprint(w, `{"data":[{"id":"z"},{"id":"a"},{"id":"a"}]}`)
|
||||
}))
|
||||
defer up.Close()
|
||||
draft := channel{BaseURL: up.URL + "/v1/", Protocol: proto, APIKey: "secret"}
|
||||
if proto == "anthropic" {
|
||||
draft.BaseURL = up.URL + "/"
|
||||
}
|
||||
w := discover(t, s, cookie, draft)
|
||||
if w.Code != 200 || strings.TrimSpace(w.Body.String()) != `{"models":["a","z"]}` {
|
||||
t.Fatal(w.Code, w.Body.String())
|
||||
}
|
||||
wantCalls := 1
|
||||
if proto == "anthropic" {
|
||||
wantCalls = 2
|
||||
}
|
||||
if calls != wantCalls {
|
||||
t.Fatal(calls)
|
||||
}
|
||||
cs, e := s.loadChannels()
|
||||
if e != nil || len(cs) != 0 {
|
||||
t.Fatal(cs, e)
|
||||
}
|
||||
draft.Name = "saved"
|
||||
draft.Models = map[string]string{"alias": "a"}
|
||||
b, _ := json.Marshal(draft)
|
||||
w = call(s, "POST", "/admin/api/channels", string(b), cookie)
|
||||
if w.Code != 200 {
|
||||
t.Fatal(w.Body.String())
|
||||
}
|
||||
var saved channel
|
||||
if e := json.Unmarshal(w.Body.Bytes(), &saved); e != nil {
|
||||
t.Fatal(e)
|
||||
}
|
||||
before, _ := s.loadChannels()
|
||||
draft.ID, draft.APIKey, draft.Name, draft.Models = saved.ID, "", "unsaved", nil
|
||||
if w = discover(t, s, cookie, draft); w.Code != 200 {
|
||||
t.Fatal(w.Code, w.Body.String())
|
||||
}
|
||||
after, _ := s.loadChannels()
|
||||
if !reflect.DeepEqual(before, after) {
|
||||
t.Fatal("discovery changed channel")
|
||||
}
|
||||
replacement := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if proto == "anthropic" {
|
||||
if r.Header.Get("x-api-key") != "replacement" {
|
||||
t.Error("draft key not used")
|
||||
}
|
||||
} else if r.Header.Get("Authorization") != "Bearer replacement" {
|
||||
t.Error("draft key not used")
|
||||
}
|
||||
fmt.Fprint(w, `{"data":[]}`)
|
||||
}))
|
||||
defer replacement.Close()
|
||||
draft.BaseURL, draft.APIKey = replacement.URL, "replacement"
|
||||
w = discover(t, s, cookie, draft)
|
||||
if w.Code != 200 || strings.TrimSpace(w.Body.String()) != `{"models":[]}` {
|
||||
t.Fatal(w.Code, w.Body.String())
|
||||
}
|
||||
after, _ = s.loadChannels()
|
||||
if !reflect.DeepEqual(before, after) {
|
||||
t.Fatal("draft replacement was saved")
|
||||
}
|
||||
var logs int
|
||||
if e := s.db.QueryRow("SELECT count(*) FROM logs").Scan(&logs); e != nil || logs != 0 {
|
||||
t.Fatal(logs, e)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestDiscoveryAuthAndValidation(t *testing.T) {
|
||||
s := testServer(t)
|
||||
cookie := loginCookie(t, s)
|
||||
if w := call(s, "POST", discoveryPath, "{}", nil); w.Code != 401 {
|
||||
t.Fatal(w.Code)
|
||||
}
|
||||
for _, header := range []string{"Origin", "X-Requested-With"} {
|
||||
r := httptest.NewRequest("POST", discoveryPath, strings.NewReader("{}"))
|
||||
r.AddCookie(cookie)
|
||||
r.Header.Set("Origin", "http://example.com")
|
||||
r.Header.Set("X-Requested-With", "ai-router")
|
||||
r.Header.Del(header)
|
||||
w := httptest.NewRecorder()
|
||||
s.ServeHTTP(w, r)
|
||||
if w.Code != 403 {
|
||||
t.Fatal(header, w.Code)
|
||||
}
|
||||
}
|
||||
for _, method := range []string{"GET", "DELETE", "PUT"} {
|
||||
if w := call(s, method, discoveryPath, "{}", cookie); w.Code != 405 {
|
||||
t.Fatal(method, w.Code)
|
||||
}
|
||||
}
|
||||
for _, base := range []string{"", "ftp://example.com", "http://", "http://user:[email protected]", "http://example.com?x=1", "http://example.com?", "http://example.com#"} {
|
||||
if w := discover(t, s, cookie, channel{BaseURL: base, Protocol: "openai-chat"}); w.Code != 400 {
|
||||
t.Fatal(base, w.Code)
|
||||
}
|
||||
}
|
||||
if w := discover(t, s, cookie, channel{BaseURL: "http://example.com", Protocol: "unknown"}); w.Code != 400 {
|
||||
t.Fatal(w.Code)
|
||||
}
|
||||
if w := discover(t, s, cookie, channel{ID: "missing", BaseURL: "http://example.com", Protocol: "anthropic"}); w.Code != 404 {
|
||||
t.Fatal(w.Code)
|
||||
}
|
||||
for _, body := range []string{"{", "{} {}", strings.Repeat(" ", maxBody+1)} {
|
||||
if w := call(s, "POST", discoveryPath, body, cookie); w.Code != 400 {
|
||||
t.Fatal(w.Code)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestDiscoveryUpstreamErrors(t *testing.T) {
|
||||
for _, tc := range []struct {
|
||||
name, body string
|
||||
status, calls int
|
||||
}{
|
||||
{"status", "secret upstream details", 401, 1},
|
||||
{"json", "secret", 200, 1},
|
||||
{"missing data", "{}", 200, 1},
|
||||
{"null", `{"data":null}`, 200, 1},
|
||||
{"empty id", `{"data":[{"id":""}]}`, 200, 1},
|
||||
{"trailing", `{"data":[]}{}`, 200, 1},
|
||||
{"oversized", strings.Repeat("x", discoveryMaxBody+1), 200, 1},
|
||||
{"missing cursor", `{"data":[],"has_more":true}`, 200, 1},
|
||||
{"repeated cursor", `{"data":[],"has_more":true,"last_id":"same"}`, 200, 2},
|
||||
{"page limit", "", 200, discoveryMaxPages},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
s := testServer(t)
|
||||
cookie := loginCookie(t, s)
|
||||
calls := 0
|
||||
up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
calls++
|
||||
w.WriteHeader(tc.status)
|
||||
if tc.name == "page limit" {
|
||||
fmt.Fprintf(w, `{"data":[],"has_more":true,"last_id":"%d"}`, calls)
|
||||
} else {
|
||||
fmt.Fprint(w, tc.body)
|
||||
}
|
||||
}))
|
||||
defer up.Close()
|
||||
w := discover(t, s, cookie, channel{BaseURL: up.URL, Protocol: "anthropic", APIKey: "secret"})
|
||||
if w.Code != 502 || strings.TrimSpace(w.Body.String()) != `{"error":"model discovery failed"}` || calls != tc.calls {
|
||||
t.Fatal(w.Code, w.Body.String(), calls)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestDiscoveryRedirectAndDeadline(t *testing.T) {
|
||||
s := testServer(t)
|
||||
cookie := loginCookie(t, s)
|
||||
target := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { t.Error("redirect followed") }))
|
||||
defer target.Close()
|
||||
up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { http.Redirect(w, r, target.URL, 307) }))
|
||||
defer up.Close()
|
||||
for _, proto := range []string{"openai-chat", "anthropic"} {
|
||||
if w := discover(t, s, cookie, channel{BaseURL: up.URL, Protocol: proto, APIKey: "secret"}); w.Code != 502 {
|
||||
t.Fatal(w.Code)
|
||||
}
|
||||
}
|
||||
slow := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { <-r.Context().Done() }))
|
||||
defer slow.Close()
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 20*time.Millisecond)
|
||||
defer cancel()
|
||||
b, _ := json.Marshal(channel{BaseURL: slow.URL, Protocol: "openai-chat"})
|
||||
r := httptest.NewRequest("POST", discoveryPath, strings.NewReader(string(b))).WithContext(ctx)
|
||||
r.AddCookie(cookie)
|
||||
r.Header.Set("Origin", "http://example.com")
|
||||
r.Header.Set("X-Requested-With", "ai-router")
|
||||
w := httptest.NewRecorder()
|
||||
s.ServeHTTP(w, r)
|
||||
if w.Code != 504 {
|
||||
t.Fatal(w.Code, w.Body.String())
|
||||
}
|
||||
}
|
||||
|
||||
func TestModelsListSelectedAliasesAndPermissions(t *testing.T) {
|
||||
s := testServer(t)
|
||||
cookie := loginCookie(t, s)
|
||||
var allowed string
|
||||
for _, name := range []string{"selected", "forbidden", "disabled"} {
|
||||
b, _ := json.Marshal(channel{Name: name, BaseURL: "http://example.com", Protocol: "openai-chat", Enabled: name != "disabled", Models: map[string]string{"z": "vendor/real", "alias": "vendor/real"}})
|
||||
w := call(s, "POST", "/admin/api/channels", string(b), cookie)
|
||||
if w.Code != 200 {
|
||||
t.Fatal(w.Body.String())
|
||||
}
|
||||
var c channel
|
||||
json.Unmarshal(w.Body.Bytes(), &c)
|
||||
if name != "forbidden" {
|
||||
if allowed != "" {
|
||||
allowed += ","
|
||||
}
|
||||
allowed += fmt.Sprintf("%q", c.ID)
|
||||
}
|
||||
}
|
||||
w := call(s, "POST", "/admin/api/keys", `{"name":"client","channels":[`+allowed+`]}`, cookie)
|
||||
if w.Code != 200 {
|
||||
t.Fatal(w.Body.String())
|
||||
}
|
||||
var key map[string]string
|
||||
json.Unmarshal(w.Body.Bytes(), &key)
|
||||
r := httptest.NewRequest("GET", "/v1/models", nil)
|
||||
r.Header.Set("Authorization", "Bearer "+key["key"])
|
||||
w = httptest.NewRecorder()
|
||||
s.ServeHTTP(w, r)
|
||||
var result struct{ Data []struct{ ID string } }
|
||||
if e := json.Unmarshal(w.Body.Bytes(), &result); e != nil || w.Code != 200 {
|
||||
t.Fatal(w.Code, e)
|
||||
}
|
||||
if len(result.Data) != 2 || result.Data[0].ID != "selected/alias" || result.Data[1].ID != "selected/z" {
|
||||
t.Fatal(w.Body.String())
|
||||
}
|
||||
for _, model := range []string{"forbidden/alias", "disabled/alias"} {
|
||||
r = httptest.NewRequest("POST", "/v1/chat/completions", strings.NewReader(fmt.Sprintf(`{"model":%q,"messages":[]}`, model)))
|
||||
r.Header.Set("Authorization", "Bearer "+key["key"])
|
||||
w = httptest.NewRecorder()
|
||||
s.ServeHTTP(w, r)
|
||||
if w.Code != 403 {
|
||||
t.Fatal(model, w.Code)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,443 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
type integrationKey struct{ raw, id string }
|
||||
|
||||
type integrationStreamRecorder struct {
|
||||
*httptest.ResponseRecorder
|
||||
flushed chan struct{}
|
||||
}
|
||||
|
||||
func (w *integrationStreamRecorder) Flush() {
|
||||
w.ResponseRecorder.Flush()
|
||||
select {
|
||||
case w.flushed <- struct{}{}:
|
||||
default:
|
||||
}
|
||||
}
|
||||
|
||||
func integrationAdmin(t *testing.T, s *Server, cookie *http.Cookie, method, path string, body any, out any) {
|
||||
t.Helper()
|
||||
b, err := json.Marshal(body)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
w := call(s, method, path, string(b), cookie)
|
||||
if w.Code != http.StatusOK {
|
||||
t.Fatalf("%s %s: %d %s", method, path, w.Code, w.Body.String())
|
||||
}
|
||||
if out != nil {
|
||||
if err := json.Unmarshal(w.Body.Bytes(), out); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func integrationChannel(t *testing.T, s *Server, cookie *http.Cookie, name, kind, base string) channel {
|
||||
t.Helper()
|
||||
if kind != "anthropic" {
|
||||
base += "/v1"
|
||||
}
|
||||
var ch channel
|
||||
integrationAdmin(t, s, cookie, "POST", "/admin/api/channels", channel{
|
||||
Name: name, Protocol: kind, BaseURL: base, Enabled: true,
|
||||
APIKey: "upstream-secret", Models: map[string]string{"gpt-5": "gpt-4", "gpt-4": "abc/gpt-4"},
|
||||
}, &ch)
|
||||
if ch.ID == "" || ch.APIKey != "" {
|
||||
t.Fatalf("invalid channel response: %+v", ch)
|
||||
}
|
||||
return ch
|
||||
}
|
||||
|
||||
func integrationCreateKey(t *testing.T, s *Server, cookie *http.Cookie, name string, ch channel) integrationKey {
|
||||
t.Helper()
|
||||
var created struct {
|
||||
Key string `json:"key"`
|
||||
}
|
||||
integrationAdmin(t, s, cookie, "POST", "/admin/api/keys", apiKey{Name: name, Channels: []string{ch.ID}}, &created)
|
||||
var keys []apiKey
|
||||
integrationAdmin(t, s, cookie, "GET", "/admin/api/keys", nil, &keys)
|
||||
for _, k := range keys {
|
||||
if k.Name == name && created.Key != "" {
|
||||
return integrationKey{created.Key, k.ID}
|
||||
}
|
||||
}
|
||||
t.Fatal("created key missing from management API")
|
||||
return integrationKey{}
|
||||
}
|
||||
|
||||
func integrationRequest(kind, model string, stream bool) string {
|
||||
if kind == "openai-responses" {
|
||||
return fmt.Sprintf(`{"model":%q,"input":"hello","max_output_tokens":32,"stream":%t}`, model, stream)
|
||||
}
|
||||
return fmt.Sprintf(`{"model":%q,"messages":[{"role":"user","content":"hello"}],"max_tokens":32,"stream":%t}`, model, stream)
|
||||
}
|
||||
|
||||
func integrationPath(kind string) string {
|
||||
return map[string]string{"openai-chat": "/v1/chat/completions", "openai-responses": "/v1/responses", "anthropic": "/v1/messages"}[kind]
|
||||
}
|
||||
|
||||
func integrationProxy(s *Server, key integrationKey, kind, model string, stream bool) *httptest.ResponseRecorder {
|
||||
r := httptest.NewRequest("POST", integrationPath(kind), strings.NewReader(integrationRequest(kind, model, stream)))
|
||||
if kind == "anthropic" {
|
||||
r.Header.Set("x-api-key", key.raw)
|
||||
} else {
|
||||
r.Header.Set("Authorization", "Bearer "+key.raw)
|
||||
}
|
||||
w := httptest.NewRecorder()
|
||||
s.ServeHTTP(w, r)
|
||||
return w
|
||||
}
|
||||
|
||||
func integrationResponse(kind string) string {
|
||||
switch kind {
|
||||
case "openai-chat":
|
||||
return `{"id":"chat-test","object":"chat.completion","model":"gpt-4","choices":[{"index":0,"message":{"role":"assistant","content":"hello"},"finish_reason":"stop"}],"usage":{"prompt_tokens":12,"completion_tokens":3,"prompt_tokens_details":{"cached_tokens":4}}}`
|
||||
case "openai-responses":
|
||||
return `{"id":"resp-test","object":"response","model":"gpt-4","status":"completed","output":[{"type":"message","id":"msg-test","role":"assistant","status":"completed","content":[{"type":"output_text","text":"hello","annotations":[]}]}],"usage":{"input_tokens":12,"output_tokens":3,"input_tokens_details":{"cached_tokens":4}}}`
|
||||
default:
|
||||
return `{"id":"msg-test","type":"message","role":"assistant","model":"gpt-4","content":[{"type":"text","text":"hello"}],"stop_reason":"end_turn","stop_sequence":null,"usage":{"input_tokens":6,"output_tokens":3,"cache_read_input_tokens":4,"cache_creation_input_tokens":2}}`
|
||||
}
|
||||
}
|
||||
|
||||
func integrationSSE(kind string) string {
|
||||
switch kind {
|
||||
case "openai-chat":
|
||||
return "data: " + `{"id":"chat-test","choices":[{"index":0,"delta":{"role":"assistant","content":"hello"},"finish_reason":null}]}` + "\n\ndata: " + `{"choices":[{"index":0,"delta":{},"finish_reason":"stop"}],"usage":{"prompt_tokens":12,"completion_tokens":3,"prompt_tokens_details":{"cached_tokens":4}}}` + "\n\ndata: [DONE]\n\n"
|
||||
case "openai-responses":
|
||||
return "event: response.created\ndata: " + `{"type":"response.created","response":{"id":"resp-test","status":"in_progress"}}` + "\n\n" +
|
||||
"event: response.output_item.added\ndata: " + `{"type":"response.output_item.added","output_index":0,"item":{"type":"message","role":"assistant","content":[]}}` + "\n\n" +
|
||||
"event: response.output_text.delta\ndata: " + `{"type":"response.output_text.delta","output_index":0,"content_index":0,"delta":"hello"}` + "\n\n" +
|
||||
"event: response.completed\ndata: " + `{"type":"response.completed","response":` + integrationResponse(kind) + "}\n\n"
|
||||
default:
|
||||
return "event: message_start\ndata: " + `{"type":"message_start","message":{"id":"msg-test","role":"assistant","content":[],"usage":{"input_tokens":6,"output_tokens":0,"cache_read_input_tokens":4,"cache_creation_input_tokens":2}}}` + "\n\n" +
|
||||
"event: content_block_start\ndata: " + `{"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}` + "\n\n" +
|
||||
"event: content_block_delta\ndata: " + `{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"hello"}}` + "\n\n" +
|
||||
"event: content_block_stop\ndata: " + `{"type":"content_block_stop","index":0}` + "\n\n" +
|
||||
"event: message_delta\ndata: " + `{"type":"message_delta","delta":{"stop_reason":"end_turn","stop_sequence":null},"usage":{"output_tokens":3}}` + "\n\n" +
|
||||
"event: message_stop\ndata: " + `{"type":"message_stop"}` + "\n\n"
|
||||
}
|
||||
}
|
||||
|
||||
func integrationLogs(t *testing.T, s *Server, c *http.Cookie, key integrationKey) []requestLog {
|
||||
t.Helper()
|
||||
if err := s.flushLogs(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
var logs []requestLog
|
||||
integrationAdmin(t, s, c, "GET", "/admin/api/logs?key_id="+key.id, nil, &logs)
|
||||
for _, l := range logs {
|
||||
if l.KeyID != key.id {
|
||||
t.Fatalf("cross-key log leak: %+v", l)
|
||||
}
|
||||
}
|
||||
return logs
|
||||
}
|
||||
|
||||
func integrationCount(t *testing.T, name string, got *int64, want int64) {
|
||||
t.Helper()
|
||||
if got == nil || *got != want {
|
||||
t.Fatalf("%s: got %v, want %d", name, got, want)
|
||||
}
|
||||
}
|
||||
|
||||
func TestIntegrationProtocolMatrix(t *testing.T) {
|
||||
kinds := []string{"openai-chat", "openai-responses", "anthropic"}
|
||||
for _, upKind := range kinds {
|
||||
for _, downKind := range kinds {
|
||||
t.Run(downKind+"_to_"+upKind, func(t *testing.T) {
|
||||
s := testServer(t)
|
||||
cookie := loginCookie(t, s)
|
||||
var calls atomic.Int64
|
||||
var expectedModel atomic.Value
|
||||
up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
calls.Add(1)
|
||||
if r.URL.Path != integrationPath(upKind) {
|
||||
t.Errorf("upstream path: %s", r.URL.Path)
|
||||
}
|
||||
if upKind == "anthropic" {
|
||||
if r.Header.Get("x-api-key") != "upstream-secret" || r.Header.Get("anthropic-version") == "" {
|
||||
t.Error("missing Anthropic credentials/version")
|
||||
}
|
||||
} else if r.Header.Get("Authorization") != "Bearer upstream-secret" {
|
||||
t.Error("wrong upstream credentials")
|
||||
}
|
||||
var body struct {
|
||||
Model string `json:"model"`
|
||||
Stream bool `json:"stream"`
|
||||
}
|
||||
b, err := io.ReadAll(r.Body)
|
||||
if err != nil || json.Unmarshal(b, &body) != nil {
|
||||
t.Error("invalid upstream request")
|
||||
w.WriteHeader(400)
|
||||
return
|
||||
}
|
||||
if body.Model != expectedModel.Load() || !strings.Contains(string(b), "hello") {
|
||||
t.Errorf("upstream request: %s; want model %s", b, expectedModel.Load())
|
||||
}
|
||||
if body.Stream {
|
||||
w.Header().Set("Content-Type", "text/event-stream")
|
||||
for _, frame := range strings.Split(integrationSSE(upKind), "\n\n") {
|
||||
if frame != "" {
|
||||
fmt.Fprint(w, frame+"\n\n")
|
||||
w.(http.Flusher).Flush()
|
||||
}
|
||||
}
|
||||
} else {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
fmt.Fprint(w, integrationResponse(upKind))
|
||||
}
|
||||
}))
|
||||
defer up.Close()
|
||||
abc := integrationChannel(t, s, cookie, "abc", upKind, up.URL)
|
||||
other := integrationChannel(t, s, cookie, "other", upKind, up.URL)
|
||||
keys := []integrationKey{integrationCreateKey(t, s, cookie, "first", abc), integrationCreateKey(t, s, cookie, "second", other)}
|
||||
for i, key := range keys {
|
||||
name, forbidden := "abc", "other"
|
||||
if i == 1 {
|
||||
name, forbidden = forbidden, name
|
||||
}
|
||||
for _, stream := range []bool{false, true} {
|
||||
for _, alias := range []string{"gpt-5", "gpt-4"} {
|
||||
expectedModel.Store("gpt-4")
|
||||
if alias == "gpt-4" {
|
||||
expectedModel.Store("abc/gpt-4")
|
||||
}
|
||||
w := integrationProxy(s, key, downKind, name+"/"+alias, stream)
|
||||
if w.Code != 200 || !strings.Contains(w.Body.String(), "hello") {
|
||||
t.Fatalf("stream=%t model=%s: %d %s", stream, alias, w.Code, w.Body.String())
|
||||
}
|
||||
if stream {
|
||||
terminal := map[string]string{"openai-chat": "data: [DONE]", "openai-responses": "event: response.completed", "anthropic": "event: message_stop"}[downKind]
|
||||
if !w.Flushed || !strings.Contains(w.Header().Get("Content-Type"), "text/event-stream") || !strings.Contains(w.Body.String(), terminal) {
|
||||
t.Fatalf("incomplete SSE: %s", w.Body.String())
|
||||
}
|
||||
} else {
|
||||
var body map[string]any
|
||||
if err := json.Unmarshal(w.Body.Bytes(), &body); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
field := map[string]string{"openai-chat": "choices", "openai-responses": "output", "anthropic": "content"}[downKind]
|
||||
if body[field] == nil {
|
||||
t.Fatalf("wrong downstream shape: %s", w.Body.String())
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
before := calls.Load()
|
||||
if w := integrationProxy(s, key, downKind, forbidden+"/gpt-5", false); w.Code != 403 {
|
||||
t.Fatalf("permission: %d %s", w.Code, w.Body.String())
|
||||
}
|
||||
if calls.Load() != before {
|
||||
t.Fatal("forbidden request reached upstream")
|
||||
}
|
||||
}
|
||||
if calls.Load() != 8 {
|
||||
t.Fatalf("upstream calls: %d", calls.Load())
|
||||
}
|
||||
for _, key := range keys {
|
||||
logs := integrationLogs(t, s, cookie, key)
|
||||
if len(logs) != 5 {
|
||||
t.Fatalf("logs: %+v", logs)
|
||||
}
|
||||
for _, l := range logs {
|
||||
if l.Status == 403 {
|
||||
if l.UsageKnown || l.Input != nil {
|
||||
t.Fatal("denied request has usage")
|
||||
}
|
||||
continue
|
||||
}
|
||||
actual := "gpt-4"
|
||||
if strings.HasSuffix(l.RequestModel, "/gpt-4") {
|
||||
actual = "abc/gpt-4"
|
||||
}
|
||||
if l.Status != 200 || l.Error != "" || !l.UsageKnown || l.ActualModel != actual || l.Downstream != downKind || l.Upstream != upKind {
|
||||
t.Fatalf("wrong log: %+v", l)
|
||||
}
|
||||
integrationCount(t, "input", l.Input, 12)
|
||||
integrationCount(t, "output", l.Output, 3)
|
||||
integrationCount(t, "cache read", l.CacheRead, 4)
|
||||
write := int64(0)
|
||||
if upKind == "anthropic" {
|
||||
write = 2
|
||||
}
|
||||
integrationCount(t, "cache write", l.CacheWrite, write)
|
||||
}
|
||||
var stats struct {
|
||||
Summary summary `json:"summary"`
|
||||
ByKey []struct {
|
||||
KeyID string `json:"key_id"`
|
||||
} `json:"by_key"`
|
||||
}
|
||||
integrationAdmin(t, s, cookie, "GET", "/admin/api/stats?key_id="+key.id, nil, &stats)
|
||||
if stats.Summary.Requests != 5 || stats.Summary.Success != 4 || stats.Summary.Failed != 1 || stats.Summary.Unknown != 1 || len(stats.ByKey) != 1 || stats.ByKey[0].KeyID != key.id {
|
||||
t.Fatalf("wrong isolated stats: %+v", stats)
|
||||
}
|
||||
integrationCount(t, "total input", stats.Summary.Input, 48)
|
||||
integrationCount(t, "total output", stats.Summary.Output, 12)
|
||||
integrationCount(t, "total cache read", stats.Summary.CacheRead, 16)
|
||||
write := int64(0)
|
||||
if upKind == "anthropic" {
|
||||
write = 8
|
||||
}
|
||||
integrationCount(t, "total cache write", stats.Summary.CacheWrite, write)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestIntegrationStreamFailures(t *testing.T) {
|
||||
for _, tc := range []struct {
|
||||
name, body string
|
||||
status int
|
||||
}{
|
||||
{"upstream_http_error", `{"error":"private upstream detail"}`, 429},
|
||||
{"error_event", "event: error\ndata: {\"type\":\"error\",\"error\":{\"message\":\"private upstream detail\"}}\n\n", 200},
|
||||
{"malformed_event", "data: {invalid}\n\n", 200},
|
||||
{"truncated_stream", "data: {\"choices\":[{\"index\":0,\"delta\":{\"content\":\"partial\"},\"finish_reason\":null}]}\n\n", 200},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
s := testServer(t)
|
||||
cookie := loginCookie(t, s)
|
||||
up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("Content-Type", "text/event-stream")
|
||||
w.WriteHeader(tc.status)
|
||||
fmt.Fprint(w, tc.body)
|
||||
}))
|
||||
defer up.Close()
|
||||
ch := integrationChannel(t, s, cookie, "abc", "openai-chat", up.URL)
|
||||
key := integrationCreateKey(t, s, cookie, "errors", ch)
|
||||
w := integrationProxy(s, key, "openai-responses", "abc/gpt-5", true)
|
||||
if w.Code != tc.status {
|
||||
t.Fatalf("status: %d body: %s", w.Code, w.Body.String())
|
||||
}
|
||||
if strings.Contains(w.Body.String(), "response.completed") || strings.Contains(w.Body.String(), "private upstream detail") {
|
||||
t.Fatalf("false success or leaked upstream detail: %s", w.Body.String())
|
||||
}
|
||||
logs := integrationLogs(t, s, cookie, key)
|
||||
if len(logs) != 1 {
|
||||
t.Fatalf("logs: %+v", logs)
|
||||
}
|
||||
l := logs[0]
|
||||
if l.Status < 400 || l.Error == "" || l.UsageKnown || l.Input != nil || l.Output != nil {
|
||||
t.Fatalf("failure not recorded correctly: %+v", l)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestAnthropicHeaderForwarding(t *testing.T) {
|
||||
s := testServer(t)
|
||||
cookie := loginCookie(t, s)
|
||||
var mu sync.Mutex
|
||||
got := map[string]string{}
|
||||
up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
mu.Lock()
|
||||
got["version"], got["beta"] = r.Header.Get("anthropic-version"), r.Header.Get("anthropic-beta")
|
||||
mu.Unlock()
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
fmt.Fprint(w, integrationResponse("anthropic"))
|
||||
}))
|
||||
defer up.Close()
|
||||
ch := integrationChannel(t, s, cookie, "abc", "anthropic", up.URL)
|
||||
key := integrationCreateKey(t, s, cookie, "beta", ch)
|
||||
request := func(version, beta string) {
|
||||
r := httptest.NewRequest("POST", "/v1/messages", strings.NewReader(integrationRequest("anthropic", "abc/gpt-5", false)))
|
||||
r.Header.Set("x-api-key", key.raw)
|
||||
if version != "" {
|
||||
r.Header.Set("anthropic-version", version)
|
||||
}
|
||||
if beta != "" {
|
||||
r.Header.Set("anthropic-beta", beta)
|
||||
}
|
||||
w := httptest.NewRecorder()
|
||||
s.ServeHTTP(w, r)
|
||||
if w.Code != 200 {
|
||||
t.Fatalf("status %d: %s", w.Code, w.Body.String())
|
||||
}
|
||||
}
|
||||
request("2023-06-01", "context-1m-2025-08-07,interleaved-thinking-2025-05-14")
|
||||
mu.Lock()
|
||||
if got["version"] != "2023-06-01" || got["beta"] != "context-1m-2025-08-07,interleaved-thinking-2025-05-14" {
|
||||
t.Fatalf("valid headers dropped: %v", got)
|
||||
}
|
||||
mu.Unlock()
|
||||
request("bad version; injected", `weird"value`)
|
||||
mu.Lock()
|
||||
if got["version"] != "2023-06-01" || got["beta"] != "" {
|
||||
t.Fatalf("unsafe headers forwarded: %v", got)
|
||||
}
|
||||
mu.Unlock()
|
||||
}
|
||||
|
||||
func TestIntegrationCancellationClosesUpstream(t *testing.T) {
|
||||
s := testServer(t)
|
||||
cookie := loginCookie(t, s)
|
||||
started, stopped := make(chan struct{}), make(chan struct{})
|
||||
release := make(chan struct{})
|
||||
up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
io.Copy(io.Discard, r.Body)
|
||||
w.Header().Set("Content-Type", "text/event-stream")
|
||||
fmt.Fprint(w, "data: {\"choices\":[{\"index\":0,\"delta\":{\"content\":\"partial\"},\"finish_reason\":null}]}\n\n")
|
||||
w.(http.Flusher).Flush()
|
||||
close(started)
|
||||
select {
|
||||
case <-r.Context().Done():
|
||||
close(stopped)
|
||||
case <-release:
|
||||
}
|
||||
}))
|
||||
defer up.Close()
|
||||
defer close(release)
|
||||
ch := integrationChannel(t, s, cookie, "abc", "openai-chat", up.URL)
|
||||
key := integrationCreateKey(t, s, cookie, "cancel", ch)
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
defer cancel()
|
||||
r := httptest.NewRequest("POST", "/v1/responses", strings.NewReader(integrationRequest("openai-responses", "abc/gpt-5", true))).WithContext(ctx)
|
||||
r.Header.Set("Authorization", "Bearer "+key.raw)
|
||||
done := make(chan struct{})
|
||||
w := &integrationStreamRecorder{ResponseRecorder: httptest.NewRecorder(), flushed: make(chan struct{}, 1)}
|
||||
go func() { s.ServeHTTP(w, r); close(done) }()
|
||||
select {
|
||||
case <-started:
|
||||
case <-time.After(3 * time.Second):
|
||||
t.Fatal("upstream did not start")
|
||||
}
|
||||
select {
|
||||
case <-w.flushed:
|
||||
case <-time.After(3 * time.Second):
|
||||
t.Fatal("downstream did not receive streaming output")
|
||||
}
|
||||
cancel()
|
||||
select {
|
||||
case <-done:
|
||||
case <-time.After(3 * time.Second):
|
||||
t.Fatal("proxy did not stop on cancellation")
|
||||
}
|
||||
select {
|
||||
case <-stopped:
|
||||
case <-time.After(3 * time.Second):
|
||||
t.Fatal("upstream request was not cancelled")
|
||||
}
|
||||
logs := integrationLogs(t, s, cookie, key)
|
||||
if len(logs) != 1 || logs[0].Status < 400 || logs[0].Error == "" || logs[0].UsageKnown || logs[0].Input != nil {
|
||||
t.Fatalf("cancel log: %+v", logs)
|
||||
}
|
||||
if strings.Contains(w.Body.String(), "response.completed") {
|
||||
t.Fatal("cancelled response marked completed")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,158 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"log"
|
||||
"strings"
|
||||
"time"
|
||||
"unicode/utf8"
|
||||
)
|
||||
|
||||
const (
|
||||
maxLogFieldBytes = 512
|
||||
maxLogErrorBytes = 4096
|
||||
maxLogUsageBytes = 16384
|
||||
logCleanupBatch = 1000
|
||||
)
|
||||
|
||||
func (s *Server) SetLogRetention(days int) {
|
||||
if days < 0 {
|
||||
days = 0
|
||||
}
|
||||
s.logRetentionDays.Store(int64(days))
|
||||
}
|
||||
|
||||
func truncateLogField(value string, limit int) string {
|
||||
if len(value) <= limit {
|
||||
return value
|
||||
}
|
||||
for limit > 0 && !utf8.RuneStart(value[limit]) {
|
||||
limit--
|
||||
}
|
||||
return strings.Clone(value[:limit])
|
||||
}
|
||||
|
||||
func boundedLog(l requestLog) requestLog {
|
||||
for _, field := range []*string{&l.CreatedAt, &l.KeyID, &l.KeyName, &l.Channel, &l.RequestModel, &l.ActualModel, &l.Downstream, &l.Upstream} {
|
||||
*field = truncateLogField(*field, maxLogFieldBytes)
|
||||
}
|
||||
l.Error = truncateLogField(l.Error, maxLogErrorBytes)
|
||||
if len(l.Usage.Raw) > maxLogUsageBytes {
|
||||
// Keep valid JSON without retaining an attacker-controlled oversized object.
|
||||
l.Usage.Raw = []byte(`{"truncated":true}`)
|
||||
}
|
||||
l.RawUsage = nil
|
||||
return l
|
||||
}
|
||||
|
||||
func (s *Server) purgeExpiredLogs(ctx context.Context, days int) error {
|
||||
if days <= 0 {
|
||||
return nil
|
||||
}
|
||||
cutoff := time.Now().UTC().AddDate(0, 0, -days).Truncate(time.Second)
|
||||
_, err := s.db.ExecContext(ctx, "DELETE FROM logs WHERE id IN (SELECT id FROM logs WHERE created_at<? AND julianday(created_at)<julianday(?) LIMIT ?)", cutoff.Format("2006-01-02T15:04:05"), cutoff.Format(time.RFC3339), logCleanupBatch)
|
||||
return err
|
||||
}
|
||||
|
||||
type logEntry struct {
|
||||
value requestLog
|
||||
barrier chan error
|
||||
}
|
||||
|
||||
func (s *Server) record(l requestLog) {
|
||||
s.logMu.RLock()
|
||||
defer s.logMu.RUnlock()
|
||||
if s.logsClosed {
|
||||
return
|
||||
}
|
||||
l = boundedLog(l)
|
||||
l.Usage.Raw = append([]byte(nil), l.Usage.Raw...)
|
||||
s.logQueue <- logEntry{value: l}
|
||||
}
|
||||
|
||||
func (s *Server) flushLogs() error {
|
||||
s.logMu.RLock()
|
||||
if s.logsClosed {
|
||||
s.logMu.RUnlock()
|
||||
<-s.logDone
|
||||
return s.logErr
|
||||
}
|
||||
done := make(chan error, 1)
|
||||
s.logQueue <- logEntry{barrier: done}
|
||||
s.logMu.RUnlock()
|
||||
return <-done
|
||||
}
|
||||
|
||||
func (s *Server) writeLogs() {
|
||||
defer close(s.logDone)
|
||||
ticker := time.NewTicker(100 * time.Millisecond)
|
||||
defer ticker.Stop()
|
||||
cleanup := time.NewTicker(time.Minute)
|
||||
defer cleanup.Stop()
|
||||
batch := make([]requestLog, 0, 100)
|
||||
flush := func() {
|
||||
if len(batch) == 0 {
|
||||
return
|
||||
}
|
||||
for attempt := 0; ; attempt++ {
|
||||
if s.persistLogs(batch) == nil {
|
||||
break
|
||||
}
|
||||
if attempt == 3 {
|
||||
s.logErr = errors.New("request log persistence failed after 3 retries")
|
||||
log.Printf("ai-router: request log persistence failed; batch_size=%d retries=3", len(batch))
|
||||
break
|
||||
}
|
||||
log.Printf("ai-router: retrying request log persistence; batch_size=%d retry=%d", len(batch), attempt+1)
|
||||
time.Sleep(time.Duration(attempt+1) * 10 * time.Millisecond)
|
||||
}
|
||||
clear(batch)
|
||||
batch = batch[:0]
|
||||
}
|
||||
for {
|
||||
select {
|
||||
case entry, ok := <-s.logQueue:
|
||||
if !ok {
|
||||
flush()
|
||||
return
|
||||
}
|
||||
if entry.barrier != nil {
|
||||
flush()
|
||||
entry.barrier <- s.logErr
|
||||
continue
|
||||
}
|
||||
batch = append(batch, entry.value)
|
||||
if len(batch) == cap(batch) {
|
||||
flush()
|
||||
}
|
||||
case <-ticker.C:
|
||||
flush()
|
||||
case <-cleanup.C:
|
||||
if days := s.logRetentionDays.Load(); days > 0 {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), time.Second)
|
||||
err := s.purgeExpiredLogs(ctx, int(days))
|
||||
cancel()
|
||||
if err != nil {
|
||||
log.Print("ai-router: request log retention cleanup failed")
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (s *Server) persistLogs(batch []requestLog) error {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||
defer cancel()
|
||||
tx, err := s.db.BeginTx(ctx, nil)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer tx.Rollback()
|
||||
for _, l := range batch {
|
||||
if err := insertLog(ctx, tx, l); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return tx.Commit()
|
||||
}
|
||||
@@ -0,0 +1,248 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"ai-router/internal/protocol"
|
||||
"bytes"
|
||||
"context"
|
||||
"database/sql"
|
||||
"log"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func waitLogCondition(t *testing.T, condition func() bool) {
|
||||
t.Helper()
|
||||
deadline := time.Now().Add(3 * time.Second)
|
||||
for !condition() {
|
||||
if time.Now().After(deadline) {
|
||||
t.Fatal("condition timed out")
|
||||
}
|
||||
time.Sleep(time.Millisecond)
|
||||
}
|
||||
}
|
||||
|
||||
func logCount(t *testing.T, db *sql.DB) int {
|
||||
t.Helper()
|
||||
var n int
|
||||
if err := db.QueryRow("SELECT COUNT(*) FROM logs").Scan(&n); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
return n
|
||||
}
|
||||
|
||||
func TestLogAsyncAndRawCopy(t *testing.T) {
|
||||
s := testServer(t)
|
||||
conn, err := s.db.Conn(context.Background())
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer conn.Close()
|
||||
raw := []byte(`{"input_tokens":12}`)
|
||||
done := make(chan struct{})
|
||||
go func() {
|
||||
s.record(requestLog{Usage: protocol.Usage{Raw: raw}})
|
||||
close(done)
|
||||
}()
|
||||
select {
|
||||
case <-done:
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("record waited for database")
|
||||
}
|
||||
raw[0] = 'x'
|
||||
conn.Close()
|
||||
if err := s.flushLogs(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
var got string
|
||||
if err := s.db.QueryRow("SELECT raw_usage FROM logs").Scan(&got); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if got != `{"input_tokens":12}` {
|
||||
t.Fatal(got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestLogTimerAndBatchDrain(t *testing.T) {
|
||||
s := testServer(t)
|
||||
s.record(requestLog{})
|
||||
waitLogCondition(t, func() bool { return logCount(t, s.db) == 1 })
|
||||
for i := 0; i < 257; i++ {
|
||||
s.record(requestLog{})
|
||||
}
|
||||
if err := s.flushLogs(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if n := logCount(t, s.db); n != 258 {
|
||||
t.Fatal(n)
|
||||
}
|
||||
}
|
||||
|
||||
func TestLogBackpressureAndConcurrentClose(t *testing.T) {
|
||||
path := filepath.Join(t.TempDir(), "logs.db")
|
||||
s, err := New(path, "password", nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() { s.Close() })
|
||||
conn, err := s.db.Conn(context.Background())
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer conn.Close()
|
||||
s.record(requestLog{})
|
||||
waitLogCondition(t, func() bool { return s.db.Stats().WaitCount > 0 })
|
||||
if cap(s.logQueue) != 1024 {
|
||||
t.Fatal(cap(s.logQueue))
|
||||
}
|
||||
for i := 0; i < 1024; i++ {
|
||||
s.record(requestLog{})
|
||||
}
|
||||
enqueued := make(chan struct{})
|
||||
go func() { s.record(requestLog{}); close(enqueued) }()
|
||||
select {
|
||||
case <-enqueued:
|
||||
t.Fatal("full queue did not apply backpressure")
|
||||
case <-time.After(20 * time.Millisecond):
|
||||
}
|
||||
var wg sync.WaitGroup
|
||||
for i := 0; i < 8; i++ {
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
if err := s.Close(); err != nil {
|
||||
t.Error(err)
|
||||
}
|
||||
}()
|
||||
}
|
||||
conn.Close()
|
||||
done := make(chan struct{})
|
||||
go func() { wg.Wait(); close(done) }()
|
||||
select {
|
||||
case <-done:
|
||||
case <-time.After(5 * time.Second):
|
||||
t.Fatal("Close deadlocked")
|
||||
}
|
||||
<-enqueued
|
||||
s.record(requestLog{})
|
||||
if err := s.flushLogs(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
db, err := sql.Open("sqlite", path)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer db.Close()
|
||||
if n := logCount(t, db); n != 1026 {
|
||||
t.Fatal(n)
|
||||
}
|
||||
}
|
||||
|
||||
func TestLogPersistenceErrorIsStickyAndRedacted(t *testing.T) {
|
||||
s := testServer(t)
|
||||
var output bytes.Buffer
|
||||
previous := log.Writer()
|
||||
log.SetOutput(&output)
|
||||
defer log.SetOutput(previous)
|
||||
if _, err := s.db.Exec(`CREATE TRIGGER reject_log BEFORE INSERT ON logs WHEN NEW.key_name='secret' BEGIN SELECT RAISE(FAIL,'private database secret'); END`); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
s.record(requestLog{KeyName: "ok"})
|
||||
s.record(requestLog{KeyName: "secret"})
|
||||
start := time.Now()
|
||||
err := s.flushLogs()
|
||||
if err == nil {
|
||||
t.Fatal("missing persistence error")
|
||||
}
|
||||
if time.Since(start) > time.Second {
|
||||
t.Fatal("unbounded retries")
|
||||
}
|
||||
if n := logCount(t, s.db); n != 0 {
|
||||
t.Fatalf("partial transaction: %d", n)
|
||||
}
|
||||
if strings.Contains(output.String(), "secret") || strings.Contains(err.Error(), "secret") {
|
||||
t.Fatal("error leaked metadata")
|
||||
}
|
||||
if n := strings.Count(output.String(), "retrying"); n != 3 {
|
||||
t.Fatalf("retries: %d", n)
|
||||
}
|
||||
if _, err := s.db.Exec("DROP TRIGGER reject_log"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
s.record(requestLog{})
|
||||
if err := s.flushLogs(); err == nil {
|
||||
t.Fatal("error was not sticky")
|
||||
}
|
||||
if n := logCount(t, s.db); n != 1 {
|
||||
t.Fatal(n)
|
||||
}
|
||||
if err := s.Close(); err == nil {
|
||||
t.Fatal("Close lost persistence error")
|
||||
}
|
||||
if err := s.Close(); err == nil {
|
||||
t.Fatal("second Close lost error")
|
||||
}
|
||||
}
|
||||
|
||||
func TestCloseWaitsForActiveRequest(t *testing.T) {
|
||||
path := filepath.Join(t.TempDir(), "logs.db")
|
||||
s, err := New(path, "test-password", nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() { s.Close() })
|
||||
started, release := make(chan struct{}), make(chan struct{})
|
||||
up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
close(started)
|
||||
<-release
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
w.Write([]byte(integrationResponse("openai-chat")))
|
||||
}))
|
||||
defer up.Close()
|
||||
defer func() {
|
||||
select {
|
||||
case <-release:
|
||||
default:
|
||||
close(release)
|
||||
}
|
||||
}()
|
||||
cookie := loginCookie(t, s)
|
||||
ch := integrationChannel(t, s, cookie, "abc", "openai-chat", up.URL)
|
||||
key := integrationCreateKey(t, s, cookie, "active", ch)
|
||||
requestDone := make(chan struct{})
|
||||
go func() { integrationProxy(s, key, "openai-chat", "abc/gpt-5", false); close(requestDone) }()
|
||||
<-started
|
||||
closed := make(chan error, 1)
|
||||
go func() { closed <- s.Close() }()
|
||||
waitLogCondition(t, func() bool { s.lifecycleMu.Lock(); defer s.lifecycleMu.Unlock(); return s.closing })
|
||||
if w := call(s, "GET", "/v1/models", "", nil); w.Code != 503 {
|
||||
t.Fatal(w.Code)
|
||||
}
|
||||
select {
|
||||
case <-closed:
|
||||
t.Fatal("Close skipped active request")
|
||||
default:
|
||||
}
|
||||
close(release)
|
||||
<-requestDone
|
||||
select {
|
||||
case err := <-closed:
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
case <-time.After(3 * time.Second):
|
||||
t.Fatal("Close did not drain")
|
||||
}
|
||||
db, err := sql.Open("sqlite", path)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer db.Close()
|
||||
if n := logCount(t, db); n != 1 {
|
||||
t.Fatal(n)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,229 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"ai-router/internal/protocol"
|
||||
"bytes"
|
||||
"context"
|
||||
"database/sql"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"io"
|
||||
"net/http"
|
||||
"strings"
|
||||
"time"
|
||||
"unicode/utf8"
|
||||
)
|
||||
|
||||
const modelTestTimeout = 30 * time.Second
|
||||
const modelTestMaxBody = 1 << 20
|
||||
|
||||
type modelTestResult struct {
|
||||
OK bool `json:"ok"`
|
||||
DurationMS int64 `json:"duration_ms"`
|
||||
Status int `json:"status"`
|
||||
Error string `json:"error,omitempty"`
|
||||
}
|
||||
|
||||
// testModel is an administrative diagnostic, excluded from call logs and usage
|
||||
// statistics. It creates no API key; upstream providers may still bill the probe.
|
||||
func (s *Server) testModel(w http.ResponseWriter, r *http.Request) {
|
||||
if r.Method != http.MethodPost {
|
||||
errJSON(w, 405, "method not allowed")
|
||||
return
|
||||
}
|
||||
var input struct {
|
||||
Model string `json:"model"`
|
||||
Prompt *string `json:"prompt"`
|
||||
}
|
||||
if decode(w, r, &input) != nil {
|
||||
errJSON(w, 400, "invalid request")
|
||||
return
|
||||
}
|
||||
name, alias, ok := strings.Cut(input.Model, "/")
|
||||
prompt := "Reply with OK."
|
||||
if input.Prompt != nil {
|
||||
prompt = *input.Prompt
|
||||
}
|
||||
if !ok || name == "" || alias == "" || strings.TrimSpace(prompt) == "" || utf8.RuneCountInString(prompt) > 2000 {
|
||||
errJSON(w, 400, "invalid model or prompt")
|
||||
return
|
||||
}
|
||||
ctx, cancel := context.WithTimeout(r.Context(), modelTestTimeout)
|
||||
defer cancel()
|
||||
var c channel
|
||||
var models string
|
||||
e := s.db.QueryRowContext(ctx, "SELECT base_url,protocol,enabled,models,api_key FROM channels WHERE name=?", name).Scan(&c.BaseURL, &c.Protocol, &c.Enabled, &models, &c.APIKey)
|
||||
if e == sql.ErrNoRows {
|
||||
errJSON(w, 403, "model or channel unavailable")
|
||||
return
|
||||
}
|
||||
if e != nil || json.Unmarshal([]byte(models), &c.Models) != nil {
|
||||
errJSON(w, 500, "database error")
|
||||
return
|
||||
}
|
||||
actual, mapped := c.Models[alias]
|
||||
if !c.Enabled || !mapped || actual == "" {
|
||||
errJSON(w, 403, "model or channel unavailable")
|
||||
return
|
||||
}
|
||||
writeJSON(w, s.probeModel(ctx, c, actual, prompt))
|
||||
}
|
||||
|
||||
func (s *Server) probeModel(ctx context.Context, c channel, actual, prompt string) (result modelTestResult) {
|
||||
start := time.Now()
|
||||
defer func() {
|
||||
result.DurationMS = time.Since(start).Milliseconds()
|
||||
if ctx.Err() != nil {
|
||||
result.OK = false
|
||||
result.Error = "model test canceled"
|
||||
if ctx.Err() == context.DeadlineExceeded {
|
||||
result.Error = "model test timed out"
|
||||
}
|
||||
}
|
||||
}()
|
||||
result.Error = "upstream request failed"
|
||||
up := protocol.Kind(c.Protocol)
|
||||
body, _ := json.Marshal(map[string]any{"messages": []any{map[string]string{"role": "user", "content": prompt}}, "stream": false, "max_completion_tokens": 128})
|
||||
payload, e := protocol.ConvertRequest(protocol.Chat, up, body, actual)
|
||||
if e != nil {
|
||||
result.Error = "invalid upstream configuration"
|
||||
return
|
||||
}
|
||||
suffix := map[protocol.Kind]string{protocol.Chat: "/chat/completions", protocol.Responses: "/responses", protocol.Anthropic: "/v1/messages"}[up]
|
||||
req, e := http.NewRequestWithContext(ctx, http.MethodPost, c.BaseURL+suffix, bytes.NewReader(payload))
|
||||
if e != nil {
|
||||
result.Error = "invalid upstream configuration"
|
||||
return
|
||||
}
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
if up == protocol.Anthropic {
|
||||
req.Header.Set("x-api-key", c.APIKey)
|
||||
req.Header.Set("anthropic-version", "2023-06-01")
|
||||
} else {
|
||||
req.Header.Set("Authorization", "Bearer "+c.APIKey)
|
||||
}
|
||||
client := &http.Client{Transport: s.client.Transport, CheckRedirect: func(*http.Request, []*http.Request) error { return http.ErrUseLastResponse }}
|
||||
res, e := client.Do(req)
|
||||
if e != nil {
|
||||
if errors.Is(e, context.DeadlineExceeded) {
|
||||
result.Error = "model test timed out"
|
||||
}
|
||||
return
|
||||
}
|
||||
defer res.Body.Close()
|
||||
result.Status = res.StatusCode
|
||||
if res.StatusCode < 200 || res.StatusCode >= 300 {
|
||||
result.Error = "upstream rejected request"
|
||||
return
|
||||
}
|
||||
body, e = io.ReadAll(io.LimitReader(res.Body, modelTestMaxBody+1))
|
||||
if len(body) > modelTestMaxBody {
|
||||
result.Error = "upstream response too large"
|
||||
return
|
||||
}
|
||||
if e != nil || !validModelTestResponse(up, body) {
|
||||
result.Error = "invalid upstream response"
|
||||
return
|
||||
}
|
||||
result.OK, result.Error = true, ""
|
||||
return
|
||||
}
|
||||
|
||||
func validModelTestResponse(kind protocol.Kind, body []byte) bool {
|
||||
var r struct {
|
||||
ID string `json:"id"`
|
||||
Object string `json:"object"`
|
||||
Type string `json:"type"`
|
||||
Role string `json:"role"`
|
||||
Error json.RawMessage `json:"error"`
|
||||
Status string `json:"status"`
|
||||
StopReason string `json:"stop_reason"`
|
||||
Incomplete struct {
|
||||
Reason string `json:"reason"`
|
||||
} `json:"incomplete_details"`
|
||||
Output []json.RawMessage `json:"output"`
|
||||
Content []json.RawMessage `json:"content"`
|
||||
Choices []struct {
|
||||
Finish string `json:"finish_reason"`
|
||||
Message struct {
|
||||
Role string `json:"role"`
|
||||
Content *string `json:"content"`
|
||||
} `json:"message"`
|
||||
} `json:"choices"`
|
||||
}
|
||||
if json.Unmarshal(body, &r) != nil || r.ID == "" || (len(r.Error) != 0 && string(r.Error) != "null") {
|
||||
return false
|
||||
}
|
||||
switch kind {
|
||||
case protocol.Chat:
|
||||
if r.Object != "chat.completion" || len(r.Choices) == 0 {
|
||||
return false
|
||||
}
|
||||
for _, c := range r.Choices {
|
||||
if c.Message.Role != "assistant" || (c.Finish != "stop" && c.Finish != "length") || (c.Message.Content == nil && c.Finish != "length") {
|
||||
return false
|
||||
}
|
||||
}
|
||||
return true
|
||||
case protocol.Responses:
|
||||
if r.Object != "response" || r.Output == nil {
|
||||
return false
|
||||
}
|
||||
exhausted := r.Status == "incomplete" && r.Incomplete.Reason == "max_output_tokens"
|
||||
if !exhausted && (r.Status != "completed" || len(r.Output) == 0) {
|
||||
return false
|
||||
}
|
||||
for _, raw := range r.Output {
|
||||
var item struct {
|
||||
Type string `json:"type"`
|
||||
Role string `json:"role"`
|
||||
Content []json.RawMessage `json:"content"`
|
||||
Summary []json.RawMessage `json:"summary"`
|
||||
}
|
||||
if json.Unmarshal(raw, &item) != nil {
|
||||
return false
|
||||
}
|
||||
if item.Type == "reasoning" && item.Summary != nil {
|
||||
continue
|
||||
}
|
||||
if item.Type != "message" || item.Role != "assistant" || !validProbeBlocks(item.Content, "output_text") {
|
||||
return false
|
||||
}
|
||||
}
|
||||
return true
|
||||
case protocol.Anthropic:
|
||||
if r.Type != "message" || r.Role != "assistant" || r.Content == nil {
|
||||
return false
|
||||
}
|
||||
if r.StopReason == "max_tokens" && len(r.Content) == 0 {
|
||||
return true
|
||||
}
|
||||
return (r.StopReason == "end_turn" || r.StopReason == "max_tokens") && validProbeBlocks(r.Content, "text")
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func validProbeBlocks(blocks []json.RawMessage, typ string) bool {
|
||||
if len(blocks) == 0 {
|
||||
return false
|
||||
}
|
||||
for _, raw := range blocks {
|
||||
var b struct {
|
||||
Type string `json:"type"`
|
||||
Text *string `json:"text"`
|
||||
Thinking *string `json:"thinking"`
|
||||
Signature *string `json:"signature"`
|
||||
Data *string `json:"data"`
|
||||
}
|
||||
if json.Unmarshal(raw, &b) != nil {
|
||||
return false
|
||||
}
|
||||
if typ == "text" && ((b.Type == "thinking" && b.Thinking != nil && b.Signature != nil) || (b.Type == "redacted_thinking" && b.Data != nil)) {
|
||||
continue
|
||||
}
|
||||
if b.Type != typ || b.Text == nil {
|
||||
return false
|
||||
}
|
||||
}
|
||||
return true
|
||||
}
|
||||
@@ -0,0 +1,267 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"ai-router/internal/protocol"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
const probePath = "/admin/api/models/test"
|
||||
const probeChat = `{"id":"c","object":"chat.completion","choices":[{"message":{"role":"assistant","content":"OK"},"finish_reason":"stop"}]}`
|
||||
|
||||
func saveProbeChannel(t *testing.T, s *Server, cookie *http.Cookie, url, kind string, enabled bool) {
|
||||
t.Helper()
|
||||
b, _ := json.Marshal(channel{Name: "saved", BaseURL: url, Protocol: kind, Enabled: enabled, Models: map[string]string{"alias": "actual/model"}, APIKey: "secret-key"})
|
||||
if w := call(s, "POST", "/admin/api/channels", string(b), cookie); w.Code != 200 {
|
||||
t.Fatal(w.Body.String())
|
||||
}
|
||||
}
|
||||
|
||||
func probeResult(t *testing.T, w *httptest.ResponseRecorder) modelTestResult {
|
||||
t.Helper()
|
||||
var result modelTestResult
|
||||
if w.Code != 200 || json.Unmarshal(w.Body.Bytes(), &result) != nil {
|
||||
t.Fatal(w.Code, w.Body.String())
|
||||
}
|
||||
if result.DurationMS < 0 || strings.Contains(w.Body.String(), "secret-key") {
|
||||
t.Fatal(w.Body.String())
|
||||
}
|
||||
return result
|
||||
}
|
||||
|
||||
func TestModelProbeProtocols(t *testing.T) {
|
||||
for _, tc := range []struct{ kind, path, limit, response string }{
|
||||
{"openai-chat", "/chat/completions", "max_completion_tokens", probeChat},
|
||||
{"openai-responses", "/responses", "max_output_tokens", `{"id":"r","object":"response","status":"completed","output":[{"type":"message","role":"assistant","content":[{"type":"output_text","text":"OK"}]}]}`},
|
||||
{"anthropic", "/v1/messages", "max_tokens", `{"id":"m","type":"message","role":"assistant","stop_reason":"end_turn","content":[{"type":"text","text":"OK"}]}`},
|
||||
} {
|
||||
t.Run(tc.kind, func(t *testing.T) {
|
||||
s := testServer(t)
|
||||
cookie := loginCookie(t, s)
|
||||
expectedPrompt := "Reply with OK."
|
||||
up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
b, _ := io.ReadAll(r.Body)
|
||||
var obj map[string]any
|
||||
if json.Unmarshal(b, &obj) != nil || obj["model"] != "actual/model" || obj[tc.limit] != float64(128) || obj["stream"] != false || !strings.Contains(string(b), expectedPrompt) {
|
||||
t.Error(string(b))
|
||||
}
|
||||
if tc.kind == "openai-chat" && obj["max_tokens"] != nil {
|
||||
t.Error("legacy token limit")
|
||||
}
|
||||
if r.Method != "POST" || r.URL.Path != tc.path {
|
||||
t.Error(r.Method, r.URL.Path)
|
||||
}
|
||||
if tc.kind == "anthropic" {
|
||||
if r.Header.Get("x-api-key") != "secret-key" || r.Header.Get("anthropic-version") != "2023-06-01" {
|
||||
t.Error("headers")
|
||||
}
|
||||
} else if r.Header.Get("Authorization") != "Bearer secret-key" {
|
||||
t.Error("authorization")
|
||||
}
|
||||
fmt.Fprint(w, tc.response)
|
||||
}))
|
||||
defer up.Close()
|
||||
saveProbeChannel(t, s, cookie, up.URL, tc.kind, true)
|
||||
for _, prompt := range []string{"", strings.Repeat("中", 2000)} {
|
||||
input := map[string]string{"model": "saved/alias"}
|
||||
if prompt != "" {
|
||||
expectedPrompt = prompt
|
||||
input["prompt"] = prompt
|
||||
}
|
||||
b, _ := json.Marshal(input)
|
||||
result := probeResult(t, call(s, "POST", probePath, string(b), cookie))
|
||||
if !result.OK || result.Status != 200 || result.Error != "" {
|
||||
t.Fatal(result)
|
||||
}
|
||||
}
|
||||
for _, table := range []string{"keys", "logs"} {
|
||||
var n int
|
||||
if e := s.db.QueryRow("SELECT count(*) FROM " + table).Scan(&n); e != nil || n != 0 {
|
||||
t.Fatal(table, n, e)
|
||||
}
|
||||
}
|
||||
if len(s.logQueue) != 0 {
|
||||
t.Fatal("queued formal log")
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestModelProbeGuardAndInput(t *testing.T) {
|
||||
s := testServer(t)
|
||||
cookie := loginCookie(t, s)
|
||||
for _, prompt := range []string{"", " ", "\t\r\n", "\u3000\u00a0"} {
|
||||
body, _ := json.Marshal(map[string]string{"model": "saved/alias", "prompt": prompt})
|
||||
if w := call(s, "POST", probePath, string(body), cookie); w.Code != 400 {
|
||||
t.Fatalf("blank prompt %q: status %d", prompt, w.Code)
|
||||
}
|
||||
}
|
||||
up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { t.Error("unexpected upstream request") }))
|
||||
defer up.Close()
|
||||
saveProbeChannel(t, s, cookie, up.URL, "openai-chat", false)
|
||||
if w := call(s, "POST", probePath, `{"model":"saved/alias"}`, nil); w.Code != 401 {
|
||||
t.Fatal(w.Code)
|
||||
}
|
||||
for _, header := range []string{"Origin", "X-Requested-With"} {
|
||||
r := httptest.NewRequest("POST", probePath, strings.NewReader(`{"model":"saved/alias"}`))
|
||||
r.AddCookie(cookie)
|
||||
r.Header.Set("Origin", "http://example.com")
|
||||
r.Header.Set("X-Requested-With", "ai-router")
|
||||
r.Header.Del(header)
|
||||
w := httptest.NewRecorder()
|
||||
s.ServeHTTP(w, r)
|
||||
if w.Code != 403 {
|
||||
t.Fatal(w.Code)
|
||||
}
|
||||
}
|
||||
for _, body := range []string{`{"model":"saved/alias"}`, `{"model":"missing/alias"}`} {
|
||||
if w := call(s, "POST", probePath, body, cookie); w.Code != 403 {
|
||||
t.Fatal(w.Code)
|
||||
}
|
||||
}
|
||||
if _, e := s.db.Exec("UPDATE channels SET enabled=1"); e != nil {
|
||||
t.Fatal(e)
|
||||
}
|
||||
if w := call(s, "POST", probePath, `{"model":"saved/unmapped"}`, cookie); w.Code != 403 {
|
||||
t.Fatal(w.Code)
|
||||
}
|
||||
for _, body := range []string{`{}`, `null`, `{"model":"alias"}`, `{"model":"/alias"}`, `{"model":"saved/"}`, `{"model":"saved/alias","prompt":2}`, `{"model":"saved/alias"} {}`, `{"model":"saved/alias","prompt":"` + strings.Repeat("中", 2001) + `"}`} {
|
||||
if w := call(s, "POST", probePath, body, cookie); w.Code != 400 {
|
||||
t.Fatal(body, w.Code)
|
||||
}
|
||||
}
|
||||
if w := call(s, "GET", probePath, "", cookie); w.Code != 405 {
|
||||
t.Fatal(w.Code)
|
||||
}
|
||||
}
|
||||
|
||||
func TestModelProbeFailures(t *testing.T) {
|
||||
for _, tc := range []struct {
|
||||
name, body string
|
||||
status int
|
||||
error string
|
||||
}{
|
||||
{"rejected", "secret-key upstream body", 401, "upstream rejected request"},
|
||||
{"redirect", "secret-key", 302, "upstream rejected request"},
|
||||
{"html", "<html>secret-key</html>", 200, "invalid upstream response"},
|
||||
{"empty", "{}", 200, "invalid upstream response"},
|
||||
{"oversize", strings.Repeat("x", modelTestMaxBody+1), 200, "upstream response too large"},
|
||||
{"limit", probeChat + strings.Repeat(" ", modelTestMaxBody-len(probeChat)), 200, ""},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
s := testServer(t)
|
||||
cookie := loginCookie(t, s)
|
||||
up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if r.URL.Path != "/chat/completions" {
|
||||
t.Error("followed redirect")
|
||||
}
|
||||
w.Header().Set("Location", "/redirected")
|
||||
w.WriteHeader(tc.status)
|
||||
fmt.Fprint(w, tc.body)
|
||||
}))
|
||||
defer up.Close()
|
||||
saveProbeChannel(t, s, cookie, up.URL, "openai-chat", true)
|
||||
result := probeResult(t, call(s, "POST", probePath, `{"model":"saved/alias"}`, cookie))
|
||||
if result.Status != tc.status || result.Error != tc.error || result.OK != (tc.error == "") {
|
||||
t.Fatal(result)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestModelProbeResponseValidation(t *testing.T) {
|
||||
if !validModelTestResponse(protocol.Anthropic, []byte(`{"id":"m","type":"message","role":"assistant","stop_reason":"max_tokens","content":[{"type":"thinking","thinking":"reasoning","signature":"sig"}]}`)) {
|
||||
t.Fatal("valid thinking-only response rejected")
|
||||
}
|
||||
for _, tc := range []struct {
|
||||
kind protocol.Kind
|
||||
body string
|
||||
ok bool
|
||||
}{
|
||||
{protocol.Chat, `{"id":"c","object":"chat.completion","choices":[{"message":{"role":"assistant","content":null},"finish_reason":"length"}]}`, true},
|
||||
{protocol.Responses, `{"id":"r","object":"response","status":"incomplete","incomplete_details":{"reason":"max_output_tokens"},"output":[{"type":"reasoning","summary":[]}]}`, true},
|
||||
{protocol.Anthropic, `{"id":"m","type":"message","role":"assistant","stop_reason":"max_tokens","content":[]}`, true},
|
||||
{protocol.Chat, `{"id":"c","object":"chat.completion","choices":[{}]}`, false},
|
||||
{protocol.Responses, `{"id":"r","object":"response","status":"failed","output":[]}`, false},
|
||||
{protocol.Responses, `{"id":"r","object":"response","status":"completed","output":[{}]}`, false},
|
||||
{protocol.Anthropic, `{"id":"m","type":"message","role":"assistant","stop_reason":"end_turn","content":[{}]}`, false},
|
||||
} {
|
||||
if got := validModelTestResponse(tc.kind, []byte(tc.body)); got != tc.ok {
|
||||
t.Error(tc.kind, tc.body, got)
|
||||
}
|
||||
}
|
||||
for _, kind := range []protocol.Kind{protocol.Chat, protocol.Responses, protocol.Anthropic} {
|
||||
for _, body := range []string{"{}", "null", "[]", probeChat + "{}", `{"error":{"message":"secret-key"}}`} {
|
||||
if validModelTestResponse(kind, []byte(body)) {
|
||||
t.Error(kind, body)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
type probeRoundTripper func(*http.Request) (*http.Response, error)
|
||||
|
||||
func (f probeRoundTripper) RoundTrip(r *http.Request) (*http.Response, error) { return f(r) }
|
||||
|
||||
func TestModelProbeTimeoutAndCancellation(t *testing.T) {
|
||||
for _, cancelRequest := range []bool{false, true} {
|
||||
t.Run(fmt.Sprint(cancelRequest), func(t *testing.T) {
|
||||
s := testServer(t)
|
||||
cookie := loginCookie(t, s)
|
||||
up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}))
|
||||
defer up.Close()
|
||||
saveProbeChannel(t, s, cookie, up.URL, "openai-chat", true)
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
defer cancel()
|
||||
s.client.Transport = probeRoundTripper(func(r *http.Request) (*http.Response, error) {
|
||||
deadline, ok := r.Context().Deadline()
|
||||
if !ok || time.Until(deadline) > 30*time.Second || time.Until(deadline) < 29*time.Second {
|
||||
t.Error("missing 30s deadline", deadline)
|
||||
}
|
||||
if cancelRequest {
|
||||
cancel()
|
||||
<-r.Context().Done()
|
||||
return nil, r.Context().Err()
|
||||
}
|
||||
return nil, context.DeadlineExceeded
|
||||
})
|
||||
r := httptest.NewRequest("POST", probePath, strings.NewReader(`{"model":"saved/alias"}`)).WithContext(ctx)
|
||||
r.AddCookie(cookie)
|
||||
r.Header.Set("Origin", "http://example.com")
|
||||
r.Header.Set("X-Requested-With", "ai-router")
|
||||
w := httptest.NewRecorder()
|
||||
s.ServeHTTP(w, r)
|
||||
result := probeResult(t, w)
|
||||
if result.OK || result.Status != 0 || result.Error == "" {
|
||||
t.Fatal(result)
|
||||
}
|
||||
})
|
||||
}
|
||||
s := testServer(t)
|
||||
canceled := make(chan struct{})
|
||||
up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.WriteHeader(200)
|
||||
w.(http.Flusher).Flush()
|
||||
<-r.Context().Done()
|
||||
close(canceled)
|
||||
}))
|
||||
defer up.Close()
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond)
|
||||
defer cancel()
|
||||
result := s.probeModel(ctx, channel{BaseURL: up.URL, Protocol: "openai-chat"}, "actual", "test")
|
||||
if result.OK || result.Status != 200 || result.Error != "model test timed out" {
|
||||
t.Fatal(result)
|
||||
}
|
||||
select {
|
||||
case <-canceled:
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("upstream was not canceled")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,203 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"ai-router/internal/protocol"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"runtime"
|
||||
"sort"
|
||||
"strings"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
// Explicitly opt in: ROUTER_PERF=1 go test ./internal/server -run TestPerformanceHTTP -v -count=1.
|
||||
func TestPerformanceHTTP(t *testing.T) {
|
||||
if os.Getenv("ROUTER_PERF") != "1" {
|
||||
t.Skip("opt-in load test")
|
||||
}
|
||||
for _, stream := range []bool{false, true} {
|
||||
up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
io.Copy(io.Discard, r.Body)
|
||||
r.Body.Close()
|
||||
if stream {
|
||||
w.Header().Set("Content-Type", "text/event-stream")
|
||||
for _, f := range strings.Split(integrationSSE("openai-chat"), "\n\n") {
|
||||
if f != "" {
|
||||
fmt.Fprint(w, f+"\n\n")
|
||||
w.(http.Flusher).Flush()
|
||||
}
|
||||
}
|
||||
} else {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
fmt.Fprint(w, integrationResponse("openai-chat"))
|
||||
}
|
||||
}))
|
||||
for _, mode := range []string{"direct", "memory", "disk"} {
|
||||
for _, down := range []string{"openai-chat", "anthropic"} {
|
||||
if mode == "direct" && down != "openai-chat" {
|
||||
continue
|
||||
}
|
||||
for _, concurrency := range []int{1, 8, 32} {
|
||||
t.Run(fmt.Sprintf("stream=%t/%s/%s/c%d", stream, mode, down, concurrency), func(t *testing.T) {
|
||||
target := up.URL + integrationPath("openai-chat")
|
||||
raw := "test-key"
|
||||
var s *Server
|
||||
var router *httptest.Server
|
||||
if mode != "direct" {
|
||||
path := ":memory:"
|
||||
if mode == "disk" {
|
||||
path = filepath.Join(t.TempDir(), "router.db")
|
||||
}
|
||||
var err error
|
||||
s, err = New(path, "test-password", nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer s.Close()
|
||||
if err := s.SetLimits(1000000, 64, 128); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
cookie := loginCookie(t, s)
|
||||
ch := integrationChannel(t, s, cookie, "abc", "openai-chat", up.URL)
|
||||
key := integrationCreateKey(t, s, cookie, "perf", ch)
|
||||
raw = key.raw
|
||||
router = httptest.NewServer(s)
|
||||
defer router.Close()
|
||||
target = router.URL + integrationPath(down)
|
||||
}
|
||||
tr := http.DefaultTransport.(*http.Transport).Clone()
|
||||
tr.MaxIdleConnsPerHost = 64
|
||||
defer tr.CloseIdleConnections()
|
||||
client := &http.Client{Transport: tr, Timeout: 10 * time.Second}
|
||||
body := integrationRequest(down, "abc/gpt-5", stream)
|
||||
request := func() (time.Duration, time.Duration, bool) {
|
||||
start := time.Now()
|
||||
r, _ := http.NewRequest("POST", target, strings.NewReader(body))
|
||||
r.Header.Set("Authorization", "Bearer "+raw)
|
||||
r.Header.Set("Content-Type", "application/json")
|
||||
resp, err := client.Do(r)
|
||||
if err != nil {
|
||||
return time.Since(start), 0, false
|
||||
}
|
||||
first := make([]byte, 1)
|
||||
_, err = resp.Body.Read(first)
|
||||
ttfb := time.Since(start)
|
||||
_, e := io.Copy(io.Discard, resp.Body)
|
||||
resp.Body.Close()
|
||||
return time.Since(start), ttfb, resp.StatusCode == 200 && (err == nil || err == io.EOF) && e == nil
|
||||
}
|
||||
for i := 0; i < 20; i++ {
|
||||
if _, _, ok := request(); !ok {
|
||||
t.Fatal("warmup failed")
|
||||
}
|
||||
}
|
||||
const n = 400
|
||||
dur := make([]float64, n)
|
||||
first := make([]float64, n)
|
||||
var index atomic.Int64
|
||||
var failures atomic.Int64
|
||||
var wg sync.WaitGroup
|
||||
runtime.GC()
|
||||
var before, after runtime.MemStats
|
||||
runtime.ReadMemStats(&before)
|
||||
start := time.Now()
|
||||
for c := 0; c < concurrency; c++ {
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
for {
|
||||
i := int(index.Add(1)) - 1
|
||||
if i >= n {
|
||||
return
|
||||
}
|
||||
d, f, ok := request()
|
||||
dur[i] = float64(d) / 1e6
|
||||
first[i] = float64(f) / 1e6
|
||||
if !ok {
|
||||
failures.Add(1)
|
||||
}
|
||||
}
|
||||
}()
|
||||
}
|
||||
wg.Wait()
|
||||
elapsed := time.Since(start)
|
||||
runtime.ReadMemStats(&after)
|
||||
sort.Float64s(dur)
|
||||
sort.Float64s(first)
|
||||
var total float64
|
||||
for _, d := range dur {
|
||||
total += d
|
||||
}
|
||||
drained := elapsed
|
||||
if s != nil {
|
||||
router.Close()
|
||||
if err := s.flushLogs(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
drained = time.Since(start)
|
||||
var count int
|
||||
if err := s.db.QueryRow("SELECT COUNT(*) FROM logs").Scan(&count); err != nil || count != n+20 {
|
||||
t.Fatalf("log count %d err %v", count, err)
|
||||
}
|
||||
}
|
||||
t.Logf("RESULT rps=%.1f drained_rps=%.1f mean_ms=%.3f p50_ms=%.3f p95_ms=%.3f p99_ms=%.3f first_p50_ms=%.3f first_p95_ms=%.3f alloc_B_req=%d errors=%d", float64(n)/elapsed.Seconds(), float64(n)/drained.Seconds(), total/n, dur[n/2], dur[n*95/100], dur[n*99/100], first[n/2], first[n*95/100], (after.TotalAlloc-before.TotalAlloc)/n, failures.Load())
|
||||
if failures.Load() != 0 {
|
||||
t.Fail()
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
up.Close()
|
||||
}
|
||||
}
|
||||
|
||||
func BenchmarkPerformanceConversion(b *testing.B) {
|
||||
for _, size := range []int{1024, 65536} {
|
||||
for _, to := range []protocol.Kind{protocol.Chat, protocol.Responses, protocol.Anthropic} {
|
||||
b.Run(fmt.Sprintf("%d/%s", size, to), func(b *testing.B) {
|
||||
body := []byte(`{"model":"abc/gpt-5","max_tokens":1024,"messages":[{"role":"user","content":"` + strings.Repeat("x", size) + `"}]}`)
|
||||
b.ReportAllocs()
|
||||
b.ResetTimer()
|
||||
for i := 0; i < b.N; i++ {
|
||||
if _, err := protocol.ConvertRequest(protocol.Chat, to, body, "gpt-4"); err != nil {
|
||||
b.Fatal(err)
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func BenchmarkPerformanceStats(b *testing.B) {
|
||||
for _, n := range []int{1000, 10000, 100000} {
|
||||
b.Run(fmt.Sprint(n), func(b *testing.B) {
|
||||
s, err := New(filepath.Join(b.TempDir(), "stats.db"), "test-password", nil)
|
||||
if err != nil {
|
||||
b.Fatal(err)
|
||||
}
|
||||
defer s.Close()
|
||||
_, err = s.db.Exec(`WITH RECURSIVE seq(x) AS (SELECT 1 UNION ALL SELECT x+1 FROM seq WHERE x<?) INSERT INTO logs(created_at,key_id,key_name,channel,request_model,actual_model,downstream,upstream,status,duration_ms,stream,input_tokens,output_tokens,cache_read_tokens,cache_write_tokens,usage_known,error,raw_usage) SELECT '2026-10-20T10:00:00Z','key','perf','abc','abc/gpt-5','gpt-4','openai-chat','openai-chat',200,10,0,1000,100,100,0,1,'','{}' FROM seq`, n)
|
||||
if err != nil {
|
||||
b.Fatal(err)
|
||||
}
|
||||
r := httptest.NewRequest("GET", "/admin/api/stats", nil)
|
||||
b.ReportAllocs()
|
||||
b.ResetTimer()
|
||||
for i := 0; i < b.N; i++ {
|
||||
w := httptest.NewRecorder()
|
||||
s.stats(w, r)
|
||||
if w.Code != 200 {
|
||||
b.Fatal(w.Code)
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,273 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"ai-router/internal/protocol"
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"io"
|
||||
"net/http"
|
||||
"sort"
|
||||
"strings"
|
||||
"time"
|
||||
)
|
||||
|
||||
func resolve(cs []channel, k apiKey, model string) (channel, string, error) {
|
||||
name, alias, ok := strings.Cut(model, "/")
|
||||
if !ok || alias == "" {
|
||||
return channel{}, "", errors.New("model must be channel/alias")
|
||||
}
|
||||
for _, c := range cs {
|
||||
if c.Name != name {
|
||||
continue
|
||||
}
|
||||
if !c.Enabled || !k.allows(c.ID) {
|
||||
return channel{}, "", errors.New("channel unavailable")
|
||||
}
|
||||
if actual, ok := c.Models[alias]; ok {
|
||||
return c, actual, nil
|
||||
}
|
||||
return channel{}, "", errors.New("model unavailable")
|
||||
}
|
||||
return channel{}, "", errors.New("channel unavailable")
|
||||
}
|
||||
func (s *Server) proxy(w http.ResponseWriter, r *http.Request) {
|
||||
kinds := map[string]protocol.Kind{"/v1/chat/completions": protocol.Chat, "/v1/responses": protocol.Responses, "/v1/messages": protocol.Anthropic}
|
||||
down, ok := kinds[r.URL.Path]
|
||||
models := r.URL.Path == "/v1/models"
|
||||
if !ok && !models {
|
||||
errJSON(w, 404, "not found")
|
||||
return
|
||||
}
|
||||
if (models && r.Method != "GET") || (!models && r.Method != "POST") {
|
||||
errJSON(w, 405, "method not allowed")
|
||||
return
|
||||
}
|
||||
ip := s.clientIP(r)
|
||||
if !s.authAllowed(ip) {
|
||||
w.Header().Set("Retry-After", "60")
|
||||
errJSON(w, 429, "request limit exceeded")
|
||||
return
|
||||
}
|
||||
k, e := s.authenticate(r)
|
||||
if e != nil {
|
||||
s.authFail(ip)
|
||||
errJSON(w, 401, "invalid API key")
|
||||
return
|
||||
}
|
||||
release, retry := s.limiter.acquire(k.ID, time.Now())
|
||||
if release == nil {
|
||||
w.Header().Set("Retry-After", retry)
|
||||
errJSON(w, 429, "request limit exceeded")
|
||||
return
|
||||
}
|
||||
start := time.Now()
|
||||
l := requestLog{CreatedAt: start.UTC().Format(time.RFC3339Nano), KeyID: k.ID, KeyName: k.Name, Downstream: string(down), Status: 500}
|
||||
// Release the concurrency slot before the log write, which can block on a
|
||||
// full queue; rate-limited and unauthenticated requests stay unlogged.
|
||||
defer func() {
|
||||
release()
|
||||
if !models {
|
||||
l.DurationMS = time.Since(start).Milliseconds()
|
||||
s.record(l)
|
||||
}
|
||||
}()
|
||||
fail := func(status int, msg string) { l.Status = status; l.Error = msg; errJSON(w, status, msg) }
|
||||
cs, e := s.loadChannels()
|
||||
if e != nil {
|
||||
fail(500, "database error")
|
||||
return
|
||||
}
|
||||
if models {
|
||||
ids := []string{}
|
||||
for _, c := range cs {
|
||||
if c.Enabled && k.allows(c.ID) {
|
||||
for alias := range c.Models {
|
||||
ids = append(ids, c.Name+"/"+alias)
|
||||
}
|
||||
}
|
||||
}
|
||||
sort.Strings(ids)
|
||||
out := []any{}
|
||||
for _, id := range ids {
|
||||
out = append(out, map[string]string{"id": id, "object": "model", "owned_by": "ai-router"})
|
||||
}
|
||||
l.Status = 200
|
||||
writeJSON(w, map[string]any{"object": "list", "data": out})
|
||||
return
|
||||
}
|
||||
var b []byte
|
||||
e = readBody(w, r, maxBody, bodyReadTimeout, func() error {
|
||||
var err error
|
||||
b, err = io.ReadAll(r.Body)
|
||||
return err
|
||||
})
|
||||
if e != nil {
|
||||
var tooBig *http.MaxBytesError
|
||||
if errors.As(e, &tooBig) {
|
||||
fail(413, "request body too large")
|
||||
} else {
|
||||
fail(400, "request body unreadable")
|
||||
}
|
||||
return
|
||||
}
|
||||
var req struct {
|
||||
Model string `json:"model"`
|
||||
Stream bool `json:"stream"`
|
||||
}
|
||||
if json.Unmarshal(b, &req) != nil || req.Model == "" {
|
||||
fail(400, "invalid request")
|
||||
return
|
||||
}
|
||||
l.RequestModel = req.Model
|
||||
l.Stream = req.Stream
|
||||
c, actual, e := resolve(cs, k, req.Model)
|
||||
if e != nil {
|
||||
fail(403, "model or channel unavailable")
|
||||
return
|
||||
}
|
||||
l.Channel = c.Name
|
||||
l.ActualModel = actual
|
||||
l.Upstream = c.Protocol
|
||||
up := protocol.Kind(c.Protocol)
|
||||
payload, e := protocol.ConvertRequest(down, up, b, actual)
|
||||
if e != nil {
|
||||
fail(400, "unsupported request conversion")
|
||||
return
|
||||
}
|
||||
if req.Stream && up == protocol.Chat {
|
||||
var obj map[string]json.RawMessage
|
||||
if json.Unmarshal(payload, &obj) == nil {
|
||||
options := map[string]json.RawMessage{}
|
||||
if raw, ok := obj["stream_options"]; ok {
|
||||
_ = json.Unmarshal(raw, &options)
|
||||
}
|
||||
if options == nil {
|
||||
options = map[string]json.RawMessage{}
|
||||
}
|
||||
options["include_usage"] = json.RawMessage("true")
|
||||
obj["stream_options"], _ = json.Marshal(options)
|
||||
payload, _ = json.Marshal(obj)
|
||||
}
|
||||
}
|
||||
suffix := map[protocol.Kind]string{protocol.Chat: "/chat/completions", protocol.Responses: "/responses", protocol.Anthropic: "/v1/messages"}[up]
|
||||
ctx, cancel := context.WithTimeout(r.Context(), 10*time.Minute)
|
||||
defer cancel()
|
||||
request, e := http.NewRequestWithContext(ctx, "POST", c.BaseURL+suffix, bytes.NewReader(payload))
|
||||
if e != nil {
|
||||
fail(502, "invalid upstream configuration")
|
||||
return
|
||||
}
|
||||
request.Header.Set("Content-Type", "application/json")
|
||||
if up == protocol.Anthropic {
|
||||
request.Header.Set("x-api-key", c.APIKey)
|
||||
request.Header.Set("anthropic-version", "2023-06-01")
|
||||
if down == up {
|
||||
if v := r.Header.Get("anthropic-version"); safeForwardedHeader(v, 32) {
|
||||
request.Header.Set("anthropic-version", v)
|
||||
}
|
||||
if v := r.Header.Get("anthropic-beta"); safeForwardedHeader(v, 256) {
|
||||
request.Header.Set("anthropic-beta", v)
|
||||
}
|
||||
}
|
||||
} else {
|
||||
request.Header.Set("Authorization", "Bearer "+c.APIKey)
|
||||
}
|
||||
response, e := s.client.Do(request)
|
||||
if e != nil {
|
||||
fail(502, "upstream request failed")
|
||||
return
|
||||
}
|
||||
defer response.Body.Close()
|
||||
if response.StatusCode < 200 || response.StatusCode >= 300 {
|
||||
status := response.StatusCode
|
||||
if status < 400 {
|
||||
status = 502
|
||||
}
|
||||
fail(status, "upstream rejected request")
|
||||
return
|
||||
}
|
||||
if req.Stream {
|
||||
w.Header().Set("Content-Type", "text/event-stream")
|
||||
w.Header().Set("Cache-Control", "no-cache")
|
||||
w.Header().Set("X-Accel-Buffering", "no")
|
||||
w.WriteHeader(response.StatusCode)
|
||||
l.Status = response.StatusCode
|
||||
streamWriter := &deadlineWriter{ResponseWriter: w}
|
||||
u, e := protocol.Stream(ctx, streamWriter, response.Body, up, down, req.Model)
|
||||
l.Usage = u
|
||||
if e != nil {
|
||||
l.Status = 502
|
||||
l.Error = "upstream stream interrupted"
|
||||
l.Usage.Known = false
|
||||
if ctx.Err() == nil {
|
||||
event := "data: {\"error\":{\"type\":\"upstream_error\",\"message\":\"upstream stream interrupted\"}}\n\n"
|
||||
if down != protocol.Chat {
|
||||
event = "event: error\n" + event
|
||||
}
|
||||
_, _ = io.WriteString(streamWriter, event)
|
||||
if f, ok := w.(http.Flusher); ok {
|
||||
f.Flush()
|
||||
}
|
||||
}
|
||||
}
|
||||
return
|
||||
}
|
||||
b, e = io.ReadAll(io.LimitReader(response.Body, maxBody+1))
|
||||
if e != nil || len(b) > maxBody {
|
||||
fail(502, "invalid upstream response")
|
||||
return
|
||||
}
|
||||
l.Usage = protocol.ExtractUsage(up, b)
|
||||
out, e := protocol.ConvertResponse(up, down, b)
|
||||
if e != nil {
|
||||
fail(502, "unsupported response conversion")
|
||||
return
|
||||
}
|
||||
l.Status = response.StatusCode
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
w.WriteHeader(response.StatusCode)
|
||||
if _, e = w.Write(out); e != nil {
|
||||
l.Status = 499
|
||||
l.Error = "client disconnected"
|
||||
}
|
||||
}
|
||||
|
||||
type deadlineWriter struct{ http.ResponseWriter }
|
||||
|
||||
// safeForwardedHeader reports whether v is a short printable token list that
|
||||
// is safe to forward upstream; anything else is dropped so clients cannot
|
||||
// smuggle arbitrary features or exotic header content past the channel
|
||||
// configuration.
|
||||
func safeForwardedHeader(v string, max int) bool {
|
||||
if len(v) == 0 || len(v) > max {
|
||||
return false
|
||||
}
|
||||
for i := 0; i < len(v); i++ {
|
||||
c := v[i]
|
||||
switch {
|
||||
case c >= 'a' && c <= 'z', c >= 'A' && c <= 'Z', c >= '0' && c <= '9':
|
||||
case c == '-' || c == '_' || c == '.' || c == ',' || c == ' ':
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
func (w *deadlineWriter) Unwrap() http.ResponseWriter { return w.ResponseWriter }
|
||||
func (w *deadlineWriter) WriteHeader(status int) {
|
||||
_ = http.NewResponseController(w.ResponseWriter).SetWriteDeadline(time.Now().Add(30 * time.Second))
|
||||
w.ResponseWriter.WriteHeader(status)
|
||||
}
|
||||
|
||||
func (w *deadlineWriter) Write(b []byte) (int, error) {
|
||||
_ = http.NewResponseController(w.ResponseWriter).SetWriteDeadline(time.Now().Add(30 * time.Second))
|
||||
return w.ResponseWriter.Write(b)
|
||||
}
|
||||
func (w *deadlineWriter) FlushError() error {
|
||||
_ = http.NewResponseController(w.ResponseWriter).SetWriteDeadline(time.Now().Add(30 * time.Second))
|
||||
return http.NewResponseController(w.ResponseWriter).Flush()
|
||||
}
|
||||
func (w *deadlineWriter) Flush() { _ = w.FlushError() }
|
||||
@@ -0,0 +1,203 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"net"
|
||||
"net/http"
|
||||
"net/netip"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
)
|
||||
|
||||
const bodyReadTimeout = 30 * time.Second
|
||||
const maxLimitEntries = 10000
|
||||
|
||||
// Login rate limits: failures only, per IP and globally; successful logins
|
||||
// never consume or hit these buckets.
|
||||
const (
|
||||
loginFailuresPerMinute = 10
|
||||
loginGlobalFailuresPerMinute = 300
|
||||
)
|
||||
|
||||
// authFailuresPerMinute bounds how many failed /v1 authentications one client
|
||||
// IP may attempt per minute before further attempts are rejected before the
|
||||
// database lookup.
|
||||
const authFailuresPerMinute = 60
|
||||
|
||||
func pruneFailures(m map[string][]time.Time, now time.Time) {
|
||||
for k, ts := range m {
|
||||
if len(ts) == 0 || now.Sub(ts[len(ts)-1]) > time.Minute {
|
||||
delete(m, k)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// trimRecent drops entries older than one minute while reusing the backing
|
||||
// array, so a long-lived failure window cannot grow without bound.
|
||||
func trimRecent(ts []time.Time, now time.Time) []time.Time {
|
||||
i := 0
|
||||
for i < len(ts) && now.Sub(ts[i]) >= time.Minute {
|
||||
i++
|
||||
}
|
||||
if i > 0 {
|
||||
ts = append(ts[:0], ts[i:]...)
|
||||
}
|
||||
return ts
|
||||
}
|
||||
|
||||
// authAllowed reports whether the client IP may attempt authentication. Only
|
||||
// failures are counted, so legitimate traffic is unaffected.
|
||||
func (s *Server) authAllowed(ip string) bool {
|
||||
now := time.Now()
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
if now.Sub(s.authFailuresSeen) >= time.Minute {
|
||||
pruneFailures(s.authFailures, now)
|
||||
s.authFailuresSeen = now
|
||||
}
|
||||
recent := 0
|
||||
for _, t := range s.authFailures[ip] {
|
||||
if now.Sub(t) < time.Minute {
|
||||
recent++
|
||||
}
|
||||
}
|
||||
return recent < authFailuresPerMinute && (recent > 0 || len(s.authFailures) < maxLimitEntries)
|
||||
}
|
||||
|
||||
// authFail records one failed authentication from ip.
|
||||
func (s *Server) authFail(ip string) {
|
||||
now := time.Now()
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
if _, ok := s.authFailures[ip]; !ok && len(s.authFailures) >= maxLimitEntries {
|
||||
return
|
||||
}
|
||||
s.authFailures[ip] = append(s.authFailures[ip], now)
|
||||
}
|
||||
|
||||
func readBody(w http.ResponseWriter, r *http.Request, limit int64, timeout time.Duration, read func() error) error {
|
||||
controller := http.NewResponseController(w)
|
||||
_ = controller.SetReadDeadline(time.Now().Add(timeout))
|
||||
defer controller.SetReadDeadline(time.Time{})
|
||||
r.Body = http.MaxBytesReader(w, r.Body, limit)
|
||||
defer r.Body.Close()
|
||||
return read()
|
||||
}
|
||||
|
||||
// SetTrustedProxies must be called before serving requests.
|
||||
func (s *Server) SetTrustedProxies(cidrs []string) error {
|
||||
prefixes := make([]netip.Prefix, 0, len(cidrs))
|
||||
for _, cidr := range cidrs {
|
||||
p, err := netip.ParsePrefix(cidr)
|
||||
if err != nil {
|
||||
return fmt.Errorf("invalid trusted proxy CIDR %q: %w", cidr, err)
|
||||
}
|
||||
if p.Bits() == 0 {
|
||||
return fmt.Errorf("trusted proxy CIDR must not trust all addresses")
|
||||
}
|
||||
prefixes = append(prefixes, p)
|
||||
}
|
||||
s.trustedProxies = prefixes
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *Server) trusted(ip netip.Addr) bool {
|
||||
for _, p := range s.trustedProxies {
|
||||
if p.Contains(ip) {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func (s *Server) clientIP(r *http.Request) string {
|
||||
host, _, err := net.SplitHostPort(r.RemoteAddr)
|
||||
if err != nil {
|
||||
host = r.RemoteAddr
|
||||
}
|
||||
peer, err := netip.ParseAddr(host)
|
||||
if err != nil {
|
||||
return host
|
||||
}
|
||||
peer = peer.Unmap()
|
||||
if !s.trusted(peer) {
|
||||
return peer.String()
|
||||
}
|
||||
parts := strings.Split(strings.Join(r.Header.Values("X-Forwarded-For"), ","), ",")
|
||||
ips := make([]netip.Addr, len(parts))
|
||||
for i, part := range parts {
|
||||
ip, err := netip.ParseAddr(strings.TrimSpace(part))
|
||||
if err != nil || ip.Zone() != "" {
|
||||
return peer.String()
|
||||
}
|
||||
ips[i] = ip.Unmap()
|
||||
}
|
||||
for i := len(ips) - 1; i >= 0; i-- {
|
||||
if !s.trusted(ips[i]) {
|
||||
return ips[i].String()
|
||||
}
|
||||
}
|
||||
return peer.String()
|
||||
}
|
||||
|
||||
type limitEntry struct {
|
||||
window time.Time
|
||||
requests, active int
|
||||
}
|
||||
type requestLimiter struct {
|
||||
mu sync.Mutex
|
||||
rpm, perKey, global, active int
|
||||
entries map[string]*limitEntry
|
||||
nextCleanup time.Time
|
||||
}
|
||||
|
||||
func newRequestLimiter(rpm, perKey, global int) *requestLimiter {
|
||||
return &requestLimiter{rpm: rpm, perKey: perKey, global: global, entries: make(map[string]*limitEntry)}
|
||||
}
|
||||
|
||||
// SetLimits must be called before serving requests.
|
||||
func (s *Server) SetLimits(rpm, perKey, global int) error {
|
||||
if rpm <= 0 || perKey <= 0 || global <= 0 {
|
||||
return fmt.Errorf("limits must be positive integers")
|
||||
}
|
||||
s.limiter = newRequestLimiter(rpm, perKey, global)
|
||||
return nil
|
||||
}
|
||||
func (l *requestLimiter) acquire(key string, now time.Time) (func(), string) {
|
||||
l.mu.Lock()
|
||||
defer l.mu.Unlock()
|
||||
if !now.Before(l.nextCleanup) {
|
||||
for k, e := range l.entries {
|
||||
if e.active == 0 && now.Sub(e.window) >= time.Minute {
|
||||
delete(l.entries, k)
|
||||
}
|
||||
}
|
||||
l.nextCleanup = now.Add(time.Minute)
|
||||
}
|
||||
e := l.entries[key]
|
||||
if e == nil {
|
||||
if len(l.entries) >= maxLimitEntries {
|
||||
return nil, "60"
|
||||
}
|
||||
e = &limitEntry{window: now}
|
||||
l.entries[key] = e
|
||||
}
|
||||
if now.Sub(e.window) >= time.Minute {
|
||||
e.window = now
|
||||
e.requests = 0
|
||||
}
|
||||
if e.requests >= l.rpm {
|
||||
seconds := int((time.Minute - now.Sub(e.window) + time.Second - 1) / time.Second)
|
||||
return nil, strconv.Itoa(seconds)
|
||||
}
|
||||
if e.active >= l.perKey || l.active >= l.global {
|
||||
return nil, "1"
|
||||
}
|
||||
e.requests++
|
||||
e.active++
|
||||
l.active++
|
||||
var once sync.Once
|
||||
return func() { once.Do(func() { l.mu.Lock(); defer l.mu.Unlock(); e.active--; l.active-- }) }, ""
|
||||
}
|
||||
@@ -0,0 +1,161 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"ai-router/internal/protocol"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
"unicode/utf8"
|
||||
)
|
||||
|
||||
func TestSecurityStatsSQLSemantics(t *testing.T) {
|
||||
s := testServer(t)
|
||||
now := time.Now().UTC().Format(time.RFC3339Nano)
|
||||
for _, l := range []requestLog{
|
||||
{CreatedAt: now, KeyID: "a", KeyName: "old", Status: 200, Usage: protocol.Usage{Known: true, Input: 8}},
|
||||
{CreatedAt: now, KeyID: "a", KeyName: "new", Status: 200, Error: "failed", Usage: protocol.Usage{Raw: []byte("{}"), Input: 100}},
|
||||
{CreatedAt: now, KeyID: "b", Status: 400},
|
||||
{CreatedAt: now, KeyID: "c", Status: 204, Usage: protocol.Usage{Known: true}},
|
||||
} {
|
||||
s.record(l)
|
||||
}
|
||||
if err := s.flushLogs(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
for _, tc := range []struct {
|
||||
query string
|
||||
requests, success, unknown int64
|
||||
input *int64
|
||||
}{
|
||||
{"", 4, 2, 2, new(int64(8))},
|
||||
{"?key_id=b", 1, 0, 1, nil},
|
||||
{"?key_id=c", 1, 1, 0, new(int64(0))},
|
||||
{"?key_id=missing", 0, 0, 0, nil},
|
||||
} {
|
||||
w := httptest.NewRecorder()
|
||||
s.stats(w, httptest.NewRequest("GET", "/admin/api/stats"+tc.query, nil))
|
||||
if w.Code != 200 {
|
||||
t.Fatal(w.Body.String())
|
||||
}
|
||||
var got struct {
|
||||
Summary summary
|
||||
ByKey []struct{ Name string } `json:"by_key"`
|
||||
}
|
||||
if err := json.Unmarshal(w.Body.Bytes(), &got); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if got.Summary.Requests != tc.requests || got.Summary.Success != tc.success || got.Summary.Unknown != tc.unknown {
|
||||
t.Fatal(w.Body.String())
|
||||
}
|
||||
if (got.Summary.Input == nil) != (tc.input == nil) || tc.input != nil && *got.Summary.Input != *tc.input {
|
||||
t.Fatal(w.Body.String())
|
||||
}
|
||||
if tc.query == "" && got.ByKey[0].Name != "new" {
|
||||
t.Fatal(w.Body.String())
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestSecurityLogFieldsBounded(t *testing.T) {
|
||||
s := testServer(t)
|
||||
huge := strings.Repeat("界", 10000)
|
||||
s.record(requestLog{RequestModel: huge, ActualModel: huge, Channel: huge, KeyID: huge, KeyName: huge, Upstream: huge, Downstream: huge, Error: huge, Usage: protocol.Usage{Known: true, Input: 12, Raw: []byte(huge)}})
|
||||
if err := s.flushLogs(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
var model, actual, channel, key, name, up, down, message, raw string
|
||||
var input int64
|
||||
if err := s.db.QueryRow("SELECT request_model,actual_model,channel,key_id,key_name,upstream,downstream,error,raw_usage,input_tokens FROM logs").Scan(&model, &actual, &channel, &key, &name, &up, &down, &message, &raw, &input); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
for _, field := range []string{model, actual, channel, key, name, up, down} {
|
||||
if len(field) > maxLogFieldBytes || !utf8.ValidString(field) {
|
||||
t.Fatal("unbounded or invalid field")
|
||||
}
|
||||
}
|
||||
if len(message) > maxLogErrorBytes || !utf8.ValidString(message) || len(raw) > maxLogUsageBytes || !json.Valid([]byte(raw)) || input != 12 {
|
||||
t.Fatal("invalid bounded log")
|
||||
}
|
||||
}
|
||||
|
||||
func TestSecurityLogRetentionOptInAndBatch(t *testing.T) {
|
||||
s := testServer(t)
|
||||
old := time.Now().Add(-91 * 24 * time.Hour).Format(time.RFC3339Nano)
|
||||
for i := 0; i < logCleanupBatch+3; i++ {
|
||||
s.record(requestLog{CreatedAt: old})
|
||||
}
|
||||
s.record(requestLog{CreatedAt: time.Now().Format(time.RFC3339Nano)})
|
||||
if err := s.flushLogs(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if s.logRetentionDays.Load() != 0 {
|
||||
t.Fatal("retention enabled by default")
|
||||
}
|
||||
if err := s.purgeExpiredLogs(context.Background(), 0); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if n := logCount(t, s.db); n != logCleanupBatch+4 {
|
||||
t.Fatal(n)
|
||||
}
|
||||
s.SetLogRetention(90)
|
||||
if s.logRetentionDays.Load() != 90 {
|
||||
t.Fatal("retention not set")
|
||||
}
|
||||
if err := s.purgeExpiredLogs(context.Background(), 90); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if n := logCount(t, s.db); n != 4 {
|
||||
t.Fatal("unbounded purge", n)
|
||||
}
|
||||
if err := s.purgeExpiredLogs(context.Background(), 90); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if n := logCount(t, s.db); n != 1 {
|
||||
t.Fatal(n)
|
||||
}
|
||||
s.SetLogRetention(0)
|
||||
if s.logRetentionDays.Load() != 0 {
|
||||
t.Fatal("retention not disabled")
|
||||
}
|
||||
s.SetLogRetention(90)
|
||||
if err := s.Close(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestStatsGroupLimitAndCancellation(t *testing.T) {
|
||||
s := testServer(t)
|
||||
tx, err := s.db.Begin()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
for i := 0; i < 10001; i++ {
|
||||
if err := insertLog(context.Background(), tx, requestLog{CreatedAt: "2026-01-01T00:00:00Z", KeyID: fmt.Sprint(i)}); err != nil {
|
||||
tx.Rollback()
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
if err := tx.Commit(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
w := httptest.NewRecorder()
|
||||
s.stats(w, httptest.NewRequest("GET", "/admin/api/stats", nil))
|
||||
if w.Code != 422 {
|
||||
t.Fatal("unbounded group result", w.Code)
|
||||
}
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
cancel()
|
||||
w = httptest.NewRecorder()
|
||||
s.stats(w, httptest.NewRequest("GET", "/admin/api/stats", nil).WithContext(ctx))
|
||||
if w.Code == 200 {
|
||||
t.Fatal("cancelled query succeeded")
|
||||
}
|
||||
var day string
|
||||
if err := s.db.QueryRow("SELECT router_log_day(NULL, 'UTC')").Scan(&day); err == nil {
|
||||
t.Fatal("invalid timestamp accepted")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,359 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"fmt"
|
||||
"io"
|
||||
"net"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func TestClientIPTrust(t *testing.T) {
|
||||
s := &Server{}
|
||||
if err := s.SetTrustedProxies([]string{"10.0.0.0/8", "::1/128"}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := s.SetTrustedProxies([]string{"bad"}); err == nil {
|
||||
t.Fatal("invalid CIDR accepted")
|
||||
}
|
||||
for _, tt := range []struct{ peer, xff, want string }{
|
||||
{"192.0.2.1:42", "198.51.100.1", "192.0.2.1"},
|
||||
{"10.0.0.1:42", "198.51.100.1, 10.0.0.2", "198.51.100.1"},
|
||||
{"10.0.0.1:42", "198.51.100.1, 192.0.2.2", "192.0.2.2"},
|
||||
{"10.0.0.1:42", "bad, 192.0.2.2", "10.0.0.1"},
|
||||
{"10.0.0.1:42", "", "10.0.0.1"},
|
||||
{"10.0.0.1:42", "10.0.0.2", "10.0.0.1"},
|
||||
{"[::1]:42", "2001:db8::1", "2001:db8::1"},
|
||||
} {
|
||||
r := httptest.NewRequest("POST", "/", nil)
|
||||
r.RemoteAddr = tt.peer
|
||||
r.Header.Set("X-Forwarded-For", tt.xff)
|
||||
r.Header.Set("X-Real-IP", "203.0.113.1")
|
||||
if got := s.clientIP(r); got != tt.want {
|
||||
t.Errorf("%+v: got %s", tt, got)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestLoginProxyIsolation(t *testing.T) {
|
||||
s, err := New(":memory:", "password", nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer s.Close()
|
||||
if err := s.SetTrustedProxies([]string{"10.0.0.0/8"}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
login := func(peer, forwarded string) int {
|
||||
r := httptest.NewRequest("POST", "/admin/api/login", strings.NewReader(`{"password":"wrong"}`))
|
||||
r.RemoteAddr = peer
|
||||
r.Header.Set("X-Forwarded-For", forwarded)
|
||||
w := httptest.NewRecorder()
|
||||
s.login(w, r)
|
||||
return w.Code
|
||||
}
|
||||
for i := 0; i < 10; i++ {
|
||||
if got := login("10.0.0.1:1", "192.0.2.1"); got != 401 {
|
||||
t.Fatal(got)
|
||||
}
|
||||
}
|
||||
if got := login("10.0.0.1:1", "192.0.2.1"); got != 429 {
|
||||
t.Fatal(got)
|
||||
}
|
||||
if got := login("10.0.0.1:1", "192.0.2.2"); got != 401 {
|
||||
t.Fatal(got)
|
||||
}
|
||||
for i := 0; i < 10; i++ {
|
||||
login("192.0.2.3:1", fmt.Sprintf("198.51.100.%d", i))
|
||||
}
|
||||
if got := login("192.0.2.3:1", "203.0.113.1"); got != 429 {
|
||||
t.Fatal(got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestBodyDeadline(t *testing.T) {
|
||||
result := make(chan error, 1)
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
err := readBody(w, r, 1024, 40*time.Millisecond, func() error { _, err := io.ReadAll(r.Body); return err })
|
||||
result <- err
|
||||
w.WriteHeader(400)
|
||||
}))
|
||||
defer srv.Close()
|
||||
conn, err := net.Dial("tcp", srv.Listener.Addr().String())
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer conn.Close()
|
||||
fmt.Fprint(conn, "POST / HTTP/1.1\r\nHost: test\r\nContent-Length: 20\r\n\r\nx")
|
||||
select {
|
||||
case err := <-result:
|
||||
if err == nil {
|
||||
t.Fatal("slow body accepted")
|
||||
}
|
||||
case <-time.After(2 * time.Second):
|
||||
t.Fatal("body deadline did not fire")
|
||||
}
|
||||
}
|
||||
|
||||
func TestBodyDeadlineCleared(t *testing.T) {
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
err := readBody(w, r, 1024, 40*time.Millisecond, func() error { _, err := io.ReadAll(r.Body); return err })
|
||||
if err != nil {
|
||||
http.Error(w, "read", 400)
|
||||
return
|
||||
}
|
||||
time.Sleep(100 * time.Millisecond)
|
||||
select {
|
||||
case <-r.Context().Done():
|
||||
http.Error(w, "canceled", 500)
|
||||
return
|
||||
default:
|
||||
}
|
||||
fmt.Fprint(w, "data: ok\n\n")
|
||||
}))
|
||||
defer srv.Close()
|
||||
conn, err := net.Dial("tcp", srv.Listener.Addr().String())
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer conn.Close()
|
||||
conn.SetDeadline(time.Now().Add(2 * time.Second))
|
||||
fmt.Fprint(conn, "POST / HTTP/1.1\r\nHost: test\r\nContent-Length: 1\r\n\r\nx")
|
||||
resp, err := http.ReadResponse(bufio.NewReader(conn), nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
if resp.StatusCode != 200 {
|
||||
t.Fatal(resp.StatusCode)
|
||||
}
|
||||
}
|
||||
|
||||
func TestLoginBodyLimit(t *testing.T) {
|
||||
r := httptest.NewRequest("POST", "/admin/api/login", strings.NewReader(`{"password":"`+strings.Repeat("x", 16<<10)+`"}`))
|
||||
var v any
|
||||
if decode(httptest.NewRecorder(), r, &v) == nil {
|
||||
t.Fatal("oversize login accepted")
|
||||
}
|
||||
}
|
||||
|
||||
func TestLimiterReleaseAndCleanup(t *testing.T) {
|
||||
now := time.Now()
|
||||
l := newRequestLimiter(2, 1, 2)
|
||||
a, _ := l.acquire("a", now)
|
||||
if a == nil {
|
||||
t.Fatal("denied")
|
||||
}
|
||||
if release, _ := l.acquire("a", now); release != nil {
|
||||
t.Fatal("per-key concurrency bypass")
|
||||
}
|
||||
b, _ := l.acquire("b", now)
|
||||
if b == nil {
|
||||
t.Fatal("denied")
|
||||
}
|
||||
if release, _ := l.acquire("c", now); release != nil {
|
||||
t.Fatal("global concurrency bypass")
|
||||
}
|
||||
a()
|
||||
a()
|
||||
b()
|
||||
a, _ = l.acquire("a", now)
|
||||
if a == nil {
|
||||
t.Fatal("not released")
|
||||
}
|
||||
a()
|
||||
if release, retry := l.acquire("a", now); release != nil || retry != "60" {
|
||||
t.Fatal("rpm bypass")
|
||||
}
|
||||
c, _ := l.acquire("c", now.Add(time.Minute))
|
||||
if c == nil {
|
||||
t.Fatal("window not reset")
|
||||
}
|
||||
c()
|
||||
if len(l.entries) != 1 {
|
||||
t.Fatalf("stale entries: %d", len(l.entries))
|
||||
}
|
||||
for i := 0; i < maxLimitEntries; i++ {
|
||||
l.entries[fmt.Sprint(i)] = &limitEntry{window: now.Add(time.Minute)}
|
||||
}
|
||||
if release, _ := l.acquire("overflow", now.Add(time.Minute)); release != nil {
|
||||
t.Fatal("map cap bypass")
|
||||
}
|
||||
}
|
||||
|
||||
func TestLimiterConcurrent(t *testing.T) {
|
||||
l := newRequestLimiter(10000, 8, 64)
|
||||
var wg sync.WaitGroup
|
||||
for i := 0; i < 100; i++ {
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
if release, _ := l.acquire("key", time.Now()); release != nil {
|
||||
release()
|
||||
}
|
||||
}()
|
||||
}
|
||||
wg.Wait()
|
||||
if l.active != 0 || l.entries["key"].active != 0 {
|
||||
t.Fatal("permit leak")
|
||||
}
|
||||
}
|
||||
|
||||
func TestProxyLimitRejectsBeforeBodyAndLogging(t *testing.T) {
|
||||
s, err := New(":memory:", "test-password", nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer s.Close()
|
||||
if err := s.SetLimits(1, 1, 1); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
cookie := loginCookie(t, s)
|
||||
ch := integrationChannel(t, s, cookie, "abc", "openai-chat", "http://127.0.0.1:1")
|
||||
key := integrationCreateKey(t, s, cookie, "limited", ch)
|
||||
request := func(body io.Reader) *httptest.ResponseRecorder {
|
||||
r := httptest.NewRequest("POST", "/v1/chat/completions", body)
|
||||
r.Header.Set("Authorization", "Bearer "+key.raw)
|
||||
w := httptest.NewRecorder()
|
||||
s.ServeHTTP(w, r)
|
||||
return w
|
||||
}
|
||||
if w := request(strings.NewReader("{}")); w.Code != 400 {
|
||||
t.Fatal(w.Code)
|
||||
}
|
||||
if s.limiter.active != 0 {
|
||||
t.Fatal("permit leaked after invalid request")
|
||||
}
|
||||
body := &observedReader{}
|
||||
w := request(body)
|
||||
if w.Code != 429 || w.Header().Get("Retry-After") == "" || body.read {
|
||||
t.Fatalf("status=%d read=%t", w.Code, body.read)
|
||||
}
|
||||
if err := s.flushLogs(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
var count int
|
||||
if err := s.db.QueryRow("SELECT count(*) FROM logs").Scan(&count); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if count != 1 {
|
||||
t.Fatalf("rejected request logged: %d logs", count)
|
||||
}
|
||||
}
|
||||
|
||||
type observedReader struct{ read bool }
|
||||
|
||||
func (r *observedReader) Read(p []byte) (int, error) { r.read = true; return 0, io.EOF }
|
||||
|
||||
func TestLoginFailuresOnlyAndGlobalCap(t *testing.T) {
|
||||
s, err := New(":memory:", "password", nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer s.Close()
|
||||
if err := s.SetTrustedProxies([]string{"10.0.0.0/8"}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
login := func(peer, forwarded, password string) int {
|
||||
r := httptest.NewRequest("POST", "/admin/api/login", strings.NewReader(`{"password":"`+password+`"}`))
|
||||
r.RemoteAddr = peer
|
||||
if forwarded != "" {
|
||||
r.Header.Set("X-Forwarded-For", forwarded)
|
||||
}
|
||||
w := httptest.NewRecorder()
|
||||
s.login(w, r)
|
||||
return w.Code
|
||||
}
|
||||
for i := 0; i < 15; i++ {
|
||||
if got := login("10.0.0.1:1", "192.0.2.1", "password"); got != 200 {
|
||||
t.Fatalf("correct login %d rejected: %d", i, got)
|
||||
}
|
||||
}
|
||||
for i := 0; i < loginFailuresPerMinute; i++ {
|
||||
if got := login("10.0.0.1:1", "192.0.2.2", "wrong"); got != 401 {
|
||||
t.Fatalf("failure %d: %d", i, got)
|
||||
}
|
||||
}
|
||||
if got := login("10.0.0.1:1", "192.0.2.2", "wrong"); got != 429 {
|
||||
t.Fatal(got)
|
||||
}
|
||||
// The per-IP lockout is uniform: it must not confirm candidate passwords.
|
||||
if got := login("10.0.0.1:1", "192.0.2.2", "password"); got != 429 {
|
||||
t.Fatal("lockout leaks password validity:", got)
|
||||
}
|
||||
// A successful login clears its IP bucket, so fat-fingers do not accumulate.
|
||||
if got := login("10.0.0.1:1", "192.0.2.3", "wrong"); got != 401 {
|
||||
t.Fatal(got)
|
||||
}
|
||||
if got := login("10.0.0.1:1", "192.0.2.3", "password"); got != 200 {
|
||||
t.Fatal(got)
|
||||
}
|
||||
for i := 0; i < loginFailuresPerMinute; i++ {
|
||||
if got := login("10.0.0.1:1", "192.0.2.3", "wrong"); got != 401 {
|
||||
t.Fatalf("post-clear failure %d: %d", i, got)
|
||||
}
|
||||
}
|
||||
if got := login("10.0.0.1:1", "192.0.2.3", "wrong"); got != 429 {
|
||||
t.Fatal(got)
|
||||
}
|
||||
for i := 0; i < loginGlobalFailuresPerMinute; i++ {
|
||||
login("10.0.0.1:1", fmt.Sprintf("198.18.%d.%d", i/256, i%256), "wrong")
|
||||
}
|
||||
if got := login("10.0.0.1:1", "203.0.113.1", "wrong"); got != 429 {
|
||||
t.Fatal(got)
|
||||
}
|
||||
// The global cap fails closed too: aggregate brute force locks every
|
||||
// login out without revealing candidate-password validity.
|
||||
if got := login("10.0.0.1:1", "203.0.113.1", "password"); got != 429 {
|
||||
t.Fatal("global lockout leaks password validity:", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAuthFailureGate(t *testing.T) {
|
||||
s := testServer(t)
|
||||
cookie := loginCookie(t, s)
|
||||
ch := integrationChannel(t, s, cookie, "abc", "openai-chat", "http://127.0.0.1:1")
|
||||
key := integrationCreateKey(t, s, cookie, "gated", ch)
|
||||
request := func(auth string) *httptest.ResponseRecorder {
|
||||
r := httptest.NewRequest("POST", "/v1/chat/completions", strings.NewReader(integrationRequest("openai-chat", "abc/gpt-5", false)))
|
||||
r.Header.Set("Authorization", "Bearer "+auth)
|
||||
w := httptest.NewRecorder()
|
||||
s.ServeHTTP(w, r)
|
||||
return w
|
||||
}
|
||||
if w := request(key.raw); w.Code != 502 {
|
||||
t.Fatal("valid key setup:", w.Code, w.Body.String())
|
||||
}
|
||||
for i := 0; i < authFailuresPerMinute; i++ {
|
||||
if w := request("wrong"); w.Code != 401 {
|
||||
t.Fatalf("%d: %d", i, w.Code)
|
||||
}
|
||||
}
|
||||
if w := request("wrong"); w.Code != 429 || w.Header().Get("Retry-After") == "" {
|
||||
t.Fatal(w.Code, w.Header())
|
||||
}
|
||||
// The gate applies before the key lookup: a valid key from the same
|
||||
// (attacker-shared) IP is throttled the same way.
|
||||
if w := request(key.raw); w.Code != 429 {
|
||||
t.Fatal(w.Code)
|
||||
}
|
||||
}
|
||||
|
||||
func TestProxyBodyTooLarge(t *testing.T) {
|
||||
s := testServer(t)
|
||||
cookie := loginCookie(t, s)
|
||||
ch := integrationChannel(t, s, cookie, "abc", "openai-chat", "http://127.0.0.1:1")
|
||||
key := integrationCreateKey(t, s, cookie, "limit", ch)
|
||||
r := httptest.NewRequest("POST", "/v1/chat/completions", strings.NewReader(`{"model":"abc/gpt-5","messages":[],"pad":"`+strings.Repeat("x", maxBody)+`"}`))
|
||||
r.Header.Set("Authorization", "Bearer "+key.raw)
|
||||
w := httptest.NewRecorder()
|
||||
s.ServeHTTP(w, r)
|
||||
if w.Code != 413 || !strings.Contains(w.Body.String(), "too large") {
|
||||
t.Fatal(w.Code, w.Body.String())
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,327 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/rand"
|
||||
"crypto/sha256"
|
||||
"crypto/subtle"
|
||||
"database/sql"
|
||||
"encoding/hex"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"io"
|
||||
"io/fs"
|
||||
_ "modernc.org/sqlite"
|
||||
"net"
|
||||
"net/http"
|
||||
"net/netip"
|
||||
"net/url"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
)
|
||||
|
||||
type Server struct {
|
||||
logRetentionDays atomic.Int64
|
||||
trustedProxies []netip.Prefix
|
||||
limiter *requestLimiter
|
||||
db *sql.DB
|
||||
password [32]byte
|
||||
assets fs.FS
|
||||
mu sync.Mutex
|
||||
sessions map[string]time.Time
|
||||
attempts map[string][]time.Time
|
||||
authFailures map[string][]time.Time
|
||||
authFailuresSeen time.Time
|
||||
loginFailures []time.Time
|
||||
client *http.Client
|
||||
publicOrigin string
|
||||
lifecycleMu sync.Mutex
|
||||
closing bool
|
||||
active sync.WaitGroup
|
||||
closeOnce sync.Once
|
||||
closeErr error
|
||||
logMu sync.RWMutex
|
||||
logsClosed bool
|
||||
logQueue chan logEntry
|
||||
logDone chan struct{}
|
||||
logErr error
|
||||
}
|
||||
|
||||
const maxBody = 16 << 20
|
||||
|
||||
func New(p, password string, a fs.FS) (*Server, error) {
|
||||
if password == "" {
|
||||
return nil, errors.New("empty admin password")
|
||||
}
|
||||
if p != ":memory:" {
|
||||
if e := os.MkdirAll(filepath.Dir(p), 0700); e != nil {
|
||||
return nil, e
|
||||
}
|
||||
}
|
||||
d, e := sql.Open("sqlite", p)
|
||||
if e != nil {
|
||||
return nil, e
|
||||
}
|
||||
d.SetMaxOpenConns(1)
|
||||
schema := `PRAGMA busy_timeout=5000; CREATE TABLE IF NOT EXISTS channels(id TEXT PRIMARY KEY,name TEXT NOT NULL UNIQUE,base_url TEXT NOT NULL,protocol TEXT NOT NULL,enabled INTEGER NOT NULL,passthrough INTEGER NOT NULL,models TEXT NOT NULL,api_key TEXT NOT NULL); CREATE TABLE IF NOT EXISTS keys(id TEXT PRIMARY KEY,name TEXT NOT NULL,prefix TEXT NOT NULL,hash TEXT NOT NULL UNIQUE,enabled INTEGER NOT NULL,all_channels INTEGER NOT NULL,channels TEXT NOT NULL,expires_at TEXT NOT NULL,created_at TEXT NOT NULL,revoked INTEGER NOT NULL DEFAULT 0); CREATE TABLE IF NOT EXISTS logs(id INTEGER PRIMARY KEY AUTOINCREMENT,created_at TEXT NOT NULL,key_id TEXT NOT NULL,key_name TEXT NOT NULL,channel TEXT NOT NULL,request_model TEXT NOT NULL,actual_model TEXT NOT NULL,downstream TEXT NOT NULL,upstream TEXT NOT NULL,status INTEGER NOT NULL,duration_ms INTEGER NOT NULL,stream INTEGER NOT NULL,input_tokens INTEGER,output_tokens INTEGER,cache_read_tokens INTEGER,cache_write_tokens INTEGER,usage_known INTEGER NOT NULL,error TEXT NOT NULL,raw_usage TEXT NOT NULL); CREATE INDEX IF NOT EXISTS logs_time ON logs(created_at); CREATE INDEX IF NOT EXISTS logs_key ON logs(key_id);`
|
||||
if _, e = d.Exec(schema); e != nil {
|
||||
d.Close()
|
||||
return nil, e
|
||||
}
|
||||
if p != ":memory:" {
|
||||
if e = os.Chmod(p, 0600); e != nil {
|
||||
d.Close()
|
||||
return nil, e
|
||||
}
|
||||
}
|
||||
tr := http.DefaultTransport.(*http.Transport).Clone()
|
||||
// Upstreams must not be able to turn compressed streams into proxy-side
|
||||
// decompression work; the router asks for identity encoding.
|
||||
tr.DisableCompression = true
|
||||
tr.ResponseHeaderTimeout = 60 * time.Second
|
||||
tr.DialContext = func(ctx context.Context, network, addr string) (net.Conn, error) {
|
||||
c, e := (&net.Dialer{Timeout: 15 * time.Second, KeepAlive: 30 * time.Second}).DialContext(ctx, network, addr)
|
||||
if e != nil {
|
||||
return nil, e
|
||||
}
|
||||
return &idleConn{Conn: c}, nil
|
||||
}
|
||||
s := &Server{db: d, password: sha256.Sum256([]byte(password)), assets: a, sessions: map[string]time.Time{}, attempts: map[string][]time.Time{}, authFailures: map[string][]time.Time{}, client: &http.Client{Transport: tr, CheckRedirect: func(*http.Request, []*http.Request) error { return http.ErrUseLastResponse }}, logQueue: make(chan logEntry, 1024), logDone: make(chan struct{})}
|
||||
go s.writeLogs()
|
||||
s.limiter = newRequestLimiter(120, 8, 64)
|
||||
return s, nil
|
||||
}
|
||||
func (s *Server) Close() error {
|
||||
s.closeOnce.Do(func() {
|
||||
s.lifecycleMu.Lock()
|
||||
s.closing = true
|
||||
s.lifecycleMu.Unlock()
|
||||
s.active.Wait()
|
||||
s.logMu.Lock()
|
||||
s.logsClosed = true
|
||||
close(s.logQueue)
|
||||
s.logMu.Unlock()
|
||||
<-s.logDone
|
||||
s.client.CloseIdleConnections()
|
||||
s.closeErr = errors.Join(s.logErr, s.db.Close())
|
||||
})
|
||||
return s.closeErr
|
||||
}
|
||||
func writeJSON(w http.ResponseWriter, v any) {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
json.NewEncoder(w).Encode(v)
|
||||
}
|
||||
func errJSON(w http.ResponseWriter, n int, m string) {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
w.WriteHeader(n)
|
||||
json.NewEncoder(w).Encode(map[string]string{"error": m})
|
||||
}
|
||||
func decode(w http.ResponseWriter, r *http.Request, v any) error {
|
||||
limit := int64(maxBody)
|
||||
if r.URL.Path == "/admin/api/login" {
|
||||
limit = 16 << 10
|
||||
}
|
||||
return readBody(w, r, limit, bodyReadTimeout, func() error {
|
||||
d := json.NewDecoder(r.Body)
|
||||
if e := d.Decode(v); e != nil {
|
||||
return e
|
||||
}
|
||||
var extra any
|
||||
if d.Decode(&extra) != io.EOF {
|
||||
return errors.New("trailing JSON")
|
||||
}
|
||||
return nil
|
||||
})
|
||||
}
|
||||
func randomID() string {
|
||||
b := make([]byte, 32)
|
||||
if _, e := rand.Read(b); e != nil {
|
||||
panic(e)
|
||||
}
|
||||
return hex.EncodeToString(b)
|
||||
}
|
||||
func hash(v string) string { h := sha256.Sum256([]byte(v)); return hex.EncodeToString(h[:]) }
|
||||
func (s *Server) session(r *http.Request) bool {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
c, e := r.Cookie("ai_router_session")
|
||||
if e != nil {
|
||||
return false
|
||||
}
|
||||
t, ok := s.sessions[c.Value]
|
||||
return ok && time.Now().Before(t)
|
||||
}
|
||||
func (s *Server) guard(r *http.Request) bool {
|
||||
if r.Header.Get("X-Requested-With") != "ai-router" {
|
||||
return false
|
||||
}
|
||||
o, e := url.Parse(r.Header.Get("Origin"))
|
||||
if s.publicOrigin != "" {
|
||||
// Origins are scheme/host case-insensitive; a mixed-case deployment
|
||||
// value must not lock the admin out.
|
||||
return e == nil && strings.EqualFold(r.Header.Get("Origin"), s.publicOrigin)
|
||||
}
|
||||
scheme := "http"
|
||||
if r.TLS != nil {
|
||||
scheme = "https"
|
||||
}
|
||||
return e == nil && o.Scheme == scheme && strings.EqualFold(o.Host, r.Host) && o.User == nil && o.RawQuery == "" && o.Fragment == "" && o.Path == ""
|
||||
}
|
||||
func (s *Server) secureCookie() bool {
|
||||
return len(s.publicOrigin) >= 8 && strings.EqualFold(s.publicOrigin[:8], "https://")
|
||||
}
|
||||
func (s *Server) ServeHTTP(w http.ResponseWriter, r *http.Request) {
|
||||
w = &deadlineWriter{ResponseWriter: w}
|
||||
_ = http.NewResponseController(w).SetReadDeadline(time.Now().Add(bodyReadTimeout))
|
||||
s.lifecycleMu.Lock()
|
||||
if s.closing {
|
||||
s.lifecycleMu.Unlock()
|
||||
errJSON(w, http.StatusServiceUnavailable, "server shutting down")
|
||||
return
|
||||
}
|
||||
s.active.Add(1)
|
||||
s.lifecycleMu.Unlock()
|
||||
defer s.active.Done()
|
||||
w.Header().Set("X-Content-Type-Options", "nosniff")
|
||||
w.Header().Set("Referrer-Policy", "no-referrer")
|
||||
w.Header().Set("X-Frame-Options", "DENY")
|
||||
w.Header().Set("Content-Security-Policy", "default-src 'none'; base-uri 'none'; frame-ancestors 'none'; form-action 'none'; connect-src 'self'; script-src 'self'; style-src 'self' 'unsafe-inline'; img-src 'self' data:")
|
||||
if strings.HasPrefix(r.URL.Path, "/admin/api/") || r.URL.Path == "/admin/api" {
|
||||
w.Header().Set("Cache-Control", "no-store")
|
||||
s.admin(w, r)
|
||||
return
|
||||
}
|
||||
if strings.HasPrefix(r.URL.Path, "/v1/") {
|
||||
s.proxy(w, r)
|
||||
return
|
||||
}
|
||||
if strings.HasPrefix(r.URL.Path, "/api/") {
|
||||
http.NotFound(w, r)
|
||||
return
|
||||
}
|
||||
if s.assets == nil || (r.Method != "GET" && r.Method != "HEAD") {
|
||||
http.NotFound(w, r)
|
||||
return
|
||||
}
|
||||
p := strings.TrimPrefix(r.URL.Path, "/")
|
||||
if p == "" {
|
||||
p = "index.html"
|
||||
}
|
||||
if st, e := fs.Stat(s.assets, p); e != nil || st.IsDir() {
|
||||
p = "index.html"
|
||||
}
|
||||
http.ServeFileFS(w, r, s.assets, p)
|
||||
}
|
||||
func (s *Server) admin(w http.ResponseWriter, r *http.Request) {
|
||||
if r.Method != "GET" && !s.guard(r) {
|
||||
errJSON(w, 403, "forbidden")
|
||||
return
|
||||
}
|
||||
if r.URL.Path == "/admin/api/login" && r.Method == "POST" {
|
||||
s.login(w, r)
|
||||
return
|
||||
}
|
||||
if r.URL.Path == "/admin/api/session" && r.Method == "GET" {
|
||||
writeJSON(w, map[string]bool{"authenticated": s.session(r)})
|
||||
return
|
||||
}
|
||||
if !s.session(r) {
|
||||
errJSON(w, 401, "unauthorized")
|
||||
return
|
||||
}
|
||||
if r.URL.Path == "/admin/api/logout" && r.Method == "POST" {
|
||||
c, _ := r.Cookie("ai_router_session")
|
||||
s.mu.Lock()
|
||||
delete(s.sessions, c.Value)
|
||||
s.mu.Unlock()
|
||||
http.SetCookie(w, &http.Cookie{Name: "ai_router_session", MaxAge: -1, Path: "/", HttpOnly: true, SameSite: http.SameSiteStrictMode, Secure: r.TLS != nil || s.secureCookie()})
|
||||
writeJSON(w, map[string]bool{"authenticated": false})
|
||||
return
|
||||
}
|
||||
switch {
|
||||
case r.URL.Path == "/admin/api/models/test":
|
||||
s.testModel(w, r)
|
||||
case r.URL.Path == "/admin/api/channels":
|
||||
s.channels(w, r)
|
||||
case r.URL.Path == "/admin/api/channels/discover-models":
|
||||
s.discoverModels(w, r)
|
||||
case strings.HasPrefix(r.URL.Path, "/admin/api/channels/"):
|
||||
s.channelID(w, r)
|
||||
case r.URL.Path == "/admin/api/keys":
|
||||
s.keys(w, r)
|
||||
case strings.HasPrefix(r.URL.Path, "/admin/api/keys/"):
|
||||
s.keyID(w, r)
|
||||
case r.URL.Path == "/admin/api/logs" && r.Method == "GET":
|
||||
s.logs(w, r)
|
||||
case r.URL.Path == "/admin/api/stats" && r.Method == "GET":
|
||||
s.stats(w, r)
|
||||
default:
|
||||
errJSON(w, 404, "not found")
|
||||
}
|
||||
}
|
||||
// login gates on failure counts BEFORE verifying the password. While an IP is
|
||||
// locked — or the global failure budget is spent — every attempt, including a
|
||||
// correct password, gets the same 429 without a comparison, so responses never
|
||||
// reveal whether a candidate password is right and brute force stays capped
|
||||
// at loginFailuresPerMinute per IP and loginGlobalFailuresPerMinute overall.
|
||||
// Successful logins are not failures: their gate slot is released by clearing
|
||||
// the IP bucket, and only wrong passwords feed the global window.
|
||||
func (s *Server) login(w http.ResponseWriter, r *http.Request) {
|
||||
ip := s.clientIP(r)
|
||||
now := time.Now()
|
||||
var x struct {
|
||||
Password string `json:"password"`
|
||||
}
|
||||
if decode(w, r, &x) != nil {
|
||||
errJSON(w, 400, "invalid request")
|
||||
return
|
||||
}
|
||||
s.mu.Lock()
|
||||
pruneFailures(s.attempts, now)
|
||||
s.loginFailures = trimRecent(s.loginFailures, now)
|
||||
recent := []time.Time{}
|
||||
for _, t := range s.attempts[ip] {
|
||||
if now.Sub(t) < time.Minute {
|
||||
recent = append(recent, t)
|
||||
}
|
||||
}
|
||||
blocked := len(recent) >= loginFailuresPerMinute || len(s.loginFailures) >= loginGlobalFailuresPerMinute || (len(recent) == 0 && len(s.attempts) >= 10000)
|
||||
if !blocked {
|
||||
// Reserve the gate slot atomically with the check so a concurrent
|
||||
// burst cannot slip past the limit; a success erases the bucket below.
|
||||
s.attempts[ip] = append(recent, now)
|
||||
}
|
||||
s.mu.Unlock()
|
||||
if blocked {
|
||||
w.Header().Set("Retry-After", "60")
|
||||
errJSON(w, 429, "too many attempts")
|
||||
return
|
||||
}
|
||||
h := sha256.Sum256([]byte(x.Password))
|
||||
if subtle.ConstantTimeCompare(h[:], s.password[:]) != 1 {
|
||||
s.mu.Lock()
|
||||
s.loginFailures = append(s.loginFailures, now)
|
||||
s.mu.Unlock()
|
||||
errJSON(w, 401, "invalid credentials")
|
||||
return
|
||||
}
|
||||
v := randomID()
|
||||
s.mu.Lock()
|
||||
delete(s.attempts, ip)
|
||||
for k, t := range s.sessions {
|
||||
if now.After(t) {
|
||||
delete(s.sessions, k)
|
||||
}
|
||||
}
|
||||
if c, e := r.Cookie("ai_router_session"); e == nil {
|
||||
delete(s.sessions, c.Value)
|
||||
}
|
||||
s.sessions[v] = now.Add(24 * time.Hour)
|
||||
s.mu.Unlock()
|
||||
http.SetCookie(w, &http.Cookie{Name: "ai_router_session", Value: v, Path: "/", MaxAge: 86400, HttpOnly: true, SameSite: http.SameSiteStrictMode, Secure: r.TLS != nil || s.secureCookie()})
|
||||
writeJSON(w, map[string]bool{"authenticated": true})
|
||||
}
|
||||
@@ -0,0 +1,340 @@
|
||||
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")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,306 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"ai-router/internal/protocol"
|
||||
"context"
|
||||
"database/sql"
|
||||
"database/sql/driver"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"sort"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
_ "time/tzdata"
|
||||
|
||||
"modernc.org/sqlite"
|
||||
)
|
||||
|
||||
func init() {
|
||||
var locationMu sync.Mutex
|
||||
var cachedZone string
|
||||
var cachedLocation *time.Location
|
||||
if err := sqlite.RegisterDeterministicScalarFunction("router_log_day", 2, func(_ *sqlite.FunctionContext, args []driver.Value) (driver.Value, error) {
|
||||
stampText, ok := args[0].(string)
|
||||
zone, zoneOK := args[1].(string)
|
||||
if !ok || !zoneOK {
|
||||
return nil, fmt.Errorf("invalid log timestamp or timezone")
|
||||
}
|
||||
stamp, err := time.Parse(time.RFC3339Nano, stampText)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
locationMu.Lock()
|
||||
defer locationMu.Unlock()
|
||||
if cachedLocation == nil || cachedZone != zone {
|
||||
loc, err := time.LoadLocation(zone)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
cachedZone, cachedLocation = zone, loc
|
||||
}
|
||||
return stamp.In(cachedLocation).Format("2006-01-02"), nil
|
||||
}); err != nil {
|
||||
panic(err)
|
||||
}
|
||||
}
|
||||
|
||||
type requestLog struct {
|
||||
ID int64 `json:"id"`
|
||||
CreatedAt string `json:"created_at"`
|
||||
KeyID string `json:"key_id"`
|
||||
KeyName string `json:"key_name"`
|
||||
Channel string `json:"channel"`
|
||||
RequestModel string `json:"request_model"`
|
||||
ActualModel string `json:"actual_model"`
|
||||
Downstream string `json:"downstream"`
|
||||
Upstream string `json:"upstream"`
|
||||
Status int `json:"status"`
|
||||
DurationMS int64 `json:"duration_ms"`
|
||||
Stream bool `json:"stream"`
|
||||
Input *int64 `json:"input_tokens"`
|
||||
Output *int64 `json:"output_tokens"`
|
||||
CacheRead *int64 `json:"cache_read_tokens"`
|
||||
CacheWrite *int64 `json:"cache_write_tokens"`
|
||||
UsageKnown bool `json:"usage_known"`
|
||||
Error string `json:"error"`
|
||||
RawUsage json.RawMessage `json:"raw_usage"`
|
||||
Usage protocol.Usage `json:"-"`
|
||||
}
|
||||
|
||||
func insertLog(ctx context.Context, tx *sql.Tx, l requestLog) error {
|
||||
l = boundedLog(l)
|
||||
var in, out, cr, cw any
|
||||
if l.Usage.Known || len(l.Usage.Raw) > 0 {
|
||||
in = l.Usage.Input
|
||||
out = l.Usage.Output
|
||||
cr = l.Usage.CacheRead
|
||||
cw = l.Usage.CacheWrite
|
||||
}
|
||||
_, err := tx.ExecContext(ctx, "INSERT INTO logs(created_at,key_id,key_name,channel,request_model,actual_model,downstream,upstream,status,duration_ms,stream,input_tokens,output_tokens,cache_read_tokens,cache_write_tokens,usage_known,error,raw_usage) VALUES(?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?)", l.CreatedAt, l.KeyID, l.KeyName, l.Channel, l.RequestModel, l.ActualModel, l.Downstream, l.Upstream, l.Status, l.DurationMS, l.Stream, in, out, cr, cw, l.Usage.Known, l.Error, string(l.Usage.Raw))
|
||||
return err
|
||||
}
|
||||
func filter(r *http.Request) (string, []any, error) {
|
||||
parts := []string{"1=1"}
|
||||
args := []any{}
|
||||
for _, name := range []string{"channel", "model"} {
|
||||
if v := r.URL.Query().Get(name); v != "" {
|
||||
col := name
|
||||
if name == "model" {
|
||||
col = "actual_model"
|
||||
}
|
||||
parts = append(parts, col+"=?")
|
||||
args = append(args, v)
|
||||
}
|
||||
}
|
||||
if k := r.URL.Query().Get("key_id"); k != "" {
|
||||
parts = append(parts, "key_id=?")
|
||||
args = append(args, k)
|
||||
}
|
||||
var from, to time.Time
|
||||
for _, n := range []string{"from", "to"} {
|
||||
if v := r.URL.Query().Get(n); v != "" {
|
||||
t, e := time.Parse(time.RFC3339Nano, v)
|
||||
if e != nil {
|
||||
return "", nil, e
|
||||
}
|
||||
if n == "from" {
|
||||
from = t
|
||||
parts = append(parts, "julianday(created_at)>=julianday(?)")
|
||||
} else {
|
||||
to = t
|
||||
parts = append(parts, "julianday(created_at)<=julianday(?)")
|
||||
}
|
||||
args = append(args, t.UTC().Format(time.RFC3339Nano))
|
||||
}
|
||||
}
|
||||
if !from.IsZero() && !to.IsZero() && from.After(to) {
|
||||
return "", nil, sql.ErrNoRows
|
||||
}
|
||||
return strings.Join(parts, " AND "), args, nil
|
||||
}
|
||||
func (s *Server) readLogs(r *http.Request, limited bool) ([]requestLog, error) {
|
||||
where, args, e := filter(r)
|
||||
if e != nil {
|
||||
return nil, e
|
||||
}
|
||||
q := "SELECT id,created_at,key_id,key_name,channel,request_model,actual_model,downstream,upstream,status,duration_ms,stream,input_tokens,output_tokens,cache_read_tokens,cache_write_tokens,usage_known,error,raw_usage FROM logs WHERE " + where + " ORDER BY id DESC"
|
||||
if limited {
|
||||
q += " LIMIT 200"
|
||||
}
|
||||
rows, e := s.db.Query(q, args...)
|
||||
if e != nil {
|
||||
return nil, e
|
||||
}
|
||||
defer rows.Close()
|
||||
out := []requestLog{}
|
||||
for rows.Next() {
|
||||
var l requestLog
|
||||
var raw string
|
||||
if e = rows.Scan(&l.ID, &l.CreatedAt, &l.KeyID, &l.KeyName, &l.Channel, &l.RequestModel, &l.ActualModel, &l.Downstream, &l.Upstream, &l.Status, &l.DurationMS, &l.Stream, &l.Input, &l.Output, &l.CacheRead, &l.CacheWrite, &l.UsageKnown, &l.Error, &raw); e != nil {
|
||||
return nil, e
|
||||
}
|
||||
if json.Valid([]byte(raw)) {
|
||||
l.RawUsage = json.RawMessage(raw)
|
||||
}
|
||||
out = append(out, l)
|
||||
}
|
||||
return out, rows.Err()
|
||||
}
|
||||
func (s *Server) logs(w http.ResponseWriter, r *http.Request) {
|
||||
if _, _, e := filter(r); e != nil {
|
||||
errJSON(w, 400, "invalid time filter")
|
||||
return
|
||||
}
|
||||
ls, e := s.readLogs(r, true)
|
||||
if e != nil {
|
||||
errJSON(w, 500, "database error")
|
||||
return
|
||||
}
|
||||
writeJSON(w, ls)
|
||||
}
|
||||
|
||||
type summary struct {
|
||||
Requests int64 `json:"requests"`
|
||||
Success int64 `json:"success"`
|
||||
Failed int64 `json:"failed"`
|
||||
Input *int64 `json:"input_tokens"`
|
||||
Output *int64 `json:"output_tokens"`
|
||||
CacheRead *int64 `json:"cache_read_tokens"`
|
||||
CacheWrite *int64 `json:"cache_write_tokens"`
|
||||
Unknown int64 `json:"unknown_usage"`
|
||||
}
|
||||
|
||||
func (a *summary) add(l requestLog) {
|
||||
a.Requests++
|
||||
if l.Status >= 200 && l.Status < 300 && l.Error == "" {
|
||||
a.Success++
|
||||
} else {
|
||||
a.Failed++
|
||||
}
|
||||
if !l.UsageKnown {
|
||||
a.Unknown++
|
||||
return
|
||||
}
|
||||
if a.Input == nil {
|
||||
a.Input = new(int64)
|
||||
a.Output = new(int64)
|
||||
a.CacheRead = new(int64)
|
||||
a.CacheWrite = new(int64)
|
||||
}
|
||||
if l.Input != nil {
|
||||
*a.Input += *l.Input
|
||||
}
|
||||
if l.Output != nil {
|
||||
*a.Output += *l.Output
|
||||
}
|
||||
if l.CacheRead != nil {
|
||||
*a.CacheRead += *l.CacheRead
|
||||
}
|
||||
if l.CacheWrite != nil {
|
||||
*a.CacheWrite += *l.CacheWrite
|
||||
}
|
||||
}
|
||||
func (s *Server) stats(w http.ResponseWriter, r *http.Request) {
|
||||
zone := r.URL.Query().Get("timezone")
|
||||
if zone == "" {
|
||||
zone = "UTC"
|
||||
}
|
||||
_, err := time.LoadLocation(zone)
|
||||
if err != nil || zone == "Local" {
|
||||
errJSON(w, 400, "invalid timezone")
|
||||
return
|
||||
}
|
||||
where, args, e := filter(r)
|
||||
if e != nil {
|
||||
errJSON(w, 400, "invalid time filter")
|
||||
return
|
||||
}
|
||||
const aggregates = `COUNT(*), COALESCE(SUM(status>=200 AND status<300 AND error=''),0),
|
||||
COALESCE(SUM(NOT(status>=200 AND status<300 AND error='')),0),
|
||||
SUM(CASE WHEN usage_known THEN COALESCE(input_tokens,0) END),
|
||||
SUM(CASE WHEN usage_known THEN COALESCE(output_tokens,0) END),
|
||||
SUM(CASE WHEN usage_known THEN COALESCE(cache_read_tokens,0) END),
|
||||
SUM(CASE WHEN usage_known THEN COALESCE(cache_write_tokens,0) END),
|
||||
COALESCE(SUM(NOT usage_known),0), MAX(id)`
|
||||
// NOT MATERIALIZED avoids copying the full log payload into a temporary table.
|
||||
query := `WITH filtered AS NOT MATERIALIZED (SELECT * FROM logs WHERE ` + where + `)
|
||||
SELECT 'total','','',` + aggregates + ` FROM filtered
|
||||
UNION ALL SELECT 'key',key_id,key_name,` + aggregates + ` FROM filtered GROUP BY key_id
|
||||
UNION ALL SELECT 'day',router_log_day(created_at,?), '',` + aggregates + ` FROM filtered GROUP BY router_log_day(created_at,?)
|
||||
UNION ALL SELECT 'model',channel,actual_model,` + aggregates + ` FROM filtered GROUP BY channel,actual_model`
|
||||
args = append(args, zone, zone)
|
||||
ctx, cancel := context.WithTimeout(r.Context(), 10*time.Second)
|
||||
defer cancel()
|
||||
rows, e := s.db.QueryContext(ctx, query, args...)
|
||||
if e != nil {
|
||||
errJSON(w, 500, "database error")
|
||||
return
|
||||
}
|
||||
defer rows.Close()
|
||||
type byKey struct {
|
||||
KeyID string `json:"key_id"`
|
||||
Name string `json:"name"`
|
||||
summary
|
||||
}
|
||||
type daily struct {
|
||||
Day string `json:"day"`
|
||||
summary
|
||||
}
|
||||
type byModel struct {
|
||||
Channel string `json:"channel"`
|
||||
Model string `json:"model"`
|
||||
summary
|
||||
}
|
||||
total := summary{}
|
||||
ks := map[string]*byKey{}
|
||||
ds := map[string]*daily{}
|
||||
ms := map[[2]string]*byModel{}
|
||||
groups := 0
|
||||
for rows.Next() {
|
||||
groups++
|
||||
if groups > 10000 {
|
||||
errJSON(w, 422, "too many statistics groups; narrow the time range")
|
||||
return
|
||||
}
|
||||
var kind, a, b string
|
||||
var v summary
|
||||
var latest sql.NullInt64
|
||||
if err := rows.Scan(&kind, &a, &b, &v.Requests, &v.Success, &v.Failed, &v.Input, &v.Output, &v.CacheRead, &v.CacheWrite, &v.Unknown, &latest); err != nil {
|
||||
errJSON(w, 500, "database error")
|
||||
return
|
||||
}
|
||||
switch kind {
|
||||
case "total":
|
||||
total = v
|
||||
case "key":
|
||||
// SQLite's sole MAX selects key_name from the latest ID, not timestamp.
|
||||
ks[a] = &byKey{KeyID: a, Name: b, summary: v}
|
||||
case "day":
|
||||
ds[a] = &daily{Day: a, summary: v}
|
||||
case "model":
|
||||
ms[[2]string{a, b}] = &byModel{Channel: a, Model: b, summary: v}
|
||||
}
|
||||
}
|
||||
if err := rows.Err(); err != nil {
|
||||
errJSON(w, 500, "database error")
|
||||
return
|
||||
}
|
||||
keys := []*byKey{}
|
||||
days := []*daily{}
|
||||
mods := []*byModel{}
|
||||
for _, v := range ks {
|
||||
keys = append(keys, v)
|
||||
}
|
||||
for _, v := range ds {
|
||||
days = append(days, v)
|
||||
}
|
||||
for _, v := range ms {
|
||||
mods = append(mods, v)
|
||||
}
|
||||
sort.Slice(keys, func(i, j int) bool { return keys[i].KeyID < keys[j].KeyID })
|
||||
sort.Slice(days, func(i, j int) bool { return days[i].Day < days[j].Day })
|
||||
sort.Slice(mods, func(i, j int) bool { return mods[i].Channel+mods[i].Model < mods[j].Channel+mods[j].Model })
|
||||
writeJSON(w, map[string]any{"summary": total, "by_key": keys, "daily": days, "by_model": mods, "timezone": zone})
|
||||
}
|
||||
@@ -0,0 +1,74 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"ai-router/internal/protocol"
|
||||
"encoding/json"
|
||||
"net/url"
|
||||
"reflect"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestStatsBrowserTimezone(t *testing.T) {
|
||||
for _, tc := range []struct {
|
||||
name, zone string
|
||||
stamps []string
|
||||
want map[string]int
|
||||
}{
|
||||
{"Shanghai midnight", "Asia/Shanghai", []string{"2026-10-20T15:59:59Z", "2026-10-20T16:00:00Z"}, map[string]int{"2026-10-20": 1, "2026-10-21": 1}},
|
||||
{"UTC default", "", []string{"2026-10-20T15:59:59Z", "2026-10-20T16:00:00Z"}, map[string]int{"2026-10-20": 2}},
|
||||
{"New York DST spring", "America/New_York", []string{"2025-03-09T04:59:59Z", "2025-03-09T05:00:00Z", "2025-03-10T03:59:59Z", "2025-03-10T04:00:00Z"}, map[string]int{"2025-03-08": 1, "2025-03-09": 2, "2025-03-10": 1}},
|
||||
{"New York DST fall", "America/New_York", []string{"2025-11-02T03:59:59Z", "2025-11-02T04:00:00Z", "2025-11-03T04:59:59Z", "2025-11-03T05:00:00Z"}, map[string]int{"2025-11-01": 1, "2025-11-02": 2, "2025-11-03": 1}},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
s := testServer(t)
|
||||
cookie := loginCookie(t, s)
|
||||
for _, stamp := range tc.stamps {
|
||||
s.record(requestLog{CreatedAt: stamp, KeyID: "one", Status: 200, Usage: protocol.Usage{Known: true, Input: 10}})
|
||||
}
|
||||
s.record(requestLog{CreatedAt: tc.stamps[0], KeyID: "other", Status: 200})
|
||||
if err := s.flushLogs(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
w := call(s, "GET", "/admin/api/stats?key_id=one&timezone="+url.QueryEscape(tc.zone), "", cookie)
|
||||
if w.Code != 200 {
|
||||
t.Fatal(w.Code, w.Body.String())
|
||||
}
|
||||
var result struct {
|
||||
Timezone string `json:"timezone"`
|
||||
Summary summary `json:"summary"`
|
||||
Daily []struct {
|
||||
Day string `json:"day"`
|
||||
Requests int `json:"requests"`
|
||||
} `json:"daily"`
|
||||
}
|
||||
if err := json.Unmarshal(w.Body.Bytes(), &result); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
got := map[string]int{}
|
||||
for _, d := range result.Daily {
|
||||
got[d.Day] = d.Requests
|
||||
}
|
||||
if !reflect.DeepEqual(got, tc.want) {
|
||||
t.Fatalf("got %v want %v", got, tc.want)
|
||||
}
|
||||
if result.Summary.Requests != int64(len(tc.stamps)) || result.Summary.Input == nil || *result.Summary.Input != int64(10*len(tc.stamps)) {
|
||||
t.Fatal(result.Summary)
|
||||
}
|
||||
var stored string
|
||||
if err := s.db.QueryRow("SELECT created_at FROM logs ORDER BY id LIMIT 1").Scan(&stored); err != nil || stored != tc.stamps[0] {
|
||||
t.Fatal("UTC storage changed", stored, err)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestStatsRejectsInvalidTimezone(t *testing.T) {
|
||||
s := testServer(t)
|
||||
c := loginCookie(t, s)
|
||||
for _, zone := range []string{"Local", "not/a-zone", "../../etc/passwd"} {
|
||||
w := call(s, "GET", "/admin/api/stats?timezone="+url.QueryEscape(zone), "", c)
|
||||
if w.Code != 400 {
|
||||
t.Fatal(zone, w.Code)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,33 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"net"
|
||||
"net/url"
|
||||
"time"
|
||||
)
|
||||
|
||||
type idleConn struct{ net.Conn }
|
||||
|
||||
func (c *idleConn) Read(b []byte) (int, error) {
|
||||
if e := c.SetReadDeadline(time.Now().Add(2 * time.Minute)); e != nil {
|
||||
return 0, e
|
||||
}
|
||||
return c.Conn.Read(b)
|
||||
}
|
||||
func (c *idleConn) Write(b []byte) (int, error) {
|
||||
if e := c.SetWriteDeadline(time.Now().Add(30 * time.Second)); e != nil {
|
||||
return 0, e
|
||||
}
|
||||
return c.Conn.Write(b)
|
||||
}
|
||||
|
||||
// SetPublicOrigin must be called before serving requests; forwarded headers are not trusted.
|
||||
func (s *Server) SetPublicOrigin(origin string) error {
|
||||
u, e := url.Parse(origin)
|
||||
if e != nil || (u.Scheme != "http" && u.Scheme != "https") || u.Host == "" || u.User != nil || u.Path != "" || u.RawQuery != "" || u.ForceQuery || u.Fragment != "" {
|
||||
return errors.New("public origin must be an http(s) origin")
|
||||
}
|
||||
s.publicOrigin = origin
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,51 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
type observedDeadlineWriter struct {
|
||||
*httptest.ResponseRecorder
|
||||
writes, reads int
|
||||
deadline time.Time
|
||||
}
|
||||
|
||||
func (w *observedDeadlineWriter) SetWriteDeadline(d time.Time) error {
|
||||
w.writes++
|
||||
w.deadline = d
|
||||
return nil
|
||||
}
|
||||
func (w *observedDeadlineWriter) SetReadDeadline(d time.Time) error { w.reads++; return nil }
|
||||
|
||||
func TestResponseWriteDeadlines(t *testing.T) {
|
||||
raw := &observedDeadlineWriter{ResponseRecorder: httptest.NewRecorder()}
|
||||
w := &deadlineWriter{ResponseWriter: raw}
|
||||
w.WriteHeader(200)
|
||||
if raw.writes != 1 || !raw.deadline.After(time.Now()) {
|
||||
t.Fatal("headers lack deadline")
|
||||
}
|
||||
if _, err := w.Write([]byte("body")); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := w.FlushError(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if raw.writes != 3 {
|
||||
t.Fatal("write/flush must refresh deadline", raw.writes)
|
||||
}
|
||||
if err := http.NewResponseController(w).SetReadDeadline(time.Time{}); err != nil || raw.reads != 1 {
|
||||
t.Fatal("wrapper must preserve read deadline control", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestNonStreamResponseHasWriteDeadline(t *testing.T) {
|
||||
s := testServer(t)
|
||||
raw := &observedDeadlineWriter{ResponseRecorder: httptest.NewRecorder()}
|
||||
s.ServeHTTP(raw, httptest.NewRequest("GET", "/admin/api/session", nil))
|
||||
if raw.Code != 200 || raw.writes == 0 || !raw.deadline.After(time.Now()) {
|
||||
t.Fatal("nonstream response lacks write deadline")
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user