修复代理与日志链路的可靠性问题

proxy/proxy.go

@@ -3,6 +3,7 @@ package proxy
 import (
 	"context"
 	"crypto/tls"
+	"errors"
 	"log"
 	"net"
 	"net/http"
@@ -25,16 +26,17 @@ type LogSubmitter interface {
 
 // Handler 构造反代 HTTP handler。
 type Handler struct {
-	rp           *httputil.ReverseProxy
-	filter       *config.Filter
-	queue        LogSubmitter
-	maxBodyBytes int64
-	trusted      []netip.Prefix
-	connMu       sync.Mutex
-	conns        map[net.Conn]struct{}
-	connChanged  chan struct{}
-	closing      bool
-	closed       bool
+	rp              *httputil.ReverseProxy
+	filter          *config.Filter
+	queue           LogSubmitter
+	maxBodyBytes    int64
+	maxRequestBytes int64
+	trusted         []netip.Prefix
+	connMu          sync.Mutex
+	conns           map[net.Conn]struct{}
+	connChanged     chan struct{}
+	closing         bool
+	closed          bool
 }
 
 // Options 控制反代连接上游时的网络行为。
@@ -44,6 +46,7 @@ type Options struct {
 	SSEIdleTimeout                time.Duration
 	UpstreamTLSInsecureSkipVerify bool
 	TrustedProxies                []netip.Prefix
+	MaxRequestBytes               int64
 }
 
 // requestState 通过 context 在 ErrorHandler / ModifyResponse / 主 handler 之间共享状态。
@@ -77,12 +80,22 @@ func New(upstream *url.URL, filter *config.Filter, queue LogSubmitter, maxBody i
 }
 
 func NewWithOptions(upstream *url.URL, filter *config.Filter, queue LogSubmitter, maxBody int64, opts Options) *Handler {
+	if opts.MaxRequestBytes <= 0 {
+		opts.MaxRequestBytes = 16 << 20
+	}
 	rp := httputil.NewSingleHostReverseProxy(upstream)
 	rp.FlushInterval = -1 // 让流式 chunk 立即转发
 	if opts.UpstreamTimeout > 0 || opts.UpstreamTLSInsecureSkipVerify {
 		transport := http.DefaultTransport.(*http.Transport).Clone()
 		if opts.UpstreamTimeout > 0 {
-			transport.DialContext = (&net.Dialer{Timeout: opts.UpstreamTimeout, KeepAlive: 30 * time.Second}).DialContext
+			dialer := &net.Dialer{Timeout: opts.UpstreamTimeout, KeepAlive: 30 * time.Second}
+			transport.DialContext = func(ctx context.Context, network, address string) (net.Conn, error) {
+				conn, err := dialer.DialContext(ctx, network, address)
+				if err != nil {
+					return nil, err
+				}
+				return &writeTimeoutConn{Conn: conn, timeout: opts.UpstreamTimeout}, nil
+			}
 			transport.ResponseHeaderTimeout = opts.UpstreamTimeout
 			transport.TLSHandshakeTimeout = opts.UpstreamTimeout
 		}
@@ -130,13 +143,14 @@ func NewWithOptions(upstream *url.URL, filter *config.Filter, queue LogSubmitter
 	}
 
 	return &Handler{
-		rp:           rp,
-		filter:       filter,
-		queue:        queue,
-		maxBodyBytes: maxBody,
-		trusted:      append([]netip.Prefix(nil), opts.TrustedProxies...),
-		conns:        make(map[net.Conn]struct{}),
-		connChanged:  make(chan struct{}),
+		rp:              rp,
+		filter:          filter,
+		queue:           queue,
+		maxBodyBytes:    maxBody,
+		maxRequestBytes: opts.MaxRequestBytes,
+		trusted:         append([]netip.Prefix(nil), opts.TrustedProxies...),
+		conns:           make(map[net.Conn]struct{}),
+		connChanged:     make(chan struct{}),
 	}
 }
 
@@ -149,11 +163,15 @@ func (h *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	}
 
 	shouldLog := h.filter.ShouldLog(r.URL.Path)
-	reqBody, reqTruncated, err := readRequestBody(r, h.maxBodyBytes)
+	reqBody, reqTruncated, err := readRequestBody(r, h.maxBodyBytes, h.maxRequestBytes)
 	if err != nil && !shouldLog {
 		log.Printf("[proxy] read request body failed: %v", err)
 		_ = r.Body.Close()
-		http.Error(w, "bad request", http.StatusBadRequest)
+		if errors.Is(err, errRequestBodyTooLarge) {
+			http.Error(w, "request body too large", http.StatusRequestEntityTooLarge)
+		} else {
+			http.Error(w, "bad request", http.StatusBadRequest)
+		}
 		return
 	}
 	if !shouldLog {
@@ -176,7 +194,7 @@ func (h *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 		_ = r.Body.Close()
 		s := "read request body: " + err.Error()
 		st.lastErr.Store(&s)
-		h.serveRequestBodyError(cw, r, st, started, reqID)
+		h.serveRequestBodyError(cw, r, st, started, reqID, errors.Is(err, errRequestBodyTooLarge))
 		return
 	}
 
@@ -230,9 +248,15 @@ func (h *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	}
 }
 
-func (h *Handler) serveRequestBodyError(cw *captureWriter, r *http.Request, st *requestState, started time.Time, requestID string) {
+func (h *Handler) serveRequestBodyError(cw *captureWriter, r *http.Request, st *requestState, started time.Time, requestID string, tooLarge bool) {
 	cw.Header().Set("X-Request-Id", requestID)
-	http.Error(cw, "bad request", http.StatusBadRequest)
+	status := http.StatusBadRequest
+	message := "bad request"
+	if tooLarge {
+		status = http.StatusRequestEntityTooLarge
+		message = "request body too large"
+	}
+	http.Error(cw, message, status)
 	if !cw.Complete() {
 		return
 	}