2026-08-19 20:52:13 -07:00
|
|
|
package proxy
|
|
|
|
|
|
|
|
|
|
import (
|
2026-09-27 11:16:40 +00:00
|
|
|
"crypto/rand"
|
|
|
|
|
"crypto/rsa"
|
|
|
|
|
"encoding/base64"
|
|
|
|
|
"encoding/json"
|
2026-08-19 20:52:13 -07:00
|
|
|
"fmt"
|
|
|
|
|
"io"
|
2026-09-27 11:16:40 +00:00
|
|
|
"math/big"
|
2026-08-19 20:52:13 -07:00
|
|
|
"net/http"
|
|
|
|
|
"net/http/httptest"
|
2026-09-27 11:16:40 +00:00
|
|
|
"runtime"
|
2026-08-19 20:52:13 -07:00
|
|
|
"strings"
|
|
|
|
|
"testing"
|
|
|
|
|
"time"
|
|
|
|
|
|
2026-09-27 11:16:40 +00:00
|
|
|
"github.com/golang-jwt/jwt/v5"
|
|
|
|
|
|
2026-08-21 20:46:42 -07:00
|
|
|
"forgejo.riotpiao.com/rock/homelab-frontend/internal/config"
|
2026-09-27 11:16:40 +00:00
|
|
|
"forgejo.riotpiao.com/rock/homelab-frontend/internal/observability"
|
2026-08-19 20:52:13 -07:00
|
|
|
)
|
|
|
|
|
|
|
|
|
|
func TestProxyBasic(t *testing.T) {
|
|
|
|
|
// Start a stub upstream
|
|
|
|
|
upstreamServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
|
|
|
w.Header().Set("X-Upstream-Header", "test-value")
|
|
|
|
|
w.WriteHeader(http.StatusOK)
|
|
|
|
|
fmt.Fprintf(w, "upstream response")
|
|
|
|
|
}))
|
|
|
|
|
defer upstreamServer.Close()
|
|
|
|
|
|
|
|
|
|
// Extract host:port from upstream URL
|
|
|
|
|
upstreamAddr := strings.TrimPrefix(upstreamServer.URL, "http://")
|
|
|
|
|
|
|
|
|
|
// Create config with route to the stub
|
|
|
|
|
cfg := &config.Config{
|
|
|
|
|
Routes: map[string]*config.Route{
|
|
|
|
|
"test-route": {
|
|
|
|
|
Name: "test-route",
|
|
|
|
|
Upstream: config.Upstream{
|
|
|
|
|
Address: upstreamAddr,
|
|
|
|
|
PathRewrite: "",
|
|
|
|
|
ConnectTimeout: 5 * time.Second,
|
|
|
|
|
ReadTimeout: 5 * time.Second,
|
|
|
|
|
WriteTimeout: 5 * time.Second,
|
|
|
|
|
MaxBodySize: 1024 * 1024,
|
|
|
|
|
AuthRequired: false,
|
|
|
|
|
},
|
|
|
|
|
},
|
|
|
|
|
},
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
handler := New(cfg)
|
|
|
|
|
defer handler.Close()
|
|
|
|
|
|
|
|
|
|
// Make a request through the proxy
|
|
|
|
|
server := httptest.NewServer(handler)
|
|
|
|
|
defer server.Close()
|
|
|
|
|
|
|
|
|
|
resp, err := http.Get(server.URL + "/test/path")
|
|
|
|
|
if err != nil {
|
|
|
|
|
t.Fatalf("request failed: %v", err)
|
|
|
|
|
}
|
|
|
|
|
defer resp.Body.Close()
|
|
|
|
|
|
|
|
|
|
// Verify status code
|
|
|
|
|
if resp.StatusCode != http.StatusOK {
|
|
|
|
|
t.Errorf("expected status 200, got %d", resp.StatusCode)
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
// Verify response header passed through
|
|
|
|
|
if resp.Header.Get("X-Upstream-Header") != "test-value" {
|
|
|
|
|
t.Errorf("upstream header not passed through")
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
// Verify response body
|
|
|
|
|
body, _ := io.ReadAll(resp.Body)
|
|
|
|
|
if string(body) != "upstream response" {
|
|
|
|
|
t.Errorf("expected body 'upstream response', got %s", string(body))
|
|
|
|
|
}
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
func TestProxyPathRewrite(t *testing.T) {
|
|
|
|
|
requestedPath := ""
|
|
|
|
|
upstreamServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
|
|
|
requestedPath = r.URL.Path
|
|
|
|
|
w.WriteHeader(http.StatusOK)
|
|
|
|
|
}))
|
|
|
|
|
defer upstreamServer.Close()
|
|
|
|
|
|
|
|
|
|
upstreamAddr := strings.TrimPrefix(upstreamServer.URL, "http://")
|
|
|
|
|
|
|
|
|
|
cfg := &config.Config{
|
|
|
|
|
Routes: map[string]*config.Route{
|
|
|
|
|
"rewrite-route": {
|
|
|
|
|
Name: "rewrite-route",
|
|
|
|
|
Upstream: config.Upstream{
|
|
|
|
|
Address: upstreamAddr,
|
|
|
|
|
PathRewrite: "/api/v2",
|
|
|
|
|
ConnectTimeout: 5 * time.Second,
|
|
|
|
|
ReadTimeout: 5 * time.Second,
|
|
|
|
|
WriteTimeout: 5 * time.Second,
|
|
|
|
|
MaxBodySize: 1024 * 1024,
|
|
|
|
|
AuthRequired: false,
|
|
|
|
|
},
|
|
|
|
|
},
|
|
|
|
|
},
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
handler := New(cfg)
|
|
|
|
|
defer handler.Close()
|
|
|
|
|
|
|
|
|
|
server := httptest.NewServer(handler)
|
|
|
|
|
defer server.Close()
|
|
|
|
|
|
2026-08-19 23:55:42 -07:00
|
|
|
// Not /v1/models: ServeHTTP serves that endpoint from config (task 2.5)
|
|
|
|
|
// and returns before routing, so it never exercises PathRewrite.
|
|
|
|
|
resp, err := http.Get(server.URL + "/some/path")
|
2026-08-19 20:52:13 -07:00
|
|
|
if err != nil {
|
|
|
|
|
t.Fatalf("request failed: %v", err)
|
|
|
|
|
}
|
|
|
|
|
resp.Body.Close()
|
|
|
|
|
|
|
|
|
|
if requestedPath != "/api/v2" {
|
|
|
|
|
t.Errorf("expected rewritten path /api/v2, got %s", requestedPath)
|
|
|
|
|
}
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
func TestProxyConnectionReuse(t *testing.T) {
|
|
|
|
|
upstreamServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
|
|
|
w.WriteHeader(http.StatusOK)
|
|
|
|
|
fmt.Fprintf(w, "response")
|
|
|
|
|
}))
|
|
|
|
|
defer upstreamServer.Close()
|
|
|
|
|
|
|
|
|
|
upstreamAddr := strings.TrimPrefix(upstreamServer.URL, "http://")
|
|
|
|
|
|
|
|
|
|
cfg := &config.Config{
|
|
|
|
|
Routes: map[string]*config.Route{
|
|
|
|
|
"test-route": {
|
|
|
|
|
Name: "test-route",
|
|
|
|
|
Upstream: config.Upstream{
|
|
|
|
|
Address: upstreamAddr,
|
|
|
|
|
PathRewrite: "",
|
|
|
|
|
ConnectTimeout: 5 * time.Second,
|
|
|
|
|
ReadTimeout: 5 * time.Second,
|
|
|
|
|
WriteTimeout: 5 * time.Second,
|
|
|
|
|
MaxBodySize: 1024 * 1024,
|
|
|
|
|
AuthRequired: false,
|
|
|
|
|
},
|
|
|
|
|
},
|
|
|
|
|
},
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
handler := New(cfg)
|
|
|
|
|
defer handler.Close()
|
|
|
|
|
|
|
|
|
|
// Verify connection reuse by checking that the same transport is used
|
2026-08-19 23:55:42 -07:00
|
|
|
|
2026-08-19 20:52:13 -07:00
|
|
|
// We can't easily count raw TCP connections in this test setup,
|
|
|
|
|
// but we can verify that the transport is being reused by checking
|
|
|
|
|
// that the same transport handles both requests
|
|
|
|
|
route := handler.routes["test-route"]
|
|
|
|
|
firstTransport := route.Transport
|
|
|
|
|
|
|
|
|
|
server := httptest.NewServer(handler)
|
|
|
|
|
defer server.Close()
|
|
|
|
|
|
|
|
|
|
// Make two sequential requests
|
|
|
|
|
http.Get(server.URL + "/path1")
|
|
|
|
|
http.Get(server.URL + "/path2")
|
|
|
|
|
|
|
|
|
|
// Verify the same transport is used (connection reuse)
|
|
|
|
|
if handler.transports[upstreamAddr] != firstTransport {
|
|
|
|
|
t.Errorf("transport changed between requests")
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
// The transport should have been created once
|
|
|
|
|
if len(handler.transports) != 1 {
|
|
|
|
|
t.Errorf("expected 1 transport, got %d", len(handler.transports))
|
|
|
|
|
}
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
func TestProxyNotFound(t *testing.T) {
|
|
|
|
|
cfg := &config.Config{
|
|
|
|
|
Routes: make(map[string]*config.Route),
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
handler := New(cfg)
|
|
|
|
|
defer handler.Close()
|
|
|
|
|
|
|
|
|
|
server := httptest.NewServer(handler)
|
|
|
|
|
defer server.Close()
|
|
|
|
|
|
|
|
|
|
resp, err := http.Get(server.URL + "/nonexistent")
|
|
|
|
|
if err != nil {
|
|
|
|
|
t.Fatalf("request failed: %v", err)
|
|
|
|
|
}
|
|
|
|
|
defer resp.Body.Close()
|
|
|
|
|
|
|
|
|
|
if resp.StatusCode != http.StatusNotFound {
|
|
|
|
|
t.Errorf("expected status 404, got %d", resp.StatusCode)
|
|
|
|
|
}
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
func TestProxyUpstreamError(t *testing.T) {
|
|
|
|
|
cfg := &config.Config{
|
|
|
|
|
Routes: map[string]*config.Route{
|
|
|
|
|
"bad-route": {
|
|
|
|
|
Name: "bad-route",
|
|
|
|
|
Upstream: config.Upstream{
|
|
|
|
|
Address: "127.0.0.1:1",
|
|
|
|
|
PathRewrite: "",
|
|
|
|
|
ConnectTimeout: 100 * time.Millisecond,
|
|
|
|
|
ReadTimeout: 100 * time.Millisecond,
|
|
|
|
|
WriteTimeout: 100 * time.Millisecond,
|
|
|
|
|
MaxBodySize: 1024 * 1024,
|
|
|
|
|
AuthRequired: false,
|
|
|
|
|
},
|
|
|
|
|
},
|
|
|
|
|
},
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
handler := New(cfg)
|
|
|
|
|
defer handler.Close()
|
|
|
|
|
|
|
|
|
|
server := httptest.NewServer(handler)
|
|
|
|
|
defer server.Close()
|
|
|
|
|
|
|
|
|
|
resp, err := http.Get(server.URL + "/test")
|
|
|
|
|
if err != nil {
|
|
|
|
|
t.Fatalf("request failed: %v", err)
|
|
|
|
|
}
|
|
|
|
|
defer resp.Body.Close()
|
|
|
|
|
|
|
|
|
|
// Should get a 502 Bad Gateway when upstream is unreachable
|
|
|
|
|
if resp.StatusCode != http.StatusBadGateway {
|
|
|
|
|
t.Errorf("expected status 502, got %d", resp.StatusCode)
|
|
|
|
|
}
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
func TestProxyMultipleRoutes(t *testing.T) {
|
|
|
|
|
// Create two different upstream servers
|
|
|
|
|
upstream1 := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
|
|
|
w.WriteHeader(http.StatusOK)
|
|
|
|
|
fmt.Fprintf(w, "upstream1")
|
|
|
|
|
}))
|
|
|
|
|
defer upstream1.Close()
|
|
|
|
|
|
|
|
|
|
upstream2 := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
|
|
|
w.WriteHeader(http.StatusOK)
|
|
|
|
|
fmt.Fprintf(w, "upstream2")
|
|
|
|
|
}))
|
|
|
|
|
defer upstream2.Close()
|
|
|
|
|
|
|
|
|
|
addr1 := strings.TrimPrefix(upstream1.URL, "http://")
|
|
|
|
|
addr2 := strings.TrimPrefix(upstream2.URL, "http://")
|
|
|
|
|
|
|
|
|
|
cfg := &config.Config{
|
|
|
|
|
Routes: map[string]*config.Route{
|
|
|
|
|
"route1": {
|
|
|
|
|
Name: "route1",
|
|
|
|
|
Upstream: config.Upstream{
|
|
|
|
|
Address: addr1,
|
|
|
|
|
ConnectTimeout: 5 * time.Second,
|
|
|
|
|
ReadTimeout: 5 * time.Second,
|
|
|
|
|
WriteTimeout: 5 * time.Second,
|
|
|
|
|
MaxBodySize: 1024 * 1024,
|
|
|
|
|
AuthRequired: false,
|
|
|
|
|
},
|
|
|
|
|
},
|
|
|
|
|
"route2": {
|
|
|
|
|
Name: "route2",
|
|
|
|
|
Upstream: config.Upstream{
|
|
|
|
|
Address: addr2,
|
|
|
|
|
ConnectTimeout: 5 * time.Second,
|
|
|
|
|
ReadTimeout: 5 * time.Second,
|
|
|
|
|
WriteTimeout: 5 * time.Second,
|
|
|
|
|
MaxBodySize: 1024 * 1024,
|
|
|
|
|
AuthRequired: false,
|
|
|
|
|
},
|
|
|
|
|
},
|
|
|
|
|
},
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
handler := New(cfg)
|
|
|
|
|
defer handler.Close()
|
|
|
|
|
|
|
|
|
|
// With multiple routes, requests should be routed somewhere
|
|
|
|
|
// (task 1.1 doesn't specify which route for undecorated requests,
|
|
|
|
|
// but task 2.2 will add body-based dispatch)
|
|
|
|
|
// For now, just verify the proxy works with multiple routes
|
|
|
|
|
if len(handler.routes) != 2 {
|
|
|
|
|
t.Errorf("expected 2 routes, got %d", len(handler.routes))
|
|
|
|
|
}
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
func TestProxyPreservesMethod(t *testing.T) {
|
|
|
|
|
method := ""
|
|
|
|
|
upstreamServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
|
|
|
method = r.Method
|
|
|
|
|
w.WriteHeader(http.StatusOK)
|
|
|
|
|
}))
|
|
|
|
|
defer upstreamServer.Close()
|
|
|
|
|
|
|
|
|
|
upstreamAddr := strings.TrimPrefix(upstreamServer.URL, "http://")
|
|
|
|
|
|
|
|
|
|
cfg := &config.Config{
|
|
|
|
|
Routes: map[string]*config.Route{
|
|
|
|
|
"test-route": {
|
|
|
|
|
Name: "test-route",
|
|
|
|
|
Upstream: config.Upstream{
|
|
|
|
|
Address: upstreamAddr,
|
|
|
|
|
ConnectTimeout: 5 * time.Second,
|
|
|
|
|
ReadTimeout: 5 * time.Second,
|
|
|
|
|
WriteTimeout: 5 * time.Second,
|
|
|
|
|
MaxBodySize: 1024 * 1024,
|
|
|
|
|
AuthRequired: false,
|
|
|
|
|
},
|
|
|
|
|
},
|
|
|
|
|
},
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
handler := New(cfg)
|
|
|
|
|
defer handler.Close()
|
|
|
|
|
|
|
|
|
|
server := httptest.NewServer(handler)
|
|
|
|
|
defer server.Close()
|
|
|
|
|
|
|
|
|
|
methods := []string{"GET", "POST", "PUT", "DELETE"}
|
|
|
|
|
for _, m := range methods {
|
|
|
|
|
req, _ := http.NewRequest(m, server.URL+"/test", nil)
|
|
|
|
|
resp, _ := http.DefaultClient.Do(req)
|
|
|
|
|
resp.Body.Close()
|
|
|
|
|
|
|
|
|
|
if method != m {
|
|
|
|
|
t.Errorf("expected method %s, got %s", m, method)
|
|
|
|
|
}
|
|
|
|
|
}
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
func TestProxyPreservesQueryString(t *testing.T) {
|
|
|
|
|
requestedURL := ""
|
|
|
|
|
upstreamServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
|
|
|
requestedURL = r.URL.String()
|
|
|
|
|
w.WriteHeader(http.StatusOK)
|
|
|
|
|
}))
|
|
|
|
|
defer upstreamServer.Close()
|
|
|
|
|
|
|
|
|
|
upstreamAddr := strings.TrimPrefix(upstreamServer.URL, "http://")
|
|
|
|
|
|
|
|
|
|
cfg := &config.Config{
|
|
|
|
|
Routes: map[string]*config.Route{
|
|
|
|
|
"test-route": {
|
|
|
|
|
Name: "test-route",
|
|
|
|
|
Upstream: config.Upstream{
|
|
|
|
|
Address: upstreamAddr,
|
|
|
|
|
ConnectTimeout: 5 * time.Second,
|
|
|
|
|
ReadTimeout: 5 * time.Second,
|
|
|
|
|
WriteTimeout: 5 * time.Second,
|
|
|
|
|
MaxBodySize: 1024 * 1024,
|
|
|
|
|
AuthRequired: false,
|
|
|
|
|
},
|
|
|
|
|
},
|
|
|
|
|
},
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
handler := New(cfg)
|
|
|
|
|
defer handler.Close()
|
|
|
|
|
|
|
|
|
|
server := httptest.NewServer(handler)
|
|
|
|
|
defer server.Close()
|
|
|
|
|
|
|
|
|
|
http.Get(server.URL + "/test?key=value&other=param")
|
|
|
|
|
|
|
|
|
|
if !strings.Contains(requestedURL, "key=value") || !strings.Contains(requestedURL, "other=param") {
|
|
|
|
|
t.Errorf("query string not preserved: %s", requestedURL)
|
|
|
|
|
}
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
func TestProxyPreservesBody(t *testing.T) {
|
|
|
|
|
receivedBody := ""
|
|
|
|
|
upstreamServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
|
|
|
body, _ := io.ReadAll(r.Body)
|
|
|
|
|
receivedBody = string(body)
|
|
|
|
|
w.WriteHeader(http.StatusOK)
|
|
|
|
|
}))
|
|
|
|
|
defer upstreamServer.Close()
|
|
|
|
|
|
|
|
|
|
upstreamAddr := strings.TrimPrefix(upstreamServer.URL, "http://")
|
|
|
|
|
|
|
|
|
|
cfg := &config.Config{
|
|
|
|
|
Routes: map[string]*config.Route{
|
|
|
|
|
"test-route": {
|
|
|
|
|
Name: "test-route",
|
|
|
|
|
Upstream: config.Upstream{
|
|
|
|
|
Address: upstreamAddr,
|
|
|
|
|
ConnectTimeout: 5 * time.Second,
|
|
|
|
|
ReadTimeout: 5 * time.Second,
|
|
|
|
|
WriteTimeout: 5 * time.Second,
|
|
|
|
|
MaxBodySize: 1024 * 1024,
|
|
|
|
|
AuthRequired: false,
|
|
|
|
|
},
|
|
|
|
|
},
|
|
|
|
|
},
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
handler := New(cfg)
|
|
|
|
|
defer handler.Close()
|
|
|
|
|
|
|
|
|
|
server := httptest.NewServer(handler)
|
|
|
|
|
defer server.Close()
|
|
|
|
|
|
|
|
|
|
testBody := `{"model": "test", "messages": []}`
|
2026-08-19 23:55:42 -07:00
|
|
|
resp, err := http.Post(server.URL+"/test", "application/json", strings.NewReader(testBody))
|
|
|
|
|
if err != nil {
|
|
|
|
|
t.Fatalf("request failed: %v", err)
|
|
|
|
|
}
|
2026-08-19 20:52:13 -07:00
|
|
|
resp.Body.Close()
|
|
|
|
|
|
|
|
|
|
if receivedBody != testBody {
|
|
|
|
|
t.Errorf("expected body %s, got %s", testBody, receivedBody)
|
|
|
|
|
}
|
|
|
|
|
}
|
2026-09-27 11:16:40 +00:00
|
|
|
|
|
|
|
|
// TestAuthenticateV1ProxyLevel exercises the three branches of authenticateV1:
|
|
|
|
|
// no-token→401, bad-token→403, valid-token→proxied (non-401/403).
|
|
|
|
|
func TestAuthenticateV1ProxyLevel(t *testing.T) {
|
|
|
|
|
// Generate RSA keypair for signing test JWTs.
|
|
|
|
|
privateKey, err := rsa.GenerateKey(rand.Reader, 2048)
|
|
|
|
|
if err != nil {
|
|
|
|
|
t.Fatalf("rsa.GenerateKey: %v", err)
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
// Build minimal JWKS document from the RSA public key.
|
|
|
|
|
nBytes := privateKey.PublicKey.N.Bytes()
|
|
|
|
|
eBig := new(big.Int).SetInt64(int64(privateKey.PublicKey.E))
|
|
|
|
|
eBytes := eBig.Bytes()
|
|
|
|
|
jwksDoc, _ := json.Marshal(map[string]interface{}{
|
|
|
|
|
"keys": []map[string]interface{}{
|
|
|
|
|
{
|
|
|
|
|
"kty": "RSA",
|
|
|
|
|
"use": "sig",
|
|
|
|
|
"kid": "test-key",
|
|
|
|
|
"alg": "RS256",
|
|
|
|
|
"n": base64.RawURLEncoding.EncodeToString(nBytes),
|
|
|
|
|
"e": base64.RawURLEncoding.EncodeToString(eBytes),
|
|
|
|
|
},
|
|
|
|
|
},
|
|
|
|
|
})
|
|
|
|
|
|
|
|
|
|
// Serve JWKS from a local httptest server.
|
|
|
|
|
jwksSrv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
|
|
|
w.Header().Set("Content-Type", "application/json")
|
|
|
|
|
w.Write(jwksDoc)
|
|
|
|
|
}))
|
|
|
|
|
defer jwksSrv.Close()
|
|
|
|
|
|
|
|
|
|
// Upstream stub — records whether it was reached.
|
|
|
|
|
reached := make(chan struct{}, 1)
|
|
|
|
|
upstreamSrv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
|
|
|
select {
|
|
|
|
|
case reached <- struct{}{}:
|
|
|
|
|
default:
|
|
|
|
|
}
|
|
|
|
|
w.WriteHeader(http.StatusOK)
|
|
|
|
|
}))
|
|
|
|
|
defer upstreamSrv.Close()
|
|
|
|
|
upstreamAddr := strings.TrimPrefix(upstreamSrv.URL, "http://")
|
|
|
|
|
|
|
|
|
|
// signJWT returns a valid RS256 JWT accepted by the Authentik-pattern issuer check.
|
|
|
|
|
signJWT := func() string {
|
|
|
|
|
tok := jwt.NewWithClaims(jwt.SigningMethodRS256, jwt.MapClaims{
|
|
|
|
|
"iss": "https://authentik.riotpiao.com/application/o/test/",
|
|
|
|
|
"aud": "test-audience",
|
|
|
|
|
"sub": "test-user",
|
|
|
|
|
"exp": time.Now().Add(time.Hour).Unix(),
|
|
|
|
|
"iat": time.Now().Unix(),
|
|
|
|
|
})
|
|
|
|
|
tok.Header["kid"] = "test-key"
|
|
|
|
|
signed, err := tok.SignedString(privateKey)
|
|
|
|
|
if err != nil {
|
|
|
|
|
t.Fatalf("sign JWT: %v", err)
|
|
|
|
|
}
|
|
|
|
|
return signed
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
// Build proxy with auth enabled.
|
|
|
|
|
cfg := &config.Config{
|
|
|
|
|
Routes: make(map[string]*config.Route),
|
|
|
|
|
Models: map[string]*config.ModelUpstream{
|
|
|
|
|
"auth-model": {
|
|
|
|
|
Name: "auth-model",
|
|
|
|
|
Address: upstreamAddr,
|
|
|
|
|
},
|
|
|
|
|
},
|
|
|
|
|
Auth: config.AuthConfig{
|
|
|
|
|
Enabled: true,
|
|
|
|
|
JWKSURL: jwksSrv.URL,
|
|
|
|
|
},
|
|
|
|
|
}
|
|
|
|
|
h := New(cfg)
|
|
|
|
|
defer h.Close()
|
|
|
|
|
srv := httptest.NewServer(h)
|
|
|
|
|
defer srv.Close()
|
|
|
|
|
|
|
|
|
|
mkReq := func(authHeader string) *http.Request {
|
|
|
|
|
body := strings.NewReader(`{"model":"auth-model","messages":[]}`)
|
|
|
|
|
req, _ := http.NewRequest("POST", srv.URL+"/v1/chat/completions", body)
|
|
|
|
|
req.Header.Set("Content-Type", "application/json")
|
|
|
|
|
if authHeader != "" {
|
|
|
|
|
req.Header.Set("Authorization", authHeader)
|
|
|
|
|
}
|
|
|
|
|
return req
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
// Branch 1: no Authorization header → 401.
|
|
|
|
|
t.Run("no_token_returns_401", func(t *testing.T) {
|
|
|
|
|
resp, err := http.DefaultClient.Do(mkReq(""))
|
|
|
|
|
if err != nil {
|
|
|
|
|
t.Fatalf("request failed: %v", err)
|
|
|
|
|
}
|
|
|
|
|
defer resp.Body.Close()
|
|
|
|
|
if resp.StatusCode != http.StatusUnauthorized {
|
|
|
|
|
t.Errorf("expected 401, got %d", resp.StatusCode)
|
|
|
|
|
}
|
|
|
|
|
})
|
|
|
|
|
|
|
|
|
|
// Branch 2: malformed/invalid JWT → 403.
|
|
|
|
|
t.Run("bad_token_returns_403", func(t *testing.T) {
|
|
|
|
|
resp, err := http.DefaultClient.Do(mkReq("Bearer this.is.garbage"))
|
|
|
|
|
if err != nil {
|
|
|
|
|
t.Fatalf("request failed: %v", err)
|
|
|
|
|
}
|
|
|
|
|
defer resp.Body.Close()
|
|
|
|
|
if resp.StatusCode != http.StatusForbidden {
|
|
|
|
|
t.Errorf("expected 403, got %d", resp.StatusCode)
|
|
|
|
|
}
|
|
|
|
|
})
|
|
|
|
|
|
|
|
|
|
// Branch 3: valid RS256 JWT → proxied to upstream (not 401/403).
|
|
|
|
|
t.Run("valid_token_proxied", func(t *testing.T) {
|
|
|
|
|
resp, err := http.DefaultClient.Do(mkReq("Bearer " + signJWT()))
|
|
|
|
|
if err != nil {
|
|
|
|
|
t.Fatalf("request failed: %v", err)
|
|
|
|
|
}
|
|
|
|
|
defer resp.Body.Close()
|
|
|
|
|
if resp.StatusCode == http.StatusUnauthorized || resp.StatusCode == http.StatusForbidden {
|
|
|
|
|
body, _ := io.ReadAll(resp.Body)
|
|
|
|
|
t.Errorf("expected proxied response, got %d body=%s", resp.StatusCode, body)
|
|
|
|
|
}
|
|
|
|
|
select {
|
|
|
|
|
case <-reached:
|
|
|
|
|
// upstream was hit — auth passed
|
|
|
|
|
default:
|
|
|
|
|
t.Error("upstream was not reached after valid JWT")
|
|
|
|
|
}
|
|
|
|
|
})
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
// TestProxyQueueFull verifies that a 503 is returned and the queue_full metric
|
|
|
|
|
// is incremented when the backend queue is at capacity.
|
|
|
|
|
//
|
|
|
|
|
// Setup: maxSlots=1, maxQueue=1. Request 1 holds the slot. Request 2 fills the
|
|
|
|
|
// queue (blocks waiting for slot). Request 3 gets ErrQueueFull → 503.
|
|
|
|
|
func TestProxyQueueFull(t *testing.T) {
|
|
|
|
|
t.Setenv("VLLM_MAX_NUM_SEQS", "1") // 1 concurrent slot
|
|
|
|
|
t.Setenv("VLLM_MAX_QUEUE_DEPTH", "1") // queue capacity 1
|
|
|
|
|
t.Setenv("WARMUP_DELAY_MS", "0")
|
|
|
|
|
observability.Global().Reset()
|
|
|
|
|
|
|
|
|
|
// Upstream blocks until released; buffers one signal when it first starts.
|
|
|
|
|
slotHeld := make(chan struct{}, 1)
|
|
|
|
|
release := make(chan struct{})
|
|
|
|
|
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
|
|
|
select {
|
|
|
|
|
case slotHeld <- struct{}{}:
|
|
|
|
|
default:
|
|
|
|
|
}
|
|
|
|
|
<-release
|
|
|
|
|
w.WriteHeader(http.StatusOK)
|
|
|
|
|
}))
|
|
|
|
|
defer upstream.Close()
|
|
|
|
|
|
|
|
|
|
upstreamAddr := strings.TrimPrefix(upstream.URL, "http://")
|
|
|
|
|
cfg := &config.Config{
|
|
|
|
|
Routes: make(map[string]*config.Route),
|
|
|
|
|
Models: map[string]*config.ModelUpstream{
|
|
|
|
|
"q-model": {
|
|
|
|
|
Name: "q-model",
|
|
|
|
|
Address: upstreamAddr,
|
|
|
|
|
},
|
|
|
|
|
},
|
|
|
|
|
}
|
|
|
|
|
h := New(cfg)
|
|
|
|
|
defer h.Close()
|
|
|
|
|
|
|
|
|
|
srv := httptest.NewServer(h)
|
|
|
|
|
defer srv.Close()
|
|
|
|
|
|
|
|
|
|
mkReq := func() *http.Request {
|
|
|
|
|
body := strings.NewReader(`{"model":"q-model","messages":[]}`)
|
|
|
|
|
req, _ := http.NewRequest("POST", srv.URL+"/v1/chat/completions", body)
|
|
|
|
|
req.Header.Set("Content-Type", "application/json")
|
|
|
|
|
return req
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
// Request 1: acquires the only slot, reaches upstream, blocks there.
|
|
|
|
|
r1ch := make(chan *http.Response, 1)
|
|
|
|
|
go func() {
|
|
|
|
|
resp, _ := http.DefaultClient.Do(mkReq())
|
|
|
|
|
r1ch <- resp
|
|
|
|
|
}()
|
|
|
|
|
<-slotHeld // slot is acquired; upstream is holding request 1
|
|
|
|
|
|
|
|
|
|
// Request 2: enters the queue (capacity 1), blocks waiting for slot.
|
|
|
|
|
r2ch := make(chan *http.Response, 1)
|
|
|
|
|
go func() {
|
|
|
|
|
resp, _ := http.DefaultClient.Do(mkReq())
|
|
|
|
|
r2ch <- resp
|
|
|
|
|
}()
|
|
|
|
|
// Spin until request 2 enters the backend queue channel.
|
|
|
|
|
for len(h.schedulers["q-model"].backends[0].queue) == 0 {
|
|
|
|
|
runtime.Gosched()
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
// Request 3: queue is full → expect 503 and metric increment.
|
|
|
|
|
resp3, err := http.DefaultClient.Do(mkReq())
|
|
|
|
|
if err != nil {
|
|
|
|
|
t.Fatalf("request 3 failed: %v", err)
|
|
|
|
|
}
|
|
|
|
|
defer resp3.Body.Close()
|
|
|
|
|
|
|
|
|
|
if resp3.StatusCode != http.StatusServiceUnavailable {
|
|
|
|
|
t.Errorf("expected 503 queue full, got %d", resp3.StatusCode)
|
|
|
|
|
}
|
|
|
|
|
if got := observability.Global().GetQueueFull(); got != 1 {
|
|
|
|
|
t.Errorf("expected GetQueueFull()=1, got %d", got)
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
// Cleanup: unblock upstream so goroutines can finish.
|
|
|
|
|
close(release)
|
|
|
|
|
if r := <-r1ch; r != nil {
|
|
|
|
|
r.Body.Close()
|
|
|
|
|
}
|
|
|
|
|
if r := <-r2ch; r != nil {
|
|
|
|
|
r.Body.Close()
|
|
|
|
|
}
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
// TestProxyCacheHitLogging verifies that when a response contains the
|
|
|
|
|
// vllm-num-cached-tokens header and a cache key was set, the prefix cache
|
|
|
|
|
// hit token count is incremented in observability metrics.
|
|
|
|
|
func TestProxyCacheHitLogging(t *testing.T) {
|
|
|
|
|
observability.Global().Reset()
|
|
|
|
|
|
|
|
|
|
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
|
|
|
w.Header().Set("vllm-num-cached-tokens", "128")
|
|
|
|
|
w.WriteHeader(http.StatusOK)
|
|
|
|
|
}))
|
|
|
|
|
defer upstream.Close()
|
|
|
|
|
|
|
|
|
|
upstreamAddr := strings.TrimPrefix(upstream.URL, "http://")
|
|
|
|
|
cfg := &config.Config{
|
|
|
|
|
Routes: make(map[string]*config.Route),
|
|
|
|
|
Models: map[string]*config.ModelUpstream{
|
|
|
|
|
"cache-model": {
|
|
|
|
|
Name: "cache-model",
|
|
|
|
|
Address: upstreamAddr,
|
|
|
|
|
},
|
|
|
|
|
},
|
|
|
|
|
}
|
|
|
|
|
h := New(cfg)
|
|
|
|
|
defer h.Close()
|
|
|
|
|
|
|
|
|
|
srv := httptest.NewServer(h)
|
|
|
|
|
defer srv.Close()
|
|
|
|
|
|
|
|
|
|
body := strings.NewReader(`{"model":"cache-model","messages":[]}`)
|
|
|
|
|
req, _ := http.NewRequest("POST", srv.URL+"/v1/chat/completions", body)
|
|
|
|
|
req.Header.Set("Content-Type", "application/json")
|
|
|
|
|
req.Header.Set("X-Poimen-Cache-Key", "testkey123")
|
|
|
|
|
|
|
|
|
|
resp, err := http.DefaultClient.Do(req)
|
|
|
|
|
if err != nil {
|
|
|
|
|
t.Fatalf("request failed: %v", err)
|
|
|
|
|
}
|
|
|
|
|
resp.Body.Close()
|
|
|
|
|
|
|
|
|
|
if got := observability.Global().GetPrefixCacheHitTokens(); got != 128 {
|
|
|
|
|
t.Errorf("expected PrefixCacheHitTokens=128, got %d", got)
|
|
|
|
|
}
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
// TestWarmupHeaderHandling verifies that X-Poimen-Warmup: true is stripped from
|
|
|
|
|
// the upstream request and replaced with X-Priority: low, while a request without
|
|
|
|
|
// X-Poimen-Warmup reaches upstream without either header.
|
|
|
|
|
//
|
|
|
|
|
// The warmup path runs only for model-based routes (cfg.Models + scheduler),
|
|
|
|
|
// so we POST to /v1/chat/completions with a JSON body that names the test model.
|
|
|
|
|
// WARMUP_DELAY_MS=0 eliminates the 500 ms sleep.
|
|
|
|
|
func TestWarmupHeaderHandling(t *testing.T) {
|
|
|
|
|
t.Setenv("WARMUP_DELAY_MS", "0")
|
|
|
|
|
|
|
|
|
|
tests := []struct {
|
|
|
|
|
name string
|
|
|
|
|
sendWarmup bool
|
|
|
|
|
wantWarmupGone bool
|
|
|
|
|
wantPriorityLow bool
|
|
|
|
|
}{
|
|
|
|
|
{
|
|
|
|
|
name: "warmup=true",
|
|
|
|
|
sendWarmup: true,
|
|
|
|
|
wantWarmupGone: true, // X-Poimen-Warmup must NOT reach upstream
|
|
|
|
|
wantPriorityLow: true, // X-Priority: low must be set
|
|
|
|
|
},
|
|
|
|
|
{
|
|
|
|
|
name: "warmup=false",
|
|
|
|
|
sendWarmup: false,
|
|
|
|
|
wantWarmupGone: true, // no header sent, so absent at upstream
|
|
|
|
|
wantPriorityLow: false, // X-Priority must NOT be set
|
|
|
|
|
},
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
for _, tc := range tests {
|
|
|
|
|
t.Run(tc.name, func(t *testing.T) {
|
|
|
|
|
var gotWarmup, gotPriority string
|
|
|
|
|
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
|
|
|
gotWarmup = r.Header.Get("X-Poimen-Warmup")
|
|
|
|
|
gotPriority = r.Header.Get("X-Priority")
|
|
|
|
|
w.WriteHeader(http.StatusOK)
|
|
|
|
|
}))
|
|
|
|
|
defer upstream.Close()
|
|
|
|
|
|
|
|
|
|
upstreamAddr := strings.TrimPrefix(upstream.URL, "http://")
|
|
|
|
|
// Warmup code only runs for model-based routes (scheduler path).
|
|
|
|
|
// Use cfg.Models so New() builds a Scheduler for "warmup-model".
|
|
|
|
|
cfg := &config.Config{
|
|
|
|
|
Routes: make(map[string]*config.Route),
|
|
|
|
|
Models: map[string]*config.ModelUpstream{
|
|
|
|
|
"warmup-model": {
|
|
|
|
|
Name: "warmup-model",
|
|
|
|
|
Address: upstreamAddr,
|
|
|
|
|
},
|
|
|
|
|
},
|
|
|
|
|
}
|
|
|
|
|
h := New(cfg)
|
|
|
|
|
defer h.Close()
|
|
|
|
|
|
|
|
|
|
srv := httptest.NewServer(h)
|
|
|
|
|
defer srv.Close()
|
|
|
|
|
|
|
|
|
|
// POST to /v1/chat/completions — triggers model-based dispatch and warmup path.
|
|
|
|
|
body := strings.NewReader(`{"model":"warmup-model","messages":[]}`)
|
|
|
|
|
req, _ := http.NewRequest("POST", srv.URL+"/v1/chat/completions", body)
|
|
|
|
|
req.Header.Set("Content-Type", "application/json")
|
|
|
|
|
if tc.sendWarmup {
|
|
|
|
|
req.Header.Set("X-Poimen-Warmup", "true")
|
|
|
|
|
}
|
|
|
|
|
resp, err := http.DefaultClient.Do(req)
|
|
|
|
|
if err != nil {
|
|
|
|
|
t.Fatalf("request failed: %v", err)
|
|
|
|
|
}
|
|
|
|
|
resp.Body.Close()
|
|
|
|
|
|
|
|
|
|
if tc.wantWarmupGone && gotWarmup != "" {
|
|
|
|
|
t.Errorf("X-Poimen-Warmup must be stripped before upstream, got %q", gotWarmup)
|
|
|
|
|
}
|
|
|
|
|
if tc.wantPriorityLow && gotPriority != "low" {
|
|
|
|
|
t.Errorf("X-Priority: low must be set when warmup=true, got %q", gotPriority)
|
|
|
|
|
}
|
|
|
|
|
if !tc.wantPriorityLow && gotPriority != "" {
|
|
|
|
|
t.Errorf("X-Priority must not be set for normal requests, got %q", gotPriority)
|
|
|
|
|
}
|
|
|
|
|
})
|
|
|
|
|
}
|
|
|
|
|
}
|