588 lines · 18.3 KB · Go
RawDownload
package 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 TestRequestLoggingMiddleware_StreamingResponsesUpstreamSections(t *testing.T) {	gin.SetMode(gin.TestMode) 	logsDir := t.TempDir()	logger := logging.NewFileRequestLogger(true, logsDir, "", 10)	cfg := &config.Config{SDKConfig: config.SDKConfig{RequestLog: true}} 	router := gin.New()	router.Use(RequestLoggingMiddleware(logger))	router.POST("/v1/responses", func(c *gin.Context) {		c.Header("Content-Type", "text/event-stream")		executorCtx := context.WithValue(context.Background(), "gin", c)		helps.RecordAPIRequest(executorCtx, cfg, helps.UpstreamRequestLog{			URL:     "https://api.example.com/v1/responses",			Method:  http.MethodPost,			Headers: http.Header{"Content-Type": []string{"application/json"}},			Body:    []byte(`{"model":"gpt-5-codex","input":[]}`),		})		helps.AppendAPIResponseChunk(executorCtx, cfg, []byte("data: {\"type\":\"response.output_item.added\"}\n\n"))		_, _ = c.Writer.Write([]byte("data: {\"type\":\"response.output_item.added\"}\n\n"))		if flusher, ok := c.Writer.(http.Flusher); ok {			flusher.Flush()		}	}) 	request := httptest.NewRequest(http.MethodPost, "/v1/responses", strings.NewReader(`{"model":"gpt-5-codex","input":[],"stream":true}`))	request.Header.Set("Content-Type", "application/json")	response := httptest.NewRecorder()	router.ServeHTTP(response, request) 	if response.Code != http.StatusOK {		t.Fatalf("response status = %d, want %d", response.Code, http.StatusOK)	} 	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(), "v1-responses-") && strings.HasSuffix(entry.Name(), ".log") {			logPath = logsDir + string(os.PathSeparator) + entry.Name()			break		}	}	if logPath == "" {		t.Fatal("streaming request log was not created")	}	content, errReadLog := os.ReadFile(logPath)	if errReadLog != nil {		t.Fatalf("read log file: %v", errReadLog)	}	if !bytes.Contains(content, []byte("=== API REQUEST 1 ===")) {		t.Fatalf("streaming log missing API REQUEST: %s", string(content))	}	if !bytes.Contains(content, []byte("=== API RESPONSE 1 ===")) {		t.Fatalf("streaming log missing API RESPONSE: %s", string(content))	}} 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))

Lines 1–500 of 588.