zhanlu_proxy/internal/server/recorder.go

package server

import (
	"bytes"
	"fmt"
	"io"
	"net/http"
	"sync"
	"time"
)

const (
	maxRecentEntries = 10
	maxBodyCapture    = 1 << 20 // 1 MB per body — full request/response capture for debugging
)

// RecentEntry captures a single API request and its response for debugging.
type RecentEntry struct {
	Timestamp      time.Time         `json:"timestamp"`
	Method         string            `json:"method"`
	Path           string            `json:"path"`
	Status         int               `json:"status"`
	DurationMs     int64             `json:"duration_ms"`
	RequestHeaders map[string]string `json:"request_headers"`
	RequestBody    string            `json:"request_body"`
	ResponseBody   string            `json:"response_body"`
}

// RecentRecorder is an in-memory ring buffer that stores the most recent API
// requests. It is safe for concurrent use.
type RecentRecorder struct {
	mu      sync.Mutex
	entries []RecentEntry
}

func NewRecentRecorder() *RecentRecorder {
	return &RecentRecorder{}
}

// Record appends an entry, evicting the oldest when the buffer is full.
func (r *RecentRecorder) Record(e RecentEntry) {
	r.mu.Lock()
	defer r.mu.Unlock()
	r.entries = append(r.entries, e)
	if len(r.entries) > maxRecentEntries {
		r.entries = r.entries[len(r.entries)-maxRecentEntries:]
	}
}

// Entries returns a copy of the buffer in newest-first order.
func (r *RecentRecorder) Entries() []RecentEntry {
	r.mu.Lock()
	defer r.mu.Unlock()
	n := len(r.entries)
	out := make([]RecentEntry, n)
	for i, e := range r.entries {
		out[n-1-i] = e // reverse: newest first
	}
	return out
}

// recordingResponseWriter wraps http.ResponseWriter to capture the status code
// and a truncated copy of the response body. It implements http.Flusher so
// streaming handlers can flush through the wrapper.
type recordingResponseWriter struct {
	http.ResponseWriter
	statusCode  int
	body        bytes.Buffer
	totalBytes  int
}

func newRecordingResponseWriter(w http.ResponseWriter) *recordingResponseWriter {
	return &recordingResponseWriter{ResponseWriter: w, statusCode: http.StatusOK}
}

func (w *recordingResponseWriter) WriteHeader(code int) {
	w.statusCode = code
	w.ResponseWriter.WriteHeader(code)
}

func (w *recordingResponseWriter) Write(b []byte) (int, error) {
	w.totalBytes += len(b)
	if w.body.Len() < maxBodyCapture {
		remaining := maxBodyCapture - w.body.Len()
		if len(b) <= remaining {
			w.body.Write(b)
		} else {
			w.body.Write(b[:remaining])
		}
	}
	return w.ResponseWriter.Write(b)
}

func (w *recordingResponseWriter) Flush() {
	if f, ok := w.ResponseWriter.(http.Flusher); ok {
		f.Flush()
	}
}

// captureRequestHeaders extracts selected request headers, masking the
// Authorization value to avoid leaking API keys.
func captureRequestHeaders(r *http.Request) map[string]string {
	headers := map[string]string{}
	for _, key := range []string{"Content-Type", "User-Agent", "Accept", "Authorization"} {
		if v := r.Header.Get(key); v != "" {
			if key == "Authorization" {
				headers[key] = mask(v)
			} else {
				headers[key] = v
			}
		}
	}
	return headers
}

// truncateBody returns the string form of b, truncated to maxBodyCapture
// bytes with a marker if the original was longer.
func truncateBody(b []byte) string {
	if len(b) > maxBodyCapture {
		return string(b[:maxBodyCapture]) + fmt.Sprintf("\n...(truncated, total %d bytes)", len(b))
	}
	return string(b)
}

// formatResponseBody returns the captured response body, with a truncation
// marker if the full response exceeded the capture limit.
func formatResponseBody(buf *bytes.Buffer, totalBytes int) string {
	s := buf.String()
	if totalBytes > maxBodyCapture {
		s += fmt.Sprintf("\n...(truncated, total %d bytes)", totalBytes)
	}
	return s
}

// withRecording wraps a handler so that each request's headers, body,
// response status, and response body (truncated) are captured into the
// server's RecentRecorder. It is applied to the /v1/ API routes so that
// both successful and error responses are recorded.
func (s *Server) withRecording(next http.HandlerFunc) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		start := time.Now()

		// Read and restore the request body so the downstream handler
		// still sees the full content.
		var reqBody []byte
		if r.Body != nil {
			reqBody, _ = io.ReadAll(r.Body)
			r.Body = io.NopCloser(bytes.NewReader(reqBody))
		}

		recW := newRecordingResponseWriter(w)
		next(recW, r)

		s.recorder.Record(RecentEntry{
			Timestamp:      start,
			Method:         r.Method,
			Path:           r.URL.Path,
			Status:         recW.statusCode,
			DurationMs:     time.Since(start).Milliseconds(),
			RequestHeaders: captureRequestHeaders(r),
			RequestBody:    truncateBody(reqBody),
			ResponseBody:   formatResponseBody(&recW.body, recW.totalBytes),
		})
	}
}

// getRecent handles GET /api/recent — returns the most recent API requests
// as JSON, newest first.
func (s *Server) getRecent(w http.ResponseWriter, r *http.Request) {
	writeJSON(w, http.StatusOK, map[string]any{
		"ok":       true,
		"requests": s.recorder.Entries(),
	})
}