m1saka/token_thief
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
m1saka
2026-07-11 17:23:32 +08:00
parent d39f5a1048
commit 9e6e7ca3c7
21 changed files with 417 additions and 83 deletions
Download Patch File Download Diff File Expand all files Collapse all files
proxy/capture.go
+11 -2
View File
@@ -5,21 +5,30 @@ import (
"crypto/rand"
"encoding/hex"
"encoding/json"
"errors"
"io"
"net/http"
"net/netip"
"strings"
)
var errRequestBodyTooLarge = errors.New("request body too large")
// readRequestBody 在转发前完整读取请求体，确保读取失败时不会向上游发送损坏请求。
func readRequestBody(r *http.Request, max int64) (captured []byte, truncated bool, err error) {
func readRequestBody(r *http.Request, max, requestLimit int64) (captured []byte, truncated bool, err error) {
if r.Body == nil || r.Body == http.NoBody {
return nil, false, nil
}
body, err := io.ReadAll(r.Body)
if r.ContentLength > requestLimit {
return nil, false, errRequestBodyTooLarge
}
body, err := io.ReadAll(io.LimitReader(r.Body, requestLimit+1))
if err != nil {
return nil, false, err
}
if int64(len(body)) > requestLimit {
return nil, false, errRequestBodyTooLarge
}
if err := r.Body.Close(); err != nil {
return nil, false, err
}
proxy/proxy.go
+47 -23
View File
@@ -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
}
proxy/sse.go
+63 -17
View File
@@ -552,11 +552,16 @@ func (g *geminiChunk) UnmarshalJSON(data []byte) error {
}
type geminiCandidate struct {
raw map[string]json.RawMessage
index int
role string
text strings.Builder
partRaw map[string]json.RawMessage
raw map[string]json.RawMessage
index int
role string
parts []*geminiPart
}
type geminiPart struct {
raw map[string]json.RawMessage
text strings.Builder
kind string
}
func assembleGeminiSSE(payloads []string) ([]byte, bool) {
@@ -588,15 +593,31 @@ func assembleGeminiSSE(payloads []string) ([]byte, bool) {
assembled.role = candidate.Content.Role
}
if len(candidate.Content.Parts) > 0 {
for _, part := range candidate.Content.Parts {
assembled.text.WriteString(part.Text)
var contentRaw struct {
Parts []map[string]json.RawMessage `json:"parts"`
}
if len(assembled.partRaw) == 0 {
var contentRaw struct {
Parts []map[string]json.RawMessage `json:"parts"`
if rawContent, ok := candidate.Raw["content"]; ok {
_ = json.Unmarshal(rawContent, &contentRaw)
}
mergeByIndex := len(assembled.parts) == len(candidate.Content.Parts)
if mergeByIndex {
for i := range candidate.Content.Parts {
if i >= len(contentRaw.Parts) || assembled.parts[i].kind != geminiPartKind(contentRaw.Parts[i]) {
mergeByIndex = false
break
}
}
if rawContent, ok := candidate.Raw["content"]; ok && json.Unmarshal(rawContent, &contentRaw) == nil && len(contentRaw.Parts) > 0 {
assembled.partRaw = cloneRawMap(contentRaw.Parts[0])
}
for i, part := range candidate.Content.Parts {
target := i
if !mergeByIndex {
target = len(assembled.parts)
assembled.parts = append(assembled.parts, &geminiPart{})
}
assembled.parts[target].text.WriteString(part.Text)
if i < len(contentRaw.Parts) {
assembled.parts[target].kind = geminiPartKind(contentRaw.Parts[i])
assembled.parts[target].raw = mergeRawMap(assembled.parts[target].raw, contentRaw.Parts[i])
}
}
}
@@ -614,12 +635,15 @@ func assembleGeminiSSE(payloads []string) ([]byte, bool) {
candidate := candidates[index]
candidateRaw := cloneRawMap(candidate.raw)
candidateRaw["index"] = mustJSON(candidate.index)
contentRaw := map[string]any{"role": candidate.role, "parts": []any{map[string]any{"text": candidate.text.String()}}}
if len(candidate.partRaw) > 0 {
partRaw := cloneRawMap(candidate.partRaw)
partRaw["text"] = mustJSON(candidate.text.String())
contentRaw["parts"] = []any{rawMapToMap(partRaw)}
parts := make([]any, 0, len(candidate.parts))
for _, part := range candidate.parts {
partRaw := cloneRawMap(part.raw)
if part.text.Len() > 0 || len(partRaw) == 0 {
partRaw["text"] = mustJSON(part.text.String())
}
parts = append(parts, rawMapToMap(partRaw))
}
contentRaw := map[string]any{"role": candidate.role, "parts": parts}
candidateRaw["content"] = mustJSON(contentRaw)
assembled["candidates"] = append(assembled["candidates"].([]any), rawMapToMap(candidateRaw))
}
@@ -631,6 +655,28 @@ func assembleGeminiSSE(payloads []string) ([]byte, bool) {
return data, err == nil
}
func geminiPartKind(part map[string]json.RawMessage) string {
for _, key := range []string{"functionCall", "functionResponse", "inlineData", "fileData", "executableCode", "codeExecutionResult", "text"} {
if _, ok := part[key]; ok {
return key
}
}
return "unknown"
}
func mergeRawMap(dst, src map[string]json.RawMessage) map[string]json.RawMessage {
if dst == nil {
dst = make(map[string]json.RawMessage, len(src))
}
for key, value := range src {
if key == "text" {
continue
}
dst[key] = append(json.RawMessage(nil), value...)
}
return dst
}
func cloneRawMap(in map[string]json.RawMessage) map[string]json.RawMessage {
out := make(map[string]json.RawMessage, len(in))
for k, v := range in {
proxy/sse_event.go
+18 -1
View File
@@ -8,12 +8,29 @@ import (
type sseEventTracker struct {
buf []byte
limit int
overflow bool
recognized bool
terminal bool
afterTerminal bool
}
func newSSEEventTracker(limit int64) sseEventTracker {
if limit > int64(^uint(0)>>1) {
limit = int64(^uint(0) >> 1)
}
return sseEventTracker{limit: int(limit)}
}
func (t *sseEventTracker) Write(p []byte) {
if t.overflow {
return
}
if t.limit > 0 && len(p) > t.limit-len(t.buf) {
t.buf = nil
t.overflow = true
return
}
t.buf = append(t.buf, p...)
for {
end, separator := completeSSEEvent(t.buf)
@@ -34,7 +51,7 @@ func (t *sseEventTracker) Write(p []byte) {
}
func (t *sseEventTracker) Complete() bool {
return len(t.buf) == 0 && !t.afterTerminal && (!t.recognized || t.terminal)
return !t.overflow && len(t.buf) == 0 && !t.afterTerminal && (!t.recognized || t.terminal)
}
func completeSSEEvent(buf []byte) (int, int) {
proxy/sse_event_test.go
+12
View File
@@ -19,3 +19,15 @@ func TestSSEEventTrackerAcceptsCROnlyEventBoundary(t *testing.T) {
t.Fatal("terminal SSE event with CR-only boundary should be complete")
}
}
func TestSSEEventTrackerStopsBufferingOverLimit(t *testing.T) {
tracker := newSSEEventTracker(8)
tracker.Write([]byte("data: 123456789"))
if tracker.Complete() {
t.Fatal("overflowed SSE tracker must not report a complete stream")
}
if len(tracker.buf) != 0 || !tracker.overflow {
t.Fatalf("overflow state=%v buffered=%d, want overflow with released buffer", tracker.overflow, len(tracker.buf))
}
}
proxy/timeout.go
+18
View File
@@ -2,12 +2,30 @@ package proxy
import (
"io"
"net"
"net/http"
"sync"
"sync/atomic"
"time"
)
type writeTimeoutConn struct {
net.Conn
timeout time.Duration
}
func (c *writeTimeoutConn) Write(p []byte) (int, error) {
if err := c.Conn.SetWriteDeadline(time.Now().Add(c.timeout)); err != nil {
return 0, err
}
n, err := c.Conn.Write(p)
clearErr := c.Conn.SetWriteDeadline(time.Time{})
if err == nil {
err = clearErr
}
return n, err
}
type trackingBody struct {
io.ReadCloser
failed *atomic.Bool
proxy/writer.go
+6 -3
View File
@@ -26,7 +26,7 @@ type captureWriter struct {
}
func newCaptureWriter(w http.ResponseWriter, max int64) *captureWriter {
return &captureWriter{ResponseWriter: w, max: max, status: http.StatusOK}
return &captureWriter{ResponseWriter: w, max: max, status: http.StatusOK, sse: newSSEEventTracker(max)}
}
func (c *captureWriter) WriteHeader(code int) {
@@ -63,7 +63,7 @@ func (c *captureWriter) Write(p []byte) (int, error) {
c.truncated = true
}
c.written += int64(n)
if isStreamResponse(c.Header()) {
if c.max > 0 && isStreamResponse(c.Header()) {
c.sse.Write(p[:n])
}
c.flush()
@@ -78,7 +78,7 @@ func (c *captureWriter) Flush() {
// ConfirmDelivery establishes an observable delivery boundary for responses
// whose headers were not followed by a body write.
func (c *captureWriter) ConfirmDelivery() {
if c.written != 0 {
if c.hijacked || c.written != 0 {
return
}
if _, ok := c.ResponseWriter.(http.Flusher); !ok {
@@ -89,6 +89,9 @@ func (c *captureWriter) ConfirmDelivery() {
}
func (c *captureWriter) flush() {
if c.hijacked {
return
}
if _, ok := c.ResponseWriter.(http.Flusher); !ok {
return
}