mirror of
https://gh-proxy.org/https://github.com/klinsmaya/CLIProxyAPI.git
synced 2026-11-10 10:26:48 +08:00
feat(logging): add deferred request body capture for error handling
- Introduced `deferredRequestBodyCapture` for spooling and retaining large error-only request bodies for improved logging. - Enhanced `RequestLoggingMiddleware` to capture request body based on size and error criteria. - Implemented decompression and truncation for large request bodies with Zstd encoding. - Added support for deferred API request logging and context tracking. - Included extensive test coverage for error-specific body capture and serialization behavior.
This commit is contained in:
@@ -8,6 +8,7 @@ import (
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"os"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
@@ -15,14 +16,18 @@ import (
|
||||
"github.com/klauspost/compress/zstd"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/logging"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/util"
|
||||
log "github.com/sirupsen/logrus"
|
||||
)
|
||||
|
||||
const maxErrorOnlyCapturedRequestBodyBytes int64 = 1 << 20 // 1 MiB
|
||||
const (
|
||||
maxErrorOnlyCapturedRequestBodyBytes int64 = 1 << 20 // 1 MiB
|
||||
maxDeferredErrorRequestBodyBytes int64 = 32 << 20 // 32 MiB
|
||||
)
|
||||
|
||||
// RequestLoggingMiddleware creates a Gin middleware that logs HTTP requests and responses.
|
||||
// It captures detailed information about the request and response, including headers and body,
|
||||
// and uses the provided RequestLogger to record this data. When full request logging is disabled,
|
||||
// body capture is limited to small known-size payloads to avoid large per-request memory spikes.
|
||||
// large and unknown-size bodies are spooled to disk and retained only for error logs.
|
||||
func RequestLoggingMiddleware(logger logging.RequestLogger) gin.HandlerFunc {
|
||||
return func(c *gin.Context) {
|
||||
if logger == nil {
|
||||
@@ -42,9 +47,10 @@ func RequestLoggingMiddleware(logger logging.RequestLogger) gin.HandlerFunc {
|
||||
}
|
||||
|
||||
loggerEnabled := logger.IsEnabled()
|
||||
captureBody := shouldCaptureRequestBody(loggerEnabled, c.Request)
|
||||
|
||||
// Capture request information
|
||||
requestInfo, err := captureRequestInfo(c, shouldCaptureRequestBody(loggerEnabled, c.Request))
|
||||
requestInfo, err := captureRequestInfo(c, captureBody)
|
||||
if err != nil {
|
||||
// Log error but continue processing
|
||||
// In a real implementation, you might want to use a proper logger here
|
||||
@@ -59,6 +65,7 @@ func RequestLoggingMiddleware(logger logging.RequestLogger) gin.HandlerFunc {
|
||||
}
|
||||
c.Writer = wrapper
|
||||
attachRequestLogSources(c, logger, loggerEnabled)
|
||||
attachDeferredRequestBodyCapture(c.Request, logger, requestInfo, loggerEnabled, captureBody)
|
||||
|
||||
// Process the request
|
||||
c.Next()
|
||||
@@ -75,6 +82,167 @@ type fileBodySourceFactory interface {
|
||||
NewFileBodySource(prefix string) (*logging.FileBodySource, error)
|
||||
}
|
||||
|
||||
type deferredRequestBodyCapture struct {
|
||||
body io.ReadCloser
|
||||
file *os.File
|
||||
source *logging.FileBodySource
|
||||
contentLength int64
|
||||
bytesRead int64
|
||||
bytesCaptured int64
|
||||
captureErr error
|
||||
finished bool
|
||||
sawEOF bool
|
||||
truncated bool
|
||||
}
|
||||
|
||||
func attachDeferredRequestBodyCapture(req *http.Request, logger logging.RequestLogger, requestInfo *RequestInfo, loggerEnabled, bodyCaptured bool) *deferredRequestBodyCapture {
|
||||
if loggerEnabled || bodyCaptured || req == nil || req.Body == nil || req.Body == http.NoBody || req.ContentLength == 0 || requestInfo == nil {
|
||||
return nil
|
||||
}
|
||||
contentType := strings.ToLower(strings.TrimSpace(req.Header.Get("Content-Type")))
|
||||
if strings.HasPrefix(contentType, "multipart/form-data") {
|
||||
return nil
|
||||
}
|
||||
factory, ok := logger.(fileBodySourceFactory)
|
||||
if !ok || factory == nil {
|
||||
return nil
|
||||
}
|
||||
source, errSource := factory.NewFileBodySource("request-body")
|
||||
if errSource != nil {
|
||||
return nil
|
||||
}
|
||||
file, errPart := source.CreatePart("body")
|
||||
if errPart != nil {
|
||||
_ = source.Cleanup()
|
||||
return nil
|
||||
}
|
||||
capture := &deferredRequestBodyCapture{
|
||||
body: req.Body,
|
||||
file: file,
|
||||
source: source,
|
||||
contentLength: req.ContentLength,
|
||||
}
|
||||
req.Body = capture
|
||||
requestInfo.deferredBodyCapture = capture
|
||||
return capture
|
||||
}
|
||||
|
||||
func (c *deferredRequestBodyCapture) Read(payload []byte) (int, error) {
|
||||
if c == nil || c.body == nil {
|
||||
return 0, io.EOF
|
||||
}
|
||||
n, errRead := c.body.Read(payload)
|
||||
if errRead == io.EOF {
|
||||
c.sawEOF = true
|
||||
}
|
||||
if n == 0 {
|
||||
return n, errRead
|
||||
}
|
||||
c.bytesRead += int64(n)
|
||||
if c.file == nil || c.captureErr != nil {
|
||||
return n, errRead
|
||||
}
|
||||
|
||||
remaining := maxDeferredErrorRequestBodyBytes - c.bytesCaptured
|
||||
if remaining <= 0 {
|
||||
c.truncated = true
|
||||
return n, errRead
|
||||
}
|
||||
writeLength := int64(n)
|
||||
if writeLength > remaining {
|
||||
writeLength = remaining
|
||||
c.truncated = true
|
||||
}
|
||||
written, errWrite := c.file.Write(payload[:int(writeLength)])
|
||||
c.bytesCaptured += int64(written)
|
||||
if errWrite != nil {
|
||||
c.captureErr = errWrite
|
||||
} else if int64(written) != writeLength {
|
||||
c.captureErr = io.ErrShortWrite
|
||||
}
|
||||
if c.captureErr != nil {
|
||||
if errClose := c.file.Close(); errClose != nil {
|
||||
c.captureErr = fmt.Errorf("%v; close capture file: %w", c.captureErr, errClose)
|
||||
}
|
||||
c.file = nil
|
||||
}
|
||||
return n, errRead
|
||||
}
|
||||
|
||||
func (c *deferredRequestBodyCapture) Close() error {
|
||||
if c == nil {
|
||||
return nil
|
||||
}
|
||||
_ = c.Finish()
|
||||
if c.body == nil {
|
||||
return nil
|
||||
}
|
||||
return c.body.Close()
|
||||
}
|
||||
|
||||
func (c *deferredRequestBodyCapture) Finish() error {
|
||||
if c == nil {
|
||||
return nil
|
||||
}
|
||||
if c.finished {
|
||||
return c.captureErr
|
||||
}
|
||||
c.finished = true
|
||||
if c.file != nil {
|
||||
if errClose := c.file.Close(); errClose != nil && c.captureErr == nil {
|
||||
c.captureErr = errClose
|
||||
}
|
||||
c.file = nil
|
||||
}
|
||||
return c.captureErr
|
||||
}
|
||||
|
||||
func (c *deferredRequestBodyCapture) Bytes() ([]byte, string, error) {
|
||||
if c == nil || c.source == nil {
|
||||
return nil, "", nil
|
||||
}
|
||||
if errFinish := c.Finish(); errFinish != nil {
|
||||
return nil, "", errFinish
|
||||
}
|
||||
body, errBytes := c.source.Bytes()
|
||||
if errBytes != nil {
|
||||
return nil, "", errBytes
|
||||
}
|
||||
return body, c.statusMarker(), nil
|
||||
}
|
||||
|
||||
func (c *deferredRequestBodyCapture) statusMarker() string {
|
||||
if c == nil {
|
||||
return ""
|
||||
}
|
||||
var markers []string
|
||||
if c.truncated {
|
||||
markers = append(markers, fmt.Sprintf("[REQUEST BODY TRUNCATED: captured first %d bytes]", c.bytesCaptured))
|
||||
}
|
||||
complete := c.sawEOF || (c.contentLength >= 0 && c.bytesRead >= c.contentLength)
|
||||
if !complete {
|
||||
if c.contentLength >= 0 {
|
||||
markers = append(markers, fmt.Sprintf("[REQUEST BODY CAPTURE INCOMPLETE: consumed %d of %d bytes]", c.bytesRead, c.contentLength))
|
||||
} else {
|
||||
markers = append(markers, fmt.Sprintf("[REQUEST BODY CAPTURE INCOMPLETE: consumed %d bytes from an unknown-length body]", c.bytesRead))
|
||||
}
|
||||
}
|
||||
return strings.Join(markers, "\n")
|
||||
}
|
||||
|
||||
func (c *deferredRequestBodyCapture) Cleanup() {
|
||||
if c == nil || c.source == nil {
|
||||
return
|
||||
}
|
||||
if errFinish := c.Finish(); errFinish != nil {
|
||||
log.WithError(errFinish).Warn("failed to finish deferred request body capture")
|
||||
}
|
||||
if errCleanup := c.source.Cleanup(); errCleanup != nil {
|
||||
log.WithError(errCleanup).Warn("failed to clean up deferred request body capture")
|
||||
}
|
||||
c.source = nil
|
||||
}
|
||||
|
||||
func attachRequestLogSources(c *gin.Context, logger logging.RequestLogger, loggerEnabled bool) {
|
||||
if c == nil || !loggerEnabled {
|
||||
return
|
||||
@@ -193,6 +361,41 @@ func decodeCapturedRequestBodyForLog(raw []byte, encoding string) []byte {
|
||||
return decoded
|
||||
}
|
||||
|
||||
func decodeCapturedRequestBodyForLogWithLimit(raw []byte, encoding string, limit int64) []byte {
|
||||
if len(raw) == 0 || limit <= 0 {
|
||||
return raw
|
||||
}
|
||||
encoding = strings.TrimSpace(encoding)
|
||||
if encoding == "" || strings.EqualFold(encoding, "identity") {
|
||||
return raw
|
||||
}
|
||||
|
||||
parts := strings.Split(encoding, ",")
|
||||
body := raw
|
||||
for i := len(parts) - 1; i >= 0; i-- {
|
||||
enc := strings.ToLower(strings.TrimSpace(parts[i]))
|
||||
switch enc {
|
||||
case "", "identity":
|
||||
continue
|
||||
case "zstd":
|
||||
decoded, truncated, errDecode := decodeCapturedZstdRequestBodyWithLimit(body, limit)
|
||||
if errDecode != nil {
|
||||
return raw
|
||||
}
|
||||
body = decoded
|
||||
if truncated {
|
||||
if len(body) > 0 && !bytes.HasSuffix(body, []byte("\n")) {
|
||||
body = append(body, '\n')
|
||||
}
|
||||
return append(body, "[DECOMPRESSED REQUEST BODY TRUNCATED]"...)
|
||||
}
|
||||
default:
|
||||
return raw
|
||||
}
|
||||
}
|
||||
return body
|
||||
}
|
||||
|
||||
func decodeCapturedRequestBody(raw []byte, encoding string) ([]byte, error) {
|
||||
encoding = strings.TrimSpace(encoding)
|
||||
if encoding == "" || strings.EqualFold(encoding, "identity") {
|
||||
@@ -233,6 +436,23 @@ func decodeCapturedZstdRequestBody(raw []byte) ([]byte, error) {
|
||||
return decoded, nil
|
||||
}
|
||||
|
||||
func decodeCapturedZstdRequestBodyWithLimit(raw []byte, limit int64) ([]byte, bool, error) {
|
||||
decoder, errNewReader := zstd.NewReader(bytes.NewReader(raw))
|
||||
if errNewReader != nil {
|
||||
return nil, false, fmt.Errorf("failed to create zstd request decoder: %w", errNewReader)
|
||||
}
|
||||
defer decoder.Close()
|
||||
|
||||
decoded, errRead := io.ReadAll(io.LimitReader(decoder, limit+1))
|
||||
if errRead != nil {
|
||||
return nil, false, fmt.Errorf("failed to decode zstd request body: %w", errRead)
|
||||
}
|
||||
if int64(len(decoded)) > limit {
|
||||
return decoded[:limit], true, nil
|
||||
}
|
||||
return decoded, false, nil
|
||||
}
|
||||
|
||||
// shouldLogRequest determines whether the request should be logged.
|
||||
// It skips management endpoints to avoid leaking secrets but allows
|
||||
// all other routes, including module-provided ones, to honor request-log.
|
||||
|
||||
@@ -2,6 +2,7 @@ package middleware
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
@@ -12,7 +13,9 @@ import (
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/klauspost/compress/zstd"
|
||||
"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) {
|
||||
@@ -153,6 +156,111 @@ func TestShouldCaptureRequestBody(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
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)
|
||||
|
||||
@@ -210,6 +318,29 @@ func cleanupFileBodySourcesFromContext(c *gin.Context) {
|
||||
}
|
||||
}
|
||||
|
||||
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)
|
||||
|
||||
|
||||
@@ -21,12 +21,13 @@ const websocketTimelineOverrideContextKey = "WEBSOCKET_TIMELINE_OVERRIDE"
|
||||
|
||||
// RequestInfo holds essential details of an incoming HTTP request for logging purposes.
|
||||
type RequestInfo struct {
|
||||
URL string // URL is the request URL.
|
||||
Method string // Method is the HTTP method (e.g., GET, POST).
|
||||
Headers map[string][]string // Headers contains the request headers.
|
||||
Body []byte // Body is the raw request body.
|
||||
RequestID string // RequestID is the unique identifier for the request.
|
||||
Timestamp time.Time // Timestamp is when the request was received.
|
||||
URL string // URL is the request URL.
|
||||
Method string // Method is the HTTP method (e.g., GET or POST).
|
||||
Headers map[string][]string // Headers contains the request headers.
|
||||
Body []byte // Body is the raw request body.
|
||||
RequestID string // RequestID is the unique identifier for the request.
|
||||
Timestamp time.Time // Timestamp is when the request was received.
|
||||
deferredBodyCapture *deferredRequestBodyCapture // deferredBodyCapture spools large error-only request bodies.
|
||||
}
|
||||
|
||||
// ResponseWriterWrapper wraps the standard gin.ResponseWriter to intercept and log response data.
|
||||
@@ -258,6 +259,9 @@ func (w *ResponseWriterWrapper) processStreamingChunks(done chan struct{}) {
|
||||
// For non-streaming responses, it logs the complete request and response details,
|
||||
// including any API-specific request/response data stored in the Gin context.
|
||||
func (w *ResponseWriterWrapper) Finalize(c *gin.Context) error {
|
||||
if w.requestInfo != nil && w.requestInfo.deferredBodyCapture != nil {
|
||||
defer w.requestInfo.deferredBodyCapture.Cleanup()
|
||||
}
|
||||
if w.logger == nil {
|
||||
return nil
|
||||
}
|
||||
@@ -361,7 +365,11 @@ func (w *ResponseWriterWrapper) Finalize(c *gin.Context) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
return w.logRequest(w.extractRequestBody(c), finalStatusCode, w.cloneHeaders(), w.extractResponseBody(c), w.extractWebsocketTimeline(c), websocketTimelineSource, w.extractAPIRequest(c), apiRequestSource, w.extractAPIResponse(c), apiResponseSource, w.extractAPIWebsocketTimeline(c), apiWebsocketTimelineSource, w.extractAPIResponseTimestamp(c), slicesAPIResponseError, forceLog)
|
||||
apiRequest := w.extractAPIRequest(c)
|
||||
if forceLog && len(apiRequest) == 0 {
|
||||
apiRequest = w.extractDeferredAPIRequest(c)
|
||||
}
|
||||
return w.logRequest(w.extractRequestBody(c), finalStatusCode, w.cloneHeaders(), w.extractResponseBody(c), w.extractWebsocketTimeline(c), websocketTimelineSource, apiRequest, apiRequestSource, w.extractAPIResponse(c), apiResponseSource, w.extractAPIWebsocketTimeline(c), apiWebsocketTimelineSource, w.extractAPIResponseTimestamp(c), slicesAPIResponseError, forceLog)
|
||||
}
|
||||
|
||||
func (w *ResponseWriterWrapper) cloneHeaders() map[string][]string {
|
||||
@@ -389,6 +397,28 @@ func (w *ResponseWriterWrapper) extractAPIRequest(c *gin.Context) []byte {
|
||||
return data
|
||||
}
|
||||
|
||||
func (w *ResponseWriterWrapper) extractDeferredAPIRequest(c *gin.Context) []byte {
|
||||
if c == nil {
|
||||
return nil
|
||||
}
|
||||
value, exists := c.Get(logging.DeferredAPIRequestContextKey)
|
||||
if !exists {
|
||||
return nil
|
||||
}
|
||||
requests, ok := value.([]logging.DeferredAPIRequest)
|
||||
if !ok || len(requests) == 0 {
|
||||
return nil
|
||||
}
|
||||
var body bytes.Buffer
|
||||
for _, buildRequest := range requests {
|
||||
if buildRequest == nil {
|
||||
continue
|
||||
}
|
||||
body.Write(buildRequest())
|
||||
}
|
||||
return body.Bytes()
|
||||
}
|
||||
|
||||
func (w *ResponseWriterWrapper) extractAPIResponse(c *gin.Context) []byte {
|
||||
apiResponse, isExist := c.Get("API_RESPONSE")
|
||||
if !isExist {
|
||||
@@ -440,10 +470,35 @@ func (w *ResponseWriterWrapper) extractRequestBody(c *gin.Context) []byte {
|
||||
if body := extractBodyOverride(c, requestBodyOverrideContextKey); len(body) > 0 {
|
||||
return body
|
||||
}
|
||||
if w.requestInfo != nil && len(w.requestInfo.Body) > 0 {
|
||||
if w.requestInfo == nil {
|
||||
return nil
|
||||
}
|
||||
if len(w.requestInfo.Body) > 0 {
|
||||
return w.requestInfo.Body
|
||||
}
|
||||
return nil
|
||||
if w.requestInfo.deferredBodyCapture == nil {
|
||||
return nil
|
||||
}
|
||||
body, statusMarker, errRead := w.requestInfo.deferredBodyCapture.Bytes()
|
||||
if errRead != nil {
|
||||
log.WithError(errRead).Warn("failed to read deferred request body capture")
|
||||
return nil
|
||||
}
|
||||
encoding := ""
|
||||
for key, values := range w.requestInfo.Headers {
|
||||
if strings.EqualFold(key, "Content-Encoding") && len(values) > 0 {
|
||||
encoding = values[0]
|
||||
break
|
||||
}
|
||||
}
|
||||
body = decodeCapturedRequestBodyForLogWithLimit(body, encoding, maxDeferredErrorRequestBodyBytes)
|
||||
if statusMarker == "" {
|
||||
return body
|
||||
}
|
||||
if len(body) > 0 && !bytes.HasSuffix(body, []byte("\n")) {
|
||||
body = append(body, '\n')
|
||||
}
|
||||
return append(body, statusMarker...)
|
||||
}
|
||||
|
||||
func (w *ResponseWriterWrapper) extractResponseBody(c *gin.Context) []byte {
|
||||
|
||||
@@ -36,11 +36,15 @@ var requestLogID atomic.Uint64
|
||||
const (
|
||||
WebsocketTimelineSourceContextKey = "WEBSOCKET_TIMELINE_SOURCE"
|
||||
APIRequestSourceContextKey = "API_REQUEST_SOURCE"
|
||||
DeferredAPIRequestContextKey = "DEFERRED_API_REQUEST"
|
||||
APIResponseSourceContextKey = "API_RESPONSE_SOURCE"
|
||||
APIResponseCapturedContextKey = "API_RESPONSE_CAPTURED"
|
||||
APIWebsocketTimelineSourceContextKey = "API_WEBSOCKET_TIMELINE_SOURCE"
|
||||
)
|
||||
|
||||
// DeferredAPIRequest builds an upstream request log only when an error log needs it.
|
||||
type DeferredAPIRequest func() []byte
|
||||
|
||||
type homeRequestLogClient interface {
|
||||
HeartbeatOK() bool
|
||||
RPushRequestLog(ctx context.Context, payload []byte) error
|
||||
|
||||
@@ -20,11 +20,13 @@ import (
|
||||
)
|
||||
|
||||
const (
|
||||
apiAttemptsKey = "API_UPSTREAM_ATTEMPTS"
|
||||
apiRequestKey = "API_REQUEST"
|
||||
apiResponseKey = "API_RESPONSE"
|
||||
apiWebsocketTimelineKey = "API_WEBSOCKET_TIMELINE"
|
||||
creditsUsedKey = "__antigravity_credits_used__"
|
||||
apiAttemptsKey = "API_UPSTREAM_ATTEMPTS"
|
||||
apiRequestKey = "API_REQUEST"
|
||||
apiResponseKey = "API_RESPONSE"
|
||||
apiWebsocketTimelineKey = "API_WEBSOCKET_TIMELINE"
|
||||
deferredAPIRequestBytesKey = "DEFERRED_API_REQUEST_BYTES"
|
||||
creditsUsedKey = "__antigravity_credits_used__"
|
||||
maxDeferredAPIRequestBodyBytes = 32 << 20 // 32 MiB
|
||||
)
|
||||
|
||||
// UpstreamRequestLog captures the outbound upstream request details for logging.
|
||||
@@ -60,34 +62,21 @@ func requestLogCaptureEnabled(cfg *config.Config) bool {
|
||||
|
||||
// RecordAPIRequest stores the upstream request metadata in Gin context for request logging.
|
||||
func RecordAPIRequest(ctx context.Context, cfg *config.Config, info UpstreamRequestLog) {
|
||||
if !requestLogCaptureEnabled(cfg) {
|
||||
if cfg == nil || cfg.CommercialMode {
|
||||
return
|
||||
}
|
||||
ginCtx := ginContextFrom(ctx)
|
||||
if ginCtx == nil {
|
||||
return
|
||||
}
|
||||
if !cfg.RequestLog {
|
||||
deferAPIRequest(ginCtx, info)
|
||||
return
|
||||
}
|
||||
|
||||
attempts := getAttempts(ginCtx)
|
||||
index := len(attempts) + 1
|
||||
|
||||
builder := &strings.Builder{}
|
||||
builder.WriteString(fmt.Sprintf("=== API REQUEST %d ===\n", index))
|
||||
builder.WriteString(fmt.Sprintf("Timestamp: %s\n", time.Now().Format(time.RFC3339Nano)))
|
||||
if info.URL != "" {
|
||||
builder.WriteString(fmt.Sprintf("Upstream URL: %s\n", info.URL))
|
||||
} else {
|
||||
builder.WriteString("Upstream URL: <unknown>\n")
|
||||
}
|
||||
if info.Method != "" {
|
||||
builder.WriteString(fmt.Sprintf("HTTP Method: %s\n", info.Method))
|
||||
}
|
||||
if auth := formatAuthInfo(info); auth != "" {
|
||||
builder.WriteString(fmt.Sprintf("Auth: %s\n", auth))
|
||||
}
|
||||
builder.WriteString("\nHeaders:\n")
|
||||
writeHeaders(builder, info.Headers)
|
||||
builder.WriteString("\nBody:\n")
|
||||
builder := newAPIRequestLogBuilder(index, info, time.Now())
|
||||
|
||||
requestText := ""
|
||||
if source, ok := apiRequestSource(ginCtx); ok {
|
||||
@@ -135,6 +124,68 @@ func RecordAPIRequest(ctx context.Context, cfg *config.Config, info UpstreamRequ
|
||||
}
|
||||
}
|
||||
|
||||
func deferAPIRequest(ginCtx *gin.Context, info UpstreamRequestLog) {
|
||||
if ginCtx == nil {
|
||||
return
|
||||
}
|
||||
var requests []logging.DeferredAPIRequest
|
||||
if value, exists := ginCtx.Get(logging.DeferredAPIRequestContextKey); exists {
|
||||
requests, _ = value.([]logging.DeferredAPIRequest)
|
||||
}
|
||||
index := len(requests) + 1
|
||||
capturedInfo := info
|
||||
capturedAt := time.Now()
|
||||
capturedBytes, _ := ginCtx.Get(deferredAPIRequestBytesKey)
|
||||
bytesUsed, _ := capturedBytes.(int)
|
||||
remaining := maxDeferredAPIRequestBodyBytes - bytesUsed
|
||||
if remaining < 0 {
|
||||
remaining = 0
|
||||
}
|
||||
captureLength := len(info.Body)
|
||||
if captureLength > remaining {
|
||||
captureLength = remaining
|
||||
}
|
||||
capturedInfo.Body = bytes.Clone(info.Body[:captureLength])
|
||||
bodyEmpty := len(info.Body) == 0
|
||||
bodyTruncated := captureLength < len(info.Body)
|
||||
ginCtx.Set(deferredAPIRequestBytesKey, bytesUsed+captureLength)
|
||||
requests = append(requests, func() []byte {
|
||||
builder := newAPIRequestLogBuilder(index, capturedInfo, capturedAt)
|
||||
if bodyEmpty {
|
||||
builder.WriteString("<empty>")
|
||||
} else {
|
||||
builder.Write(capturedInfo.Body)
|
||||
if bodyTruncated {
|
||||
builder.WriteString(fmt.Sprintf("\n[API REQUEST BODY TRUNCATED: captured first %d bytes]", captureLength))
|
||||
}
|
||||
}
|
||||
builder.WriteString("\n\n")
|
||||
return []byte(builder.String())
|
||||
})
|
||||
ginCtx.Set(logging.DeferredAPIRequestContextKey, requests)
|
||||
}
|
||||
|
||||
func newAPIRequestLogBuilder(index int, info UpstreamRequestLog, timestamp time.Time) *strings.Builder {
|
||||
builder := &strings.Builder{}
|
||||
builder.WriteString(fmt.Sprintf("=== API REQUEST %d ===\n", index))
|
||||
builder.WriteString(fmt.Sprintf("Timestamp: %s\n", timestamp.Format(time.RFC3339Nano)))
|
||||
if info.URL != "" {
|
||||
builder.WriteString(fmt.Sprintf("Upstream URL: %s\n", info.URL))
|
||||
} else {
|
||||
builder.WriteString("Upstream URL: <unknown>\n")
|
||||
}
|
||||
if info.Method != "" {
|
||||
builder.WriteString(fmt.Sprintf("HTTP Method: %s\n", info.Method))
|
||||
}
|
||||
if auth := formatAuthInfo(info); auth != "" {
|
||||
builder.WriteString(fmt.Sprintf("Auth: %s\n", auth))
|
||||
}
|
||||
builder.WriteString("\nHeaders:\n")
|
||||
writeHeaders(builder, info.Headers)
|
||||
builder.WriteString("\nBody:\n")
|
||||
return builder
|
||||
}
|
||||
|
||||
// RecordAPIResponseMetadata captures upstream response status/header information for the latest attempt.
|
||||
func RecordAPIResponseMetadata(ctx context.Context, cfg *config.Config, status int, headers http.Header) {
|
||||
logging.SetResponseHeaders(ctx, headers)
|
||||
|
||||
@@ -3,12 +3,43 @@ package helps
|
||||
import (
|
||||
"context"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/config"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/logging"
|
||||
)
|
||||
|
||||
func TestRecordAPIRequestClonesDeferredBodyWhenRequestLogDisabled(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
recorder := httptest.NewRecorder()
|
||||
ginCtx, _ := gin.CreateTestContext(recorder)
|
||||
ctx := context.WithValue(context.Background(), "gin", ginCtx)
|
||||
body := []byte(`{"model":"original"}`)
|
||||
|
||||
RecordAPIRequest(ctx, &config.Config{}, UpstreamRequestLog{
|
||||
URL: "https://api.example.com/v1/responses",
|
||||
Method: http.MethodPost,
|
||||
Body: body,
|
||||
})
|
||||
body[10] = 'X'
|
||||
|
||||
value, exists := ginCtx.Get(logging.DeferredAPIRequestContextKey)
|
||||
if !exists {
|
||||
t.Fatal("deferred API request was not captured")
|
||||
}
|
||||
requests, ok := value.([]logging.DeferredAPIRequest)
|
||||
if !ok || len(requests) != 1 {
|
||||
t.Fatalf("deferred API requests = %#v, want one request", value)
|
||||
}
|
||||
captured := string(requests[0]())
|
||||
if !strings.Contains(captured, `{"model":"original"}`) {
|
||||
t.Fatalf("captured API request = %q, want original body", captured)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRecordAPIResponseMetadataStoresHeadersWhenRequestLogDisabled(t *testing.T) {
|
||||
ctx := logging.WithResponseHeadersHolder(context.Background())
|
||||
headers := http.Header{}
|
||||
|
||||
Reference in New Issue
Block a user