Something went wrong. Try again.
sloproxy
Something went wrong. Try again.
16 kB · 528 lines
Go
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529package middleware
import ( "bytes" "context" "io" "net/http" "net/http/httptest" "net/url" "os" "strings" "testing"
"github.com/gin-gonic/gin" "github.com/klauspost/compress/zstd" "github.com/router-for-me/CLIProxyAPI/v7/internal/clienterror" "github.com/router-for-me/CLIProxyAPI/v7/internal/config" "github.com/router-for-me/CLIProxyAPI/v7/internal/logging" "github.com/router-for-me/CLIProxyAPI/v7/internal/runtime/executor/helps")
func TestShouldSkipMethodForRequestLogging(t *testing.T) { tests := []struct { name string req *http.Request skip bool }{ { name: "nil request", req: nil, skip: true, }, { name: "post request should not skip", req: &http.Request{ Method: http.MethodPost, URL: &url.URL{Path: "/v1/responses"}, }, skip: false, }, { name: "plain get should skip", req: &http.Request{ Method: http.MethodGet, URL: &url.URL{Path: "/v1/models"}, Header: http.Header{}, }, skip: true, }, { name: "responses websocket upgrade should not skip", req: &http.Request{ Method: http.MethodGet, URL: &url.URL{Path: "/v1/responses"}, Header: http.Header{"Upgrade": []string{"websocket"}}, }, skip: false, }, { name: "codex responses websocket upgrade should not skip", req: &http.Request{ Method: http.MethodGet, URL: &url.URL{Path: "/backend-api/codex/responses"}, Header: http.Header{"Upgrade": []string{"websocket"}}, }, skip: false, }, { name: "responses get without upgrade should skip", req: &http.Request{ Method: http.MethodGet, URL: &url.URL{Path: "/v1/responses"}, Header: http.Header{}, }, skip: true, }, }
for i := range tests { got := shouldSkipMethodForRequestLogging(tests[i].req) if got != tests[i].skip { t.Fatalf("%s: got skip=%t, want %t", tests[i].name, got, tests[i].skip) } }}
func TestShouldCaptureRequestBody(t *testing.T) { tests := []struct { name string loggerEnabled bool req *http.Request want bool }{ { name: "logger enabled always captures", loggerEnabled: true, req: &http.Request{ Body: io.NopCloser(strings.NewReader("{}")), ContentLength: -1, Header: http.Header{"Content-Type": []string{"application/json"}}, }, want: true, }, { name: "nil request", loggerEnabled: false, req: nil, want: false, }, { name: "small known size json in error-only mode", loggerEnabled: false, req: &http.Request{ Body: io.NopCloser(strings.NewReader("{}")), ContentLength: 2, Header: http.Header{"Content-Type": []string{"application/json"}}, }, want: true, }, { name: "large known size skipped in error-only mode", loggerEnabled: false, req: &http.Request{ Body: io.NopCloser(strings.NewReader("x")), ContentLength: maxErrorOnlyCapturedRequestBodyBytes + 1, Header: http.Header{"Content-Type": []string{"application/json"}}, }, want: false, }, { name: "unknown size skipped in error-only mode", loggerEnabled: false, req: &http.Request{ Body: io.NopCloser(strings.NewReader("x")), ContentLength: -1, Header: http.Header{"Content-Type": []string{"application/json"}}, }, want: false, }, { name: "multipart skipped in error-only mode", loggerEnabled: false, req: &http.Request{ Body: io.NopCloser(strings.NewReader("x")), ContentLength: 1, Header: http.Header{"Content-Type": []string{"multipart/form-data; boundary=abc"}}, }, want: false, }, }
for i := range tests { got := shouldCaptureRequestBody(tests[i].loggerEnabled, tests[i].req) if got != tests[i].want { t.Fatalf("%s: got %t, want %t", tests[i].name, got, tests[i].want) } }}
func TestDeferredRequestBodyCaptureDoesNotDrainUnreadBody(t *testing.T) { gin.SetMode(gin.TestMode)
logger := logging.NewFileRequestLogger(false, t.TempDir(), "", 10) request := httptest.NewRequest(http.MethodPost, "/v1/responses", strings.NewReader("remaining-body")) request.ContentLength = -1 request.Header.Set("Content-Type", "application/json") requestInfo := &RequestInfo{Headers: map[string][]string{"Content-Type": {"application/json"}}} capture := attachDeferredRequestBodyCapture(request, logger, requestInfo, false, false) if capture == nil { t.Fatal("deferred request body capture was not attached") } defer capture.Cleanup()
firstByte := make([]byte, 1) if _, errRead := request.Body.Read(firstByte); errRead != nil { t.Fatalf("read first request byte: %v", errRead) } captured, marker, errCaptured := capture.Bytes() if errCaptured != nil { t.Fatalf("read captured body: %v", errCaptured) } if string(captured) != "r" { t.Fatalf("captured body = %q, want %q", string(captured), "r") } if !strings.Contains(marker, "REQUEST BODY CAPTURE INCOMPLETE") { t.Fatalf("capture marker = %q, want incomplete marker", marker) } remaining, errRemaining := io.ReadAll(capture.body) if errRemaining != nil { t.Fatalf("read remaining body: %v", errRemaining) } if string(remaining) != "emaining-body" { t.Fatalf("remaining body = %q, want %q", string(remaining), "emaining-body") }}
func TestRequestLoggingMiddlewareCapturesLargeErrorRequestAndDeferredAPIRequest(t *testing.T) { gin.SetMode(gin.TestMode)
logsDir := t.TempDir() logger := logging.NewFileRequestLogger(false, logsDir, "", 10) payload := append([]byte(`{"marker":"large-error-body","padding":"`), bytes.Repeat([]byte("x"), int(maxErrorOnlyCapturedRequestBodyBytes))...) payload = append(payload, []byte(`"}`)...) upstreamBody := []byte(`{"model":"upstream-model","input":"translated"}`)
router := gin.New() router.Use(RequestLoggingMiddleware(logger)) router.POST("/v1/responses", func(c *gin.Context) { body, errRead := io.ReadAll(c.Request.Body) if errRead != nil { c.Status(http.StatusInternalServerError) return } if !bytes.Equal(body, payload) { c.Status(http.StatusInternalServerError) return } executorCtx := context.WithValue(context.Background(), "gin", c) helps.RecordAPIRequest(executorCtx, &config.Config{}, helps.UpstreamRequestLog{ URL: "https://api.example.com/v1/responses", Method: http.MethodPost, Headers: http.Header{"Content-Type": []string{"application/json"}}, Body: upstreamBody, }) c.JSON(http.StatusBadRequest, gin.H{"error": "upstream rejected request"}) })
request := httptest.NewRequest(http.MethodPost, "/v1/responses", bytes.NewReader(payload)) request.Header.Set("Content-Type", "application/json") response := httptest.NewRecorder() router.ServeHTTP(response, request)
if response.Code != http.StatusBadRequest { t.Fatalf("response status = %d, want %d", response.Code, http.StatusBadRequest) } entries, errReadDir := os.ReadDir(logsDir) if errReadDir != nil { t.Fatalf("read logs dir: %v", errReadDir) } var logPath string for _, entry := range entries { if strings.HasPrefix(entry.Name(), "error-") && strings.HasSuffix(entry.Name(), ".log") { logPath = logsDir + string(os.PathSeparator) + entry.Name() break } } if logPath == "" { t.Fatal("forced error log was not created") } content, errReadLog := os.ReadFile(logPath) if errReadLog != nil { t.Fatalf("read error log: %v", errReadLog) } if !bytes.Contains(content, payload) { t.Fatal("error log does not contain the complete large request body") } if !bytes.Contains(content, []byte("=== API REQUEST 1 ===")) { t.Fatal("error log does not contain the deferred API request section") } if !bytes.Contains(content, upstreamBody) { t.Fatal("error log does not contain the deferred upstream request body") }}
func TestAttachRequestLogSourcesUsesLoggerLogsDir(t *testing.T) { gin.SetMode(gin.TestMode)
logsDir := t.TempDir() logger := logging.NewFileRequestLogger(true, logsDir, "", 0) recorder := httptest.NewRecorder() c, _ := gin.CreateTestContext(recorder) c.Request = httptest.NewRequest(http.MethodGet, "/backend-api/codex/responses", nil) c.Request.Header.Set("Upgrade", "websocket")
attachRequestLogSources(c, logger, true) defer cleanupFileBodySourcesFromContext(c)
for _, key := range []string{ logging.WebsocketTimelineSourceContextKey, logging.APIWebsocketTimelineSourceContextKey, } { value, exists := c.Get(key) if !exists { t.Fatalf("expected %s source to be attached", key) } source, ok := value.(*logging.FileBodySource) if !ok || source == nil { t.Fatalf("%s source type = %T", key, value) } file, errPart := source.CreatePart("probe") if errPart != nil { t.Fatalf("CreatePart(%s): %v", key, errPart) } path := file.Name() if errClose := file.Close(); errClose != nil { t.Fatalf("close part: %v", errClose) } if !strings.HasPrefix(path, logsDir+string(os.PathSeparator)) { t.Fatalf("%s part path %s is not under logs dir %s", key, path, logsDir) } }}
func cleanupFileBodySourcesFromContext(c *gin.Context) { if c == nil { return } for _, key := range []string{ logging.WebsocketTimelineSourceContextKey, logging.APIWebsocketTimelineSourceContextKey, } { value, exists := c.Get(key) if !exists { continue } if source, ok := value.(*logging.FileBodySource); ok && source != nil { _ = source.Cleanup() } }}
func TestDecodeCapturedRequestBodyForLogWithLimitTruncatesZstdExpansion(t *testing.T) { payload := bytes.Repeat([]byte("x"), 1024) var compressed bytes.Buffer encoder, errNewWriter := zstd.NewWriter(&compressed) if errNewWriter != nil { t.Fatalf("zstd.NewWriter: %v", errNewWriter) } if _, errWrite := encoder.Write(payload); errWrite != nil { t.Fatalf("zstd write: %v", errWrite) } if errClose := encoder.Close(); errClose != nil { t.Fatalf("zstd close: %v", errClose) }
decoded := decodeCapturedRequestBodyForLogWithLimit(compressed.Bytes(), "zstd", 64) if len(decoded) > 128 { t.Fatalf("limited decoded body length = %d, want bounded output", len(decoded)) } if !bytes.Contains(decoded, []byte("DECOMPRESSED REQUEST BODY TRUNCATED")) { t.Fatalf("decoded body = %q, want truncation marker", string(decoded)) }}
func TestCaptureRequestInfoDecodesZstdRequestBodyForLog(t *testing.T) { gin.SetMode(gin.TestMode)
payload := []byte(`{"model":"test-model","stream":true}`) var compressed bytes.Buffer encoder, errNewWriter := zstd.NewWriter(&compressed) if errNewWriter != nil { t.Fatalf("zstd.NewWriter: %v", errNewWriter) } if _, errWrite := encoder.Write(payload); errWrite != nil { t.Fatalf("zstd write: %v", errWrite) } if errClose := encoder.Close(); errClose != nil { t.Fatalf("zstd close: %v", errClose) } compressedBytes := compressed.Bytes()
recorder := httptest.NewRecorder() c, _ := gin.CreateTestContext(recorder) req := httptest.NewRequest(http.MethodPost, "/v1/responses", bytes.NewReader(compressedBytes)) req.Header.Set("Content-Encoding", "zstd") c.Request = req
info, errCapture := captureRequestInfo(c, true) if errCapture != nil { t.Fatalf("captureRequestInfo: %v", errCapture) } if !bytes.Equal(info.Body, payload) { t.Fatalf("logged request body = %q, want %q", string(info.Body), string(payload)) }
restoredBody, errRead := io.ReadAll(c.Request.Body) if errRead != nil { t.Fatalf("read restored request body: %v", errRead) } if !bytes.Equal(restoredBody, compressedBytes) { t.Fatal("request body was not restored with the original compressed bytes") }}
func TestRequestLoggingMiddleware_ClientCancellationExclusion(t *testing.T) { gin.SetMode(gin.TestMode)
t.Run("499 status does not create error log when request-log is false", func(t *testing.T) { logsDir := t.TempDir() logger := logging.NewFileRequestLogger(false, logsDir, "", 10)
router := gin.New() router.Use(RequestLoggingMiddleware(logger)) router.POST("/v1/responses", func(c *gin.Context) { c.AbortWithStatus(clienterror.StatusClientClosedRequest) })
req := httptest.NewRequest(http.MethodPost, "/v1/responses", strings.NewReader(`{"model":"gpt-4"}`)) req.Header.Set("Content-Type", "application/json") resp := httptest.NewRecorder() router.ServeHTTP(resp, req)
if resp.Code != clienterror.StatusClientClosedRequest { t.Fatalf("status = %d, want %d", resp.Code, clienterror.StatusClientClosedRequest) }
entries, errRead := os.ReadDir(logsDir) if errRead != nil { t.Fatalf("read logs dir: %v", errRead) } if len(entries) != 0 { t.Fatalf("expected 0 log files for 499 cancellation in error-only mode, got %d files", len(entries)) } })
t.Run("context canceled does not create error log when request-log is false", func(t *testing.T) { logsDir := t.TempDir() logger := logging.NewFileRequestLogger(false, logsDir, "", 10)
router := gin.New() router.Use(RequestLoggingMiddleware(logger)) router.POST("/v1/responses", func(c *gin.Context) { // Simulate client closing connection mid-flight ctx, cancel := context.WithCancel(c.Request.Context()) cancel() c.Request = c.Request.WithContext(ctx) c.Status(http.StatusOK) })
req := httptest.NewRequest(http.MethodPost, "/v1/responses", strings.NewReader(`{"model":"gpt-4"}`)) req.Header.Set("Content-Type", "application/json") resp := httptest.NewRecorder() router.ServeHTTP(resp, req)
entries, errRead := os.ReadDir(logsDir) if errRead != nil { t.Fatalf("read logs dir: %v", errRead) } if len(entries) != 0 { t.Fatalf("expected 0 log files for canceled context in error-only mode, got %d files", len(entries)) } })
t.Run("400 bad request creates error log when request-log is false", func(t *testing.T) { logsDir := t.TempDir() logger := logging.NewFileRequestLogger(false, logsDir, "", 10)
router := gin.New() router.Use(RequestLoggingMiddleware(logger)) router.POST("/v1/responses", func(c *gin.Context) { c.JSON(http.StatusBadRequest, gin.H{"error": "invalid parameter"}) })
req := httptest.NewRequest(http.MethodPost, "/v1/responses", strings.NewReader(`{"bad":"param"}`)) req.Header.Set("Content-Type", "application/json") resp := httptest.NewRecorder() router.ServeHTTP(resp, req)
if resp.Code != http.StatusBadRequest { t.Fatalf("status = %d, want %d", resp.Code, http.StatusBadRequest) }
entries, errRead := os.ReadDir(logsDir) if errRead != nil { t.Fatalf("read logs dir: %v", errRead) } var errorLogCount int for _, entry := range entries { if strings.HasPrefix(entry.Name(), "error-") && strings.HasSuffix(entry.Name(), ".log") { errorLogCount++ } } if errorLogCount != 1 { t.Fatalf("expected 1 error log file for 400 Bad Request, got %d", errorLogCount) } })
t.Run("499 status logs standard request when request-log is true", func(t *testing.T) { logsDir := t.TempDir() logger := logging.NewFileRequestLogger(true, logsDir, "", 10)
router := gin.New() router.Use(RequestLoggingMiddleware(logger)) router.POST("/v1/responses", func(c *gin.Context) { c.AbortWithStatus(clienterror.StatusClientClosedRequest) })
req := httptest.NewRequest(http.MethodPost, "/v1/responses", strings.NewReader(`{"model":"gpt-4"}`)) req.Header.Set("Content-Type", "application/json") resp := httptest.NewRecorder() router.ServeHTTP(resp, req)
entries, errRead := os.ReadDir(logsDir) if errRead != nil { t.Fatalf("read logs dir: %v", errRead) } var standardLogCount int for _, entry := range entries { if !strings.HasPrefix(entry.Name(), "error-") && strings.HasSuffix(entry.Name(), ".log") { standardLogCount++ } } if standardLogCount != 1 { t.Fatalf("expected 1 standard request log file when request-log=true, got %d", standardLogCount) } })}
func TestCaptureRequestInfo_HeadersDeepCopy(t *testing.T) { gin.SetMode(gin.TestMode) w := httptest.NewRecorder() c, _ := gin.CreateTestContext(w) req := httptest.NewRequest(http.MethodPost, "/v1/test", nil) req.Header.Set("X-Audit", "original-value") c.Request = req
info, err := captureRequestInfo(c, false) if err != nil { t.Fatalf("captureRequestInfo failed: %v", err) }
// Mutate the original request header slice in place c.Request.Header["X-Audit"][0] = "mutated-value"
if got := info.Headers["X-Audit"][0]; got != "original-value" { t.Fatalf("header slice was aliased: got %q, want %q", got, "original-value") }}