Benson Wong
62aea0e83d
- refactor shared http functionality into internal/shared/http.go - remove stripping of Authorization and x-api-key - add Request Context middleware to internal/server - add /ui and /metrics behind auth middleware, fixes #717 Fix #717 Updates: #834
427 lines
12 KiB
Go
427 lines
12 KiB
Go
package shared
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"encoding/base64"
|
|
"io"
|
|
"mime/multipart"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"net/url"
|
|
"strings"
|
|
"testing"
|
|
)
|
|
|
|
func TestExtractContext_GET(t *testing.T) {
|
|
tests := []struct {
|
|
name string
|
|
query string
|
|
wantModel string
|
|
wantErr bool
|
|
}{
|
|
{"model present", "model=llama3", "llama3", false},
|
|
{"model with slashes", "model=author/model-7b", "author/model-7b", false},
|
|
{"model missing", "", "", false},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
r, _ := http.NewRequest(http.MethodGet, "/?"+tt.query, nil)
|
|
got, err := extractContext(r)
|
|
if (err != nil) != tt.wantErr {
|
|
t.Fatalf("wantErr=%v got err=%v", tt.wantErr, err)
|
|
}
|
|
if got.Model != tt.wantModel {
|
|
t.Errorf("want %q got %q", tt.wantModel, got.Model)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestExtractContext_JSON(t *testing.T) {
|
|
tests := []struct {
|
|
name string
|
|
body string
|
|
wantModel string
|
|
wantErr bool
|
|
}{
|
|
{"model present", `{"model":"llama3","stream":true}`, "llama3", false},
|
|
{"model with slashes", `{"model":"author/model-7b"}`, "author/model-7b", false},
|
|
{"model empty string", `{"model":""}`, "", false},
|
|
{"model key missing", `{"stream":true}`, "", false},
|
|
{"invalid json", `not-json`, "", false},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
r, _ := http.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(tt.body))
|
|
r.Header.Set("Content-Type", "application/json")
|
|
got, err := extractContext(r)
|
|
if (err != nil) != tt.wantErr {
|
|
t.Fatalf("wantErr=%v got err=%v", tt.wantErr, err)
|
|
}
|
|
if got.Model != tt.wantModel {
|
|
t.Errorf("want %q got %q", tt.wantModel, got.Model)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestExtractContext_URLEncodedForm(t *testing.T) {
|
|
tests := []struct {
|
|
name string
|
|
formModel string
|
|
wantModel string
|
|
wantErr bool
|
|
}{
|
|
{"model present", "whisper-1", "whisper-1", false},
|
|
{"model missing", "", "", false},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
form := url.Values{}
|
|
if tt.formModel != "" {
|
|
form.Set("model", tt.formModel)
|
|
}
|
|
r, _ := http.NewRequest(http.MethodPost, "/v1/audio/transcriptions", strings.NewReader(form.Encode()))
|
|
r.Header.Set("Content-Type", "application/x-www-form-urlencoded")
|
|
got, err := extractContext(r)
|
|
if (err != nil) != tt.wantErr {
|
|
t.Fatalf("wantErr=%v got err=%v", tt.wantErr, err)
|
|
}
|
|
if got.Model != tt.wantModel {
|
|
t.Errorf("want %q got %q", tt.wantModel, got.Model)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestExtractContext_MultipartForm(t *testing.T) {
|
|
tests := []struct {
|
|
name string
|
|
formModel string
|
|
wantModel string
|
|
wantErr bool
|
|
}{
|
|
{"model present", "whisper-1", "whisper-1", false},
|
|
{"model missing", "", "", false},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
var buf bytes.Buffer
|
|
mw := multipart.NewWriter(&buf)
|
|
if tt.formModel != "" {
|
|
fw, _ := mw.CreateFormField("model")
|
|
fw.Write([]byte(tt.formModel))
|
|
}
|
|
mw.Close()
|
|
|
|
r, _ := http.NewRequest(http.MethodPost, "/v1/audio/transcriptions", &buf)
|
|
r.Header.Set("Content-Type", mw.FormDataContentType())
|
|
got, err := extractContext(r)
|
|
if (err != nil) != tt.wantErr {
|
|
t.Fatalf("wantErr=%v got err=%v", tt.wantErr, err)
|
|
}
|
|
if got.Model != tt.wantModel {
|
|
t.Errorf("want %q got %q", tt.wantModel, got.Model)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestExtractContext_JSONBodyRestored(t *testing.T) {
|
|
body := `{"model":"llama3","stream":true}`
|
|
r, _ := http.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(body))
|
|
r.Header.Set("Content-Type", "application/json")
|
|
|
|
if _, err := extractContext(r); err != nil {
|
|
t.Fatalf("ExtractContext: %v", err)
|
|
}
|
|
|
|
remaining, err := io.ReadAll(r.Body)
|
|
if err != nil {
|
|
t.Fatalf("reading body after ExtractContext: %v", err)
|
|
}
|
|
if string(remaining) != body {
|
|
t.Errorf("body not restored: want %q got %q", body, string(remaining))
|
|
}
|
|
}
|
|
|
|
func TestExtractContext_MultipartBodyRestored(t *testing.T) {
|
|
var buf bytes.Buffer
|
|
mw := multipart.NewWriter(&buf)
|
|
fw, _ := mw.CreateFormField("model")
|
|
fw.Write([]byte("whisper-1"))
|
|
ff, _ := mw.CreateFormFile("file", "audio.wav")
|
|
ff.Write([]byte("fake-audio-bytes"))
|
|
mw.Close()
|
|
|
|
original := buf.Bytes()
|
|
|
|
r, _ := http.NewRequest(http.MethodPost, "/v1/audio/transcriptions", bytes.NewReader(original))
|
|
r.Header.Set("Content-Type", mw.FormDataContentType())
|
|
|
|
if _, err := extractContext(r); err != nil {
|
|
t.Fatalf("ExtractContext: %v", err)
|
|
}
|
|
|
|
remaining, err := io.ReadAll(r.Body)
|
|
if err != nil {
|
|
t.Fatalf("reading body after ExtractContext: %v", err)
|
|
}
|
|
if !bytes.Equal(remaining, original) {
|
|
t.Errorf("multipart body not restored: want %d bytes got %d bytes", len(original), len(remaining))
|
|
}
|
|
}
|
|
|
|
func TestExtractContext_URLEncodedBodyRestored(t *testing.T) {
|
|
body := "model=whisper-1&extra=value"
|
|
r, _ := http.NewRequest(http.MethodPost, "/v1/audio/transcriptions", strings.NewReader(body))
|
|
r.Header.Set("Content-Type", "application/x-www-form-urlencoded")
|
|
|
|
if _, err := extractContext(r); err != nil {
|
|
t.Fatalf("ExtractContext: %v", err)
|
|
}
|
|
|
|
remaining, err := io.ReadAll(r.Body)
|
|
if err != nil {
|
|
t.Fatalf("reading body after ExtractContext: %v", err)
|
|
}
|
|
if string(remaining) != body {
|
|
t.Errorf("url-encoded body not restored: want %q got %q", body, string(remaining))
|
|
}
|
|
}
|
|
|
|
func TestSetContext(t *testing.T) {
|
|
ctx := SetContext(context.Background(), ReqContextData{Model: "llama3", ModelID: "llama3"})
|
|
data, ok := ctx.Value(ReqContextKey).(ReqContextData)
|
|
if !ok {
|
|
t.Fatalf("ContextKey not set or wrong type")
|
|
}
|
|
if data.Model != "llama3" {
|
|
t.Errorf("want %q got %q", "llama3", data.Model)
|
|
}
|
|
if data.ModelID != "llama3" {
|
|
t.Errorf("want %q got %q", "llama3", data.ModelID)
|
|
}
|
|
}
|
|
|
|
func TestSetContext_WithAlias(t *testing.T) {
|
|
ctx := SetContext(context.Background(), ReqContextData{Model: "llama", ModelID: "llama3"})
|
|
data, _ := ctx.Value(ReqContextKey).(ReqContextData)
|
|
if data.Model != "llama" {
|
|
t.Errorf("want requested %q got %q", "llama", data.Model)
|
|
}
|
|
if data.ModelID != "llama3" {
|
|
t.Errorf("want real %q got %q", "llama3", data.ModelID)
|
|
}
|
|
}
|
|
|
|
func TestSetContext_DoesNotMutateParent(t *testing.T) {
|
|
parent := context.Background()
|
|
_ = SetContext(parent, ReqContextData{Model: "llama3", ModelID: "llama3"})
|
|
if v := parent.Value(ReqContextKey); v != nil {
|
|
t.Errorf("parent context was mutated: %v", v)
|
|
}
|
|
}
|
|
|
|
func TestReadContext(t *testing.T) {
|
|
tests := []struct {
|
|
name string
|
|
ctx context.Context
|
|
wantReq string
|
|
wantReal string
|
|
wantBool bool
|
|
}{
|
|
{
|
|
name: "model present, same name",
|
|
ctx: SetContext(context.Background(), ReqContextData{Model: "llama3", ModelID: "llama3"}),
|
|
wantReq: "llama3",
|
|
wantReal: "llama3",
|
|
wantBool: true,
|
|
},
|
|
{
|
|
name: "model present, aliased",
|
|
ctx: SetContext(context.Background(), ReqContextData{Model: "llama", ModelID: "llama3"}),
|
|
wantReq: "llama",
|
|
wantReal: "llama3",
|
|
wantBool: true,
|
|
},
|
|
{
|
|
name: "model absent",
|
|
ctx: context.Background(),
|
|
wantReq: "",
|
|
wantReal: "",
|
|
wantBool: false,
|
|
},
|
|
{
|
|
name: "model is empty string",
|
|
ctx: SetContext(context.Background(), ReqContextData{Model: "", ModelID: ""}),
|
|
wantReq: "",
|
|
wantReal: "",
|
|
wantBool: true,
|
|
},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
gotData, ok := ReadContext(tt.ctx)
|
|
if gotData.Model != tt.wantReq || gotData.ModelID != tt.wantReal || ok != tt.wantBool {
|
|
t.Errorf("want (%q, %q, %v) got (%q, %q, %v)", tt.wantReq, tt.wantReal, tt.wantBool, gotData.Model, gotData.ModelID, ok)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestExtractContext_Streaming_GET(t *testing.T) {
|
|
tests := []struct {
|
|
name string
|
|
query string
|
|
wantStreaming bool
|
|
}{
|
|
{"streaming true", "model=llama3&stream=true", true},
|
|
{"streaming false", "model=llama3&stream=false", false},
|
|
{"no stream param", "model=llama3", false},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
r, _ := http.NewRequest(http.MethodGet, "/?"+tt.query, nil)
|
|
got, err := extractContext(r)
|
|
if err != nil {
|
|
t.Fatalf("ExtractContext: %v", err)
|
|
}
|
|
if got.Streaming != tt.wantStreaming {
|
|
t.Errorf("Streaming: want %v, got %v", tt.wantStreaming, got.Streaming)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestExtractContext_Streaming_JSON(t *testing.T) {
|
|
tests := []struct {
|
|
name string
|
|
body string
|
|
wantStreaming bool
|
|
}{
|
|
{"streaming true", `{"model":"llama3","stream":true}`, true},
|
|
{"streaming false", `{"model":"llama3","stream":false}`, false},
|
|
{"no stream param", `{"model":"llama3"}`, false},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
r, _ := http.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(tt.body))
|
|
r.Header.Set("Content-Type", "application/json")
|
|
got, err := extractContext(r)
|
|
if err != nil {
|
|
t.Fatalf("ExtractContext: %v", err)
|
|
}
|
|
if got.Streaming != tt.wantStreaming {
|
|
t.Errorf("Streaming: want %v, got %v", tt.wantStreaming, got.Streaming)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestExtractContext_Streaming_URLEncodedForm(t *testing.T) {
|
|
r, _ := http.NewRequest(http.MethodPost, "/v1/audio/transcriptions", strings.NewReader("model=whisper-1&stream=true"))
|
|
r.Header.Set("Content-Type", "application/x-www-form-urlencoded")
|
|
got, err := extractContext(r)
|
|
if err != nil {
|
|
t.Fatalf("ExtractContext: %v", err)
|
|
}
|
|
if !got.Streaming {
|
|
t.Error("Streaming should be true")
|
|
}
|
|
}
|
|
|
|
func TestExtractContext_ApiKey(t *testing.T) {
|
|
basicHeader := func(user, pass string) string {
|
|
return "Basic " + base64.StdEncoding.EncodeToString([]byte(user+":"+pass))
|
|
}
|
|
cases := []struct {
|
|
name string
|
|
method string
|
|
ct string
|
|
body string
|
|
auth string
|
|
xapi string
|
|
wantKey string
|
|
}{
|
|
{"GET bearer", http.MethodGet, "", "", "Bearer sk-get", "", "sk-get"},
|
|
{"GET x-api-key", http.MethodGet, "", "", "", "xk-get", "xk-get"},
|
|
{"GET basic", http.MethodGet, "", "", basicHeader("u", "pw-get"), "", "pw-get"},
|
|
{"JSON bearer", http.MethodPost, "application/json", `{"model":"m"}`, "Bearer sk-json", "", "sk-json"},
|
|
{"JSON x-api-key", http.MethodPost, "application/json", `{"model":"m"}`, "", "xk-json", "xk-json"},
|
|
{"form bearer", http.MethodPost, "application/x-www-form-urlencoded", "model=m", "Bearer sk-form", "", "sk-form"},
|
|
{"no key", http.MethodGet, "", "", "", "", ""},
|
|
}
|
|
for _, c := range cases {
|
|
t.Run(c.name, func(t *testing.T) {
|
|
var body io.Reader
|
|
if c.body != "" {
|
|
body = strings.NewReader(c.body)
|
|
}
|
|
r, _ := http.NewRequest(c.method, "/", body)
|
|
if c.ct != "" {
|
|
r.Header.Set("Content-Type", c.ct)
|
|
}
|
|
if c.auth != "" {
|
|
r.Header.Set("Authorization", c.auth)
|
|
}
|
|
if c.xapi != "" {
|
|
r.Header.Set("x-api-key", c.xapi)
|
|
}
|
|
got, err := extractContext(r)
|
|
if err != nil {
|
|
t.Fatalf("extractContext: %v", err)
|
|
}
|
|
if got.ApiKey != c.wantKey {
|
|
t.Errorf("ApiKey = %q, want %q", got.ApiKey, c.wantKey)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestServer_ExtractAPIKey(t *testing.T) {
|
|
basicHeader := func(user, pass string) string {
|
|
return "Basic " + base64.StdEncoding.EncodeToString([]byte(user+":"+pass))
|
|
}
|
|
cases := []struct {
|
|
name string
|
|
auth string
|
|
xapi string
|
|
want string
|
|
}{
|
|
{"none", "", "", ""},
|
|
{"bearer", "Bearer tok123", "", "tok123"},
|
|
{"basic", basicHeader("user", "pw-key"), "", "pw-key"},
|
|
{"x-api-key", "", "xkey", "xkey"},
|
|
{"basic beats bearer", basicHeader("u", "bk"), "", "bk"},
|
|
{"bearer beats x-api-key", "Bearer btok", "xkey", "btok"},
|
|
{"malformed basic falls back to x-api-key", "Basic !!!notbase64", "xkey", "xkey"},
|
|
{"lowercase bearer", "bearer tok123", "", "tok123"},
|
|
{"lowercase basic", "basic " + base64.StdEncoding.EncodeToString([]byte("user:pw-key")), "", "pw-key"},
|
|
{"mixed case BEARER", "BEARER tok456", "", "tok456"},
|
|
{"mixed case bAsIc", "bAsIc " + base64.StdEncoding.EncodeToString([]byte("u:bk")), "", "bk"},
|
|
}
|
|
for _, c := range cases {
|
|
t.Run(c.name, func(t *testing.T) {
|
|
r := httptest.NewRequest(http.MethodGet, "/", nil)
|
|
if c.auth != "" {
|
|
r.Header.Set("Authorization", c.auth)
|
|
}
|
|
if c.xapi != "" {
|
|
r.Header.Set("x-api-key", c.xapi)
|
|
}
|
|
if got := ExtractAPIKey(r); got != c.want {
|
|
t.Errorf("extractAPIKey() = %q, want %q", got, c.want)
|
|
}
|
|
})
|
|
}
|
|
}
|