fix: restore reviewable migration evidence
This commit is contained in:
MiMoCode
@@ -0,0 +1,104 @@
|
||||
package proxy
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"crypto/rand"
|
||||
"encoding/hex"
|
||||
"encoding/json"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/netip"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// readRequestBody 在转发前完整读取请求体，确保读取失败时不会向上游发送损坏请求。
|
||||
func readRequestBody(r *http.Request, max int64) (captured []byte, truncated bool, err error) {
|
||||
if r.Body == nil || r.ContentLength == 0 {
|
||||
return nil, false, nil
|
||||
}
|
||||
body, err := io.ReadAll(r.Body)
|
||||
if err != nil {
|
||||
return nil, false, err
|
||||
}
|
||||
if err := r.Body.Close(); err != nil {
|
||||
return nil, false, err
|
||||
}
|
||||
r.Body = io.NopCloser(bytes.NewReader(body))
|
||||
if int64(len(body)) > max {
|
||||
return body[:max], true, nil
|
||||
}
|
||||
return body, false, nil
|
||||
}
|
||||
|
||||
func headersJSON(h http.Header) []byte {
|
||||
if len(h) == 0 {
|
||||
return nil
|
||||
}
|
||||
b, err := json.Marshal(h)
|
||||
if err != nil {
|
||||
return nil
|
||||
}
|
||||
return b
|
||||
}
|
||||
|
||||
func newRequestID() string {
|
||||
var b [16]byte
|
||||
if _, err := rand.Read(b[:]); err != nil {
|
||||
return "unknown"
|
||||
}
|
||||
return hex.EncodeToString(b[:])
|
||||
}
|
||||
|
||||
// clientIP 从请求中提取客户端 IP。
|
||||
func clientIP(r *http.Request, trusted []netip.Prefix) string {
|
||||
peer, ok := parsePeerAddr(r.RemoteAddr)
|
||||
if !ok {
|
||||
return r.RemoteAddr
|
||||
}
|
||||
if !isTrusted(peer, trusted) {
|
||||
return peer.String()
|
||||
}
|
||||
xff := strings.Split(r.Header.Get("X-Forwarded-For"), ",")
|
||||
if len(xff) == 1 && strings.TrimSpace(xff[0]) == "" {
|
||||
if realIP, err := netip.ParseAddr(strings.TrimSpace(r.Header.Get("X-Real-IP"))); err == nil {
|
||||
return realIP.Unmap().String()
|
||||
}
|
||||
return peer.String()
|
||||
}
|
||||
chain := make([]netip.Addr, len(xff))
|
||||
for i, raw := range xff {
|
||||
addr, err := netip.ParseAddr(strings.TrimSpace(raw))
|
||||
if err != nil {
|
||||
return peer.String()
|
||||
}
|
||||
chain[i] = addr.Unmap()
|
||||
}
|
||||
client := peer
|
||||
for i := len(chain) - 1; i >= 0 && isTrusted(client, trusted); i-- {
|
||||
client = chain[i]
|
||||
}
|
||||
return client.String()
|
||||
}
|
||||
|
||||
func parsePeerAddr(remote string) (netip.Addr, bool) {
|
||||
if addrPort, err := netip.ParseAddrPort(remote); err == nil {
|
||||
return addrPort.Addr().Unmap(), true
|
||||
}
|
||||
addr, err := netip.ParseAddr(remote)
|
||||
return addr.Unmap(), err == nil
|
||||
}
|
||||
|
||||
func isTrusted(addr netip.Addr, prefixes []netip.Prefix) bool {
|
||||
for _, prefix := range prefixes {
|
||||
if prefix.Contains(addr) {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// isStreamResponse 通过响应头判断是否为流式响应。
|
||||
func isStreamResponse(h http.Header) bool {
|
||||
ct := strings.ToLower(strings.TrimSpace(strings.SplitN(h.Get("Content-Type"), ";", 2)[0]))
|
||||
return ct == "text/event-stream"
|
||||
}
|
||||
+339
@@ -0,0 +1,339 @@
|
||||
package proxy
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/tls"
|
||||
"log"
|
||||
"net"
|
||||
"net/http"
|
||||
"net/http/httputil"
|
||||
"net/netip"
|
||||
"net/url"
|
||||
"strings"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
|
||||
"git.misaka.ren/M1saka/token_thief/config"
|
||||
"git.misaka.ren/M1saka/token_thief/logger"
|
||||
)
|
||||
|
||||
// LogSubmitter 是 proxy 唯一依赖的日志接收方接口。
|
||||
type LogSubmitter interface {
|
||||
Submit(*logger.LogEntry)
|
||||
}
|
||||
|
||||
// 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{}
|
||||
}
|
||||
|
||||
// Options 控制反代连接上游时的网络行为。
|
||||
type Options struct {
|
||||
UpstreamTimeout time.Duration
|
||||
ResponseTimeout time.Duration
|
||||
SSEIdleTimeout time.Duration
|
||||
UpstreamTLSInsecureSkipVerify bool
|
||||
TrustedProxies []netip.Prefix
|
||||
}
|
||||
|
||||
// requestState 通过 context 在 ErrorHandler / ModifyResponse / 主 handler 之间共享状态。
|
||||
type requestState struct {
|
||||
requestID string
|
||||
lastErr atomic.Pointer[string]
|
||||
upgrade atomic.Pointer[upgradeResponse]
|
||||
readFailed atomic.Bool
|
||||
}
|
||||
|
||||
type upgradeResponse struct {
|
||||
status int
|
||||
header http.Header
|
||||
}
|
||||
|
||||
type ctxKey struct{}
|
||||
|
||||
func newRequestState(id string) *requestState { return &requestState{requestID: id} }
|
||||
|
||||
func stateFromCtx(ctx context.Context) *requestState {
|
||||
v, _ := ctx.Value(ctxKey{}).(*requestState)
|
||||
return v
|
||||
}
|
||||
|
||||
func New(upstream *url.URL, filter *config.Filter, queue LogSubmitter, maxBody int64, upstreamTimeout ...time.Duration) *Handler {
|
||||
opts := Options{}
|
||||
if len(upstreamTimeout) > 0 {
|
||||
opts.UpstreamTimeout = upstreamTimeout[0]
|
||||
}
|
||||
return NewWithOptions(upstream, filter, queue, maxBody, opts)
|
||||
}
|
||||
|
||||
func NewWithOptions(upstream *url.URL, filter *config.Filter, queue LogSubmitter, maxBody int64, opts Options) *Handler {
|
||||
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
|
||||
transport.ResponseHeaderTimeout = opts.UpstreamTimeout
|
||||
transport.TLSHandshakeTimeout = opts.UpstreamTimeout
|
||||
}
|
||||
if opts.UpstreamTLSInsecureSkipVerify {
|
||||
transport.TLSClientConfig = &tls.Config{InsecureSkipVerify: true}
|
||||
}
|
||||
rp.Transport = transport
|
||||
}
|
||||
|
||||
origDirector := rp.Director
|
||||
rp.Director = func(r *http.Request) {
|
||||
origDirector(r)
|
||||
r.Host = upstream.Host
|
||||
}
|
||||
|
||||
// ModifyResponse 在响应头写回客户端之前调用，确保 X-Request-Id 一定生效。
|
||||
rp.ModifyResponse = func(resp *http.Response) error {
|
||||
if st := stateFromCtx(resp.Request.Context()); st != nil {
|
||||
resp.Header.Set("X-Request-Id", st.requestID)
|
||||
if resp.StatusCode == http.StatusSwitchingProtocols &&
|
||||
strings.EqualFold(resp.Request.Header.Get("Upgrade"), "websocket") &&
|
||||
strings.EqualFold(resp.Header.Get("Upgrade"), "websocket") {
|
||||
st.upgrade.Store(&upgradeResponse{status: resp.StatusCode, header: resp.Header.Clone()})
|
||||
}
|
||||
}
|
||||
if timeout := responseBodyTimeout(resp, opts); timeout > 0 {
|
||||
resp.Body = newTimeoutBody(resp.Body, timeout, isStreamResponse(resp.Header))
|
||||
}
|
||||
if st := stateFromCtx(resp.Request.Context()); st != nil && resp.StatusCode != http.StatusSwitchingProtocols {
|
||||
resp.Body = &trackingBody{ReadCloser: resp.Body, failed: &st.readFailed}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
rp.ErrorHandler = func(w http.ResponseWriter, r *http.Request, err error) {
|
||||
log.Printf("[proxy] upstream error for %s %s: %v", r.Method, r.URL.Path, err)
|
||||
// 把错误暴露给主 handler，使其能写入日志。
|
||||
if st := stateFromCtx(r.Context()); st != nil {
|
||||
s := err.Error()
|
||||
st.lastErr.Store(&s)
|
||||
// ErrorHandler 路径下 ModifyResponse 不会被调用，这里手动写 X-Request-Id。
|
||||
w.Header().Set("X-Request-Id", st.requestID)
|
||||
}
|
||||
http.Error(w, "bad gateway", http.StatusBadGateway)
|
||||
}
|
||||
|
||||
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{}),
|
||||
}
|
||||
}
|
||||
|
||||
func (h *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
|
||||
// 健康检查不参与反代与日志。
|
||||
if r.URL.Path == "/healthz" {
|
||||
w.WriteHeader(http.StatusOK)
|
||||
_, _ = w.Write([]byte("ok"))
|
||||
return
|
||||
}
|
||||
|
||||
shouldLog := h.filter.ShouldLog(r.URL.Path)
|
||||
reqBody, reqTruncated, err := readRequestBody(r, h.maxBodyBytes)
|
||||
if err != nil && !shouldLog {
|
||||
log.Printf("[proxy] read request body failed: %v", err)
|
||||
_ = r.Body.Close()
|
||||
http.Error(w, "bad request", http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
if !shouldLog {
|
||||
log.Printf("[proxy] skip log by filter method=%s path=%s", r.Method, r.URL.Path)
|
||||
cw := newCaptureWriter(w, 0)
|
||||
cw.OnHijack(h.trackConn)
|
||||
h.rp.ServeHTTP(cw, r)
|
||||
return
|
||||
}
|
||||
|
||||
started := time.Now()
|
||||
reqID := newRequestID()
|
||||
log.Printf("[proxy] capture start request_id=%s method=%s path=%s", reqID, r.Method, r.URL.Path)
|
||||
st := newRequestState(reqID)
|
||||
r = r.WithContext(context.WithValue(r.Context(), ctxKey{}, st))
|
||||
cw := newCaptureWriter(w, h.maxBodyBytes)
|
||||
|
||||
if err != nil {
|
||||
log.Printf("[proxy] read request body failed: %v", err)
|
||||
_ = r.Body.Close()
|
||||
s := "read request body: " + err.Error()
|
||||
st.lastErr.Store(&s)
|
||||
h.serveRequestBodyError(cw, r, st, started, reqID)
|
||||
return
|
||||
}
|
||||
|
||||
reqHeadersJSON := headersJSON(r.Header)
|
||||
clientAddr := clientIP(r, h.trusted)
|
||||
method := r.Method
|
||||
path := r.URL.Path
|
||||
query := r.URL.RawQuery
|
||||
|
||||
var submitOnce sync.Once
|
||||
submit := func(finished time.Time) {
|
||||
entry, ok := h.buildLogEntry(cw, st, logEntryInput{
|
||||
requestID: reqID,
|
||||
method: method,
|
||||
path: path,
|
||||
query: query,
|
||||
clientAddr: clientAddr,
|
||||
requestHeaders: reqHeadersJSON,
|
||||
requestBody: reqBody,
|
||||
requestTruncated: reqTruncated,
|
||||
started: started,
|
||||
finished: finished,
|
||||
})
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
h.queue.Submit(entry)
|
||||
log.Printf("[proxy] capture finish request_id=%s method=%s path=%s status=%d is_stream=%v latency_ms=%d",
|
||||
reqID, method, path, entry.StatusCode, entry.IsStream, entry.LatencyMS)
|
||||
}
|
||||
cw.OnHijack(func(conn net.Conn) net.Conn {
|
||||
tracked := h.trackConn(conn)
|
||||
if upgrade := st.upgrade.Load(); upgrade != nil {
|
||||
cw.SetHijackedResponse(upgrade.status, upgrade.header)
|
||||
submitOnce.Do(func() { submit(time.Now()) })
|
||||
}
|
||||
return tracked
|
||||
})
|
||||
h.rp.ServeHTTP(cw, r)
|
||||
|
||||
finished := time.Now()
|
||||
responseComplete := !isStreamResponse(cw.Header()) || cw.SSEComplete()
|
||||
if cw.Complete() && responseComplete && !st.readFailed.Load() && r.Context().Err() == nil {
|
||||
submitOnce.Do(func() { submit(finished) })
|
||||
}
|
||||
}
|
||||
|
||||
func (h *Handler) serveRequestBodyError(cw *captureWriter, r *http.Request, st *requestState, started time.Time, requestID string) {
|
||||
cw.Header().Set("X-Request-Id", requestID)
|
||||
http.Error(cw, "bad request", http.StatusBadRequest)
|
||||
if !cw.Complete() {
|
||||
return
|
||||
}
|
||||
entry, ok := h.buildLogEntry(cw, st, logEntryInput{
|
||||
requestID: requestID,
|
||||
method: r.Method,
|
||||
path: r.URL.Path,
|
||||
query: r.URL.RawQuery,
|
||||
clientAddr: clientIP(r, h.trusted),
|
||||
requestHeaders: headersJSON(r.Header),
|
||||
started: started,
|
||||
finished: time.Now(),
|
||||
})
|
||||
if ok {
|
||||
h.queue.Submit(entry)
|
||||
}
|
||||
}
|
||||
|
||||
// Shutdown closes all active hijacked connections and waits for their release.
|
||||
func (h *Handler) Shutdown(ctx context.Context) error {
|
||||
for {
|
||||
h.connMu.Lock()
|
||||
if len(h.conns) == 0 {
|
||||
h.connMu.Unlock()
|
||||
return nil
|
||||
}
|
||||
conns := make([]net.Conn, 0, len(h.conns))
|
||||
for conn := range h.conns {
|
||||
conns = append(conns, conn)
|
||||
}
|
||||
changed := h.connChanged
|
||||
h.connMu.Unlock()
|
||||
for _, conn := range conns {
|
||||
_ = conn.Close()
|
||||
}
|
||||
select {
|
||||
case <-changed:
|
||||
case <-ctx.Done():
|
||||
return ctx.Err()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (h *Handler) trackConn(conn net.Conn) net.Conn {
|
||||
tracked := &trackedConn{Conn: conn}
|
||||
tracked.onClose = func() {
|
||||
h.connMu.Lock()
|
||||
delete(h.conns, tracked)
|
||||
close(h.connChanged)
|
||||
h.connChanged = make(chan struct{})
|
||||
h.connMu.Unlock()
|
||||
}
|
||||
h.connMu.Lock()
|
||||
h.conns[tracked] = struct{}{}
|
||||
close(h.connChanged)
|
||||
h.connChanged = make(chan struct{})
|
||||
h.connMu.Unlock()
|
||||
return tracked
|
||||
}
|
||||
|
||||
type logEntryInput struct {
|
||||
requestID string
|
||||
method string
|
||||
path string
|
||||
query string
|
||||
clientAddr string
|
||||
requestHeaders []byte
|
||||
requestBody []byte
|
||||
requestTruncated bool
|
||||
started time.Time
|
||||
finished time.Time
|
||||
}
|
||||
|
||||
func (h *Handler) buildLogEntry(cw *captureWriter, st *requestState, in logEntryInput) (*logger.LogEntry, bool) {
|
||||
if cw.Hijacked() && cw.Status() != http.StatusSwitchingProtocols {
|
||||
return nil, false
|
||||
}
|
||||
|
||||
var errMsg string
|
||||
if p := st.lastErr.Load(); p != nil {
|
||||
errMsg = *p
|
||||
}
|
||||
|
||||
isStream := isStreamResponse(cw.Header())
|
||||
responseBody := append([]byte(nil), cw.Body()...)
|
||||
responseTruncated := cw.Truncated()
|
||||
if isStream && !responseTruncated {
|
||||
if assembled, ok := assembleSSEJSON(responseBody); ok {
|
||||
responseBody = assembled
|
||||
}
|
||||
}
|
||||
|
||||
return &logger.LogEntry{
|
||||
RequestID: in.requestID,
|
||||
Method: in.method,
|
||||
Path: in.path,
|
||||
Query: in.query,
|
||||
ClientIP: in.clientAddr,
|
||||
RequestHeaders: in.requestHeaders,
|
||||
RequestBody: in.requestBody,
|
||||
RequestTruncated: in.requestTruncated,
|
||||
StatusCode: cw.Status(),
|
||||
ResponseHeaders: headersJSON(cw.Header()),
|
||||
ResponseBody: responseBody,
|
||||
ResponseTruncated: responseTruncated,
|
||||
IsStream: isStream,
|
||||
LatencyMS: in.finished.Sub(in.started).Milliseconds(),
|
||||
StartedAt: in.started,
|
||||
FinishedAt: in.finished,
|
||||
Error: errMsg,
|
||||
}, true
|
||||
}
|
||||
+661
@@ -0,0 +1,661 @@
|
||||
package proxy
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"bytes"
|
||||
"encoding/json"
|
||||
"strings"
|
||||
)
|
||||
|
||||
func assembleSSEJSON(body []byte) ([]byte, bool) {
|
||||
payloads := sseDataPayloads(body)
|
||||
if len(payloads) == 0 {
|
||||
return nil, false
|
||||
}
|
||||
if assembled, ok := assembleOpenAICompletionsSSE(payloads); ok {
|
||||
return assembled, true
|
||||
}
|
||||
if assembled, ok := assembleOpenAIChatSSE(payloads); ok {
|
||||
return assembled, true
|
||||
}
|
||||
if assembled, ok := assembleOpenAIResponsesSSE(payloads); ok {
|
||||
return assembled, true
|
||||
}
|
||||
if assembled, ok := assembleAnthropicSSE(payloads); ok {
|
||||
return assembled, true
|
||||
}
|
||||
if assembled, ok := assembleGeminiSSE(payloads); ok {
|
||||
return assembled, true
|
||||
}
|
||||
return nil, false
|
||||
}
|
||||
|
||||
func sseDataPayloads(body []byte) []string {
|
||||
payloads := make([]string, 0)
|
||||
scanner := bufio.NewScanner(bytes.NewReader(body))
|
||||
scanner.Buffer(make([]byte, 0, 64*1024), len(body)+1)
|
||||
for scanner.Scan() {
|
||||
line := strings.TrimSpace(scanner.Text())
|
||||
if !strings.HasPrefix(line, "data:") {
|
||||
continue
|
||||
}
|
||||
payload := strings.TrimSpace(strings.TrimPrefix(line, "data:"))
|
||||
if payload == "" || payload == "[DONE]" {
|
||||
continue
|
||||
}
|
||||
payloads = append(payloads, payload)
|
||||
}
|
||||
return payloads
|
||||
}
|
||||
|
||||
type openAICompletionChunk struct {
|
||||
ID string `json:"id,omitempty"`
|
||||
Object string `json:"object,omitempty"`
|
||||
Created int64 `json:"created,omitempty"`
|
||||
Model string `json:"model,omitempty"`
|
||||
Usage json.RawMessage `json:"usage"`
|
||||
Choices []struct {
|
||||
Index int `json:"index"`
|
||||
Text *string `json:"text"`
|
||||
FinishReason *string `json:"finish_reason"`
|
||||
} `json:"choices"`
|
||||
}
|
||||
|
||||
type openAICompletionChoice struct {
|
||||
index int
|
||||
text strings.Builder
|
||||
finishReason string
|
||||
}
|
||||
|
||||
func assembleOpenAICompletionsSSE(payloads []string) ([]byte, bool) {
|
||||
choices := map[int]*openAICompletionChoice{}
|
||||
order := make([]int, 0, 1)
|
||||
var id, object, model string
|
||||
var created int64
|
||||
var usage json.RawMessage
|
||||
matched := false
|
||||
|
||||
for _, payload := range payloads {
|
||||
var chunk openAICompletionChunk
|
||||
if err := json.Unmarshal([]byte(payload), &chunk); err != nil {
|
||||
return nil, false
|
||||
}
|
||||
if len(chunk.Choices) == 0 && len(chunk.Usage) == 0 {
|
||||
return nil, false
|
||||
}
|
||||
for _, choice := range chunk.Choices {
|
||||
if choice.Text == nil {
|
||||
return nil, false
|
||||
}
|
||||
}
|
||||
matched = true
|
||||
if id == "" {
|
||||
id = chunk.ID
|
||||
}
|
||||
if object == "" {
|
||||
object = chunk.Object
|
||||
}
|
||||
if created == 0 {
|
||||
created = chunk.Created
|
||||
}
|
||||
if model == "" {
|
||||
model = chunk.Model
|
||||
}
|
||||
if len(chunk.Usage) > 0 && string(chunk.Usage) != "null" {
|
||||
usage = chunk.Usage
|
||||
}
|
||||
|
||||
for _, choice := range chunk.Choices {
|
||||
assembled := choices[choice.Index]
|
||||
if assembled == nil {
|
||||
assembled = &openAICompletionChoice{index: choice.Index}
|
||||
choices[choice.Index] = assembled
|
||||
order = append(order, choice.Index)
|
||||
}
|
||||
assembled.text.WriteString(*choice.Text)
|
||||
if choice.FinishReason != nil {
|
||||
assembled.finishReason = *choice.FinishReason
|
||||
}
|
||||
}
|
||||
}
|
||||
if !matched {
|
||||
return nil, false
|
||||
}
|
||||
|
||||
assembled := struct {
|
||||
ID string `json:"id,omitempty"`
|
||||
Object string `json:"object,omitempty"`
|
||||
Created int64 `json:"created,omitempty"`
|
||||
Model string `json:"model,omitempty"`
|
||||
Usage json.RawMessage `json:"usage,omitempty"`
|
||||
Choices []struct {
|
||||
Index int `json:"index"`
|
||||
Text string `json:"text"`
|
||||
FinishReason string `json:"finish_reason,omitempty"`
|
||||
} `json:"choices"`
|
||||
}{ID: id, Object: object, Created: created, Model: model, Usage: usage}
|
||||
for _, index := range order {
|
||||
choice := choices[index]
|
||||
assembled.Choices = append(assembled.Choices, struct {
|
||||
Index int `json:"index"`
|
||||
Text string `json:"text"`
|
||||
FinishReason string `json:"finish_reason,omitempty"`
|
||||
}{Index: choice.index, Text: choice.text.String(), FinishReason: choice.finishReason})
|
||||
}
|
||||
|
||||
data, err := json.Marshal(assembled)
|
||||
return data, err == nil
|
||||
}
|
||||
|
||||
type openAIChatChunk struct {
|
||||
ID string `json:"id,omitempty"`
|
||||
Object string `json:"object,omitempty"`
|
||||
Created int64 `json:"created,omitempty"`
|
||||
Model string `json:"model,omitempty"`
|
||||
SystemFingerprint string `json:"system_fingerprint,omitempty"`
|
||||
Usage json.RawMessage `json:"usage"`
|
||||
Choices []struct {
|
||||
Index int `json:"index"`
|
||||
Delta struct {
|
||||
Role string `json:"role,omitempty"`
|
||||
Content string `json:"content,omitempty"`
|
||||
ReasoningContent string `json:"reasoning_content,omitempty"`
|
||||
ToolCalls []struct {
|
||||
Index int `json:"index"`
|
||||
ID string `json:"id,omitempty"`
|
||||
Type string `json:"type,omitempty"`
|
||||
Function struct {
|
||||
Name string `json:"name,omitempty"`
|
||||
Arguments string `json:"arguments,omitempty"`
|
||||
} `json:"function,omitempty"`
|
||||
} `json:"tool_calls,omitempty"`
|
||||
} `json:"delta"`
|
||||
FinishReason *string `json:"finish_reason"`
|
||||
NativeFinishReason *string `json:"native_finish_reason"`
|
||||
} `json:"choices"`
|
||||
}
|
||||
|
||||
type openAIChatChoice struct {
|
||||
index int
|
||||
role string
|
||||
content strings.Builder
|
||||
reasoning strings.Builder
|
||||
toolCalls map[int]*openAIToolCall
|
||||
toolOrder []int
|
||||
finishReason string
|
||||
nativeFinish string
|
||||
}
|
||||
|
||||
type openAIToolCall struct {
|
||||
id string
|
||||
callType string
|
||||
name string
|
||||
arguments strings.Builder
|
||||
}
|
||||
|
||||
type openAIChatResponse struct {
|
||||
ID string `json:"id,omitempty"`
|
||||
Object string `json:"object,omitempty"`
|
||||
Created int64 `json:"created,omitempty"`
|
||||
Model string `json:"model,omitempty"`
|
||||
SystemFingerprint string `json:"system_fingerprint,omitempty"`
|
||||
Usage json.RawMessage `json:"usage,omitempty"`
|
||||
Choices []openAIResponseChoice `json:"choices"`
|
||||
}
|
||||
|
||||
type openAIResponseChoice struct {
|
||||
Index int `json:"index"`
|
||||
Message openAIResponseMessage `json:"message"`
|
||||
FinishReason string `json:"finish_reason,omitempty"`
|
||||
NativeFinishReason string `json:"native_finish_reason,omitempty"`
|
||||
}
|
||||
|
||||
type openAIResponseMessage struct {
|
||||
Role string `json:"role,omitempty"`
|
||||
Content string `json:"content"`
|
||||
ReasoningContent string `json:"reasoning_content,omitempty"`
|
||||
ToolCalls []openAIResponseTool `json:"tool_calls,omitempty"`
|
||||
}
|
||||
|
||||
type openAIResponseTool struct {
|
||||
ID string `json:"id,omitempty"`
|
||||
Type string `json:"type,omitempty"`
|
||||
Function openAIResponseToolFunction `json:"function"`
|
||||
}
|
||||
|
||||
type openAIResponseToolFunction struct {
|
||||
Name string `json:"name,omitempty"`
|
||||
Arguments string `json:"arguments"`
|
||||
}
|
||||
|
||||
func assembleOpenAIChatSSE(payloads []string) ([]byte, bool) {
|
||||
choices := map[int]*openAIChatChoice{}
|
||||
order := make([]int, 0, 1)
|
||||
var id, object string
|
||||
var created int64
|
||||
var model string
|
||||
var systemFingerprint string
|
||||
var usage json.RawMessage
|
||||
matched := false
|
||||
|
||||
for _, payload := range payloads {
|
||||
var chunk openAIChatChunk
|
||||
if err := json.Unmarshal([]byte(payload), &chunk); err != nil {
|
||||
return nil, false
|
||||
}
|
||||
if len(chunk.Choices) == 0 && len(chunk.Usage) == 0 {
|
||||
return nil, false
|
||||
}
|
||||
matched = true
|
||||
if id == "" {
|
||||
id = chunk.ID
|
||||
}
|
||||
if object == "" {
|
||||
object = strings.TrimSuffix(chunk.Object, ".chunk")
|
||||
}
|
||||
if created == 0 {
|
||||
created = chunk.Created
|
||||
}
|
||||
if model == "" {
|
||||
model = chunk.Model
|
||||
}
|
||||
if systemFingerprint == "" {
|
||||
systemFingerprint = chunk.SystemFingerprint
|
||||
}
|
||||
if len(chunk.Usage) > 0 && string(chunk.Usage) != "null" {
|
||||
usage = chunk.Usage
|
||||
}
|
||||
|
||||
for _, choice := range chunk.Choices {
|
||||
assembled := choices[choice.Index]
|
||||
if assembled == nil {
|
||||
assembled = &openAIChatChoice{index: choice.Index}
|
||||
choices[choice.Index] = assembled
|
||||
order = append(order, choice.Index)
|
||||
}
|
||||
if choice.Delta.Role != "" {
|
||||
assembled.role = choice.Delta.Role
|
||||
}
|
||||
if choice.Delta.Content != "" {
|
||||
assembled.content.WriteString(choice.Delta.Content)
|
||||
}
|
||||
if choice.Delta.ReasoningContent != "" {
|
||||
assembled.reasoning.WriteString(choice.Delta.ReasoningContent)
|
||||
}
|
||||
for _, toolCall := range choice.Delta.ToolCalls {
|
||||
if assembled.toolCalls == nil {
|
||||
assembled.toolCalls = map[int]*openAIToolCall{}
|
||||
}
|
||||
assembledTool := assembled.toolCalls[toolCall.Index]
|
||||
if assembledTool == nil {
|
||||
assembledTool = &openAIToolCall{}
|
||||
assembled.toolCalls[toolCall.Index] = assembledTool
|
||||
assembled.toolOrder = append(assembled.toolOrder, toolCall.Index)
|
||||
}
|
||||
if toolCall.ID != "" {
|
||||
assembledTool.id = toolCall.ID
|
||||
}
|
||||
if toolCall.Type != "" {
|
||||
assembledTool.callType = toolCall.Type
|
||||
}
|
||||
if toolCall.Function.Name != "" {
|
||||
assembledTool.name = toolCall.Function.Name
|
||||
}
|
||||
if toolCall.Function.Arguments != "" {
|
||||
assembledTool.arguments.WriteString(toolCall.Function.Arguments)
|
||||
}
|
||||
}
|
||||
if choice.FinishReason != nil {
|
||||
assembled.finishReason = *choice.FinishReason
|
||||
}
|
||||
if choice.NativeFinishReason != nil {
|
||||
assembled.nativeFinish = *choice.NativeFinishReason
|
||||
}
|
||||
}
|
||||
}
|
||||
if !matched {
|
||||
return nil, false
|
||||
}
|
||||
|
||||
assembled := openAIChatResponse{ID: id, Object: object, Created: created, Model: model, SystemFingerprint: systemFingerprint, Usage: usage}
|
||||
for _, index := range order {
|
||||
choice := choices[index]
|
||||
out := openAIResponseChoice{Index: choice.index, FinishReason: choice.finishReason, NativeFinishReason: choice.nativeFinish}
|
||||
out.Message.Role = choice.role
|
||||
out.Message.Content = choice.content.String()
|
||||
out.Message.ReasoningContent = choice.reasoning.String()
|
||||
for _, toolIndex := range choice.toolOrder {
|
||||
toolCall := choice.toolCalls[toolIndex]
|
||||
outTool := openAIResponseTool{ID: toolCall.id, Type: toolCall.callType}
|
||||
outTool.Function.Name = toolCall.name
|
||||
outTool.Function.Arguments = toolCall.arguments.String()
|
||||
out.Message.ToolCalls = append(out.Message.ToolCalls, outTool)
|
||||
}
|
||||
assembled.Choices = append(assembled.Choices, out)
|
||||
}
|
||||
|
||||
data, err := json.Marshal(assembled)
|
||||
return data, err == nil
|
||||
}
|
||||
|
||||
type openAIResponsesEvent struct {
|
||||
Type string `json:"type"`
|
||||
Response json.RawMessage `json:"response"`
|
||||
}
|
||||
|
||||
func assembleOpenAIResponsesSSE(payloads []string) ([]byte, bool) {
|
||||
var completed json.RawMessage
|
||||
matched := false
|
||||
for _, payload := range payloads {
|
||||
var event openAIResponsesEvent
|
||||
if err := json.Unmarshal([]byte(payload), &event); err != nil {
|
||||
return nil, false
|
||||
}
|
||||
if !strings.HasPrefix(event.Type, "response.") {
|
||||
return nil, false
|
||||
}
|
||||
matched = true
|
||||
if event.Type == "response.completed" && len(event.Response) > 0 {
|
||||
completed = append(json.RawMessage(nil), event.Response...)
|
||||
}
|
||||
}
|
||||
if !matched || len(completed) == 0 {
|
||||
return nil, false
|
||||
}
|
||||
return completed, true
|
||||
}
|
||||
|
||||
type anthropicEvent struct {
|
||||
Type string `json:"type"`
|
||||
Message *struct {
|
||||
ID string `json:"id"`
|
||||
Type string `json:"type"`
|
||||
Role string `json:"role"`
|
||||
Model string `json:"model,omitempty"`
|
||||
StopReason string `json:"stop_reason"`
|
||||
StopSequence string `json:"stop_sequence"`
|
||||
Usage struct {
|
||||
InputTokens int `json:"input_tokens"`
|
||||
OutputTokens int `json:"output_tokens"`
|
||||
} `json:"usage"`
|
||||
} `json:"message"`
|
||||
Index int `json:"index"`
|
||||
ContentBlock *struct {
|
||||
Type string `json:"type"`
|
||||
Text string `json:"text"`
|
||||
ID string `json:"id"`
|
||||
Name string `json:"name"`
|
||||
} `json:"content_block"`
|
||||
Delta *struct {
|
||||
StopReason string `json:"stop_reason"`
|
||||
StopSequence string `json:"stop_sequence"`
|
||||
Text string `json:"text"`
|
||||
PartialJSON string `json:"partial_json"`
|
||||
} `json:"delta"`
|
||||
Usage *struct {
|
||||
OutputTokens int `json:"output_tokens"`
|
||||
} `json:"usage"`
|
||||
}
|
||||
|
||||
type anthropicContentBlock struct {
|
||||
Type string `json:"type"`
|
||||
Text string `json:"text,omitempty"`
|
||||
ID string `json:"id,omitempty"`
|
||||
Name string `json:"name,omitempty"`
|
||||
Input map[string]any `json:"input,omitempty"`
|
||||
|
||||
partialJSON strings.Builder
|
||||
}
|
||||
|
||||
func assembleAnthropicSSE(payloads []string) ([]byte, bool) {
|
||||
var assembled struct {
|
||||
ID string `json:"id"`
|
||||
Type string `json:"type"`
|
||||
Role string `json:"role"`
|
||||
Content []anthropicContentBlock `json:"content"`
|
||||
Model string `json:"model,omitempty"`
|
||||
StopReason string `json:"stop_reason,omitempty"`
|
||||
StopSequence string `json:"stop_sequence,omitempty"`
|
||||
Usage struct {
|
||||
InputTokens int `json:"input_tokens"`
|
||||
OutputTokens int `json:"output_tokens"`
|
||||
} `json:"usage"`
|
||||
}
|
||||
contents := map[int]*anthropicContentBlock{}
|
||||
order := make([]int, 0, 1)
|
||||
matched := false
|
||||
|
||||
for _, payload := range payloads {
|
||||
var event anthropicEvent
|
||||
if err := json.Unmarshal([]byte(payload), &event); err != nil {
|
||||
return nil, false
|
||||
}
|
||||
if !strings.HasPrefix(event.Type, "message_") && !strings.HasPrefix(event.Type, "content_block_") {
|
||||
return nil, false
|
||||
}
|
||||
matched = true
|
||||
|
||||
if event.Message != nil {
|
||||
assembled.ID = event.Message.ID
|
||||
assembled.Type = event.Message.Type
|
||||
assembled.Role = event.Message.Role
|
||||
assembled.Model = event.Message.Model
|
||||
assembled.StopReason = event.Message.StopReason
|
||||
assembled.StopSequence = event.Message.StopSequence
|
||||
assembled.Usage.InputTokens = event.Message.Usage.InputTokens
|
||||
assembled.Usage.OutputTokens = event.Message.Usage.OutputTokens
|
||||
}
|
||||
if event.ContentBlock != nil {
|
||||
block := contents[event.Index]
|
||||
if block == nil {
|
||||
block = &anthropicContentBlock{Type: event.ContentBlock.Type, ID: event.ContentBlock.ID, Name: event.ContentBlock.Name}
|
||||
contents[event.Index] = block
|
||||
order = append(order, event.Index)
|
||||
}
|
||||
if event.ContentBlock.ID != "" {
|
||||
block.ID = event.ContentBlock.ID
|
||||
}
|
||||
if event.ContentBlock.Name != "" {
|
||||
block.Name = event.ContentBlock.Name
|
||||
}
|
||||
block.Text += event.ContentBlock.Text
|
||||
}
|
||||
if event.Delta != nil {
|
||||
if event.Delta.Text != "" {
|
||||
block := contents[event.Index]
|
||||
if block == nil {
|
||||
block = &anthropicContentBlock{Type: "text"}
|
||||
contents[event.Index] = block
|
||||
order = append(order, event.Index)
|
||||
}
|
||||
block.Text += event.Delta.Text
|
||||
}
|
||||
if event.Delta.PartialJSON != "" {
|
||||
block := contents[event.Index]
|
||||
if block == nil {
|
||||
block = &anthropicContentBlock{Type: "tool_use"}
|
||||
contents[event.Index] = block
|
||||
order = append(order, event.Index)
|
||||
}
|
||||
block.partialJSON.WriteString(event.Delta.PartialJSON)
|
||||
}
|
||||
if event.Delta.StopReason != "" {
|
||||
assembled.StopReason = event.Delta.StopReason
|
||||
}
|
||||
if event.Delta.StopSequence != "" {
|
||||
assembled.StopSequence = event.Delta.StopSequence
|
||||
}
|
||||
}
|
||||
if event.Usage != nil {
|
||||
assembled.Usage.OutputTokens = event.Usage.OutputTokens
|
||||
}
|
||||
}
|
||||
if !matched || assembled.Type == "" {
|
||||
return nil, false
|
||||
}
|
||||
for _, index := range order {
|
||||
block := contents[index]
|
||||
if block.partialJSON.Len() > 0 {
|
||||
var input map[string]any
|
||||
if err := json.Unmarshal([]byte(block.partialJSON.String()), &input); err != nil {
|
||||
return nil, false
|
||||
}
|
||||
block.Input = input
|
||||
}
|
||||
assembled.Content = append(assembled.Content, *block)
|
||||
}
|
||||
|
||||
data, err := json.Marshal(assembled)
|
||||
return data, err == nil
|
||||
}
|
||||
|
||||
type geminiChunk struct {
|
||||
Raw map[string]json.RawMessage `json:"-"`
|
||||
Candidates []struct {
|
||||
Raw map[string]json.RawMessage `json:"-"`
|
||||
Content struct {
|
||||
Parts []struct {
|
||||
Text string `json:"text"`
|
||||
} `json:"parts"`
|
||||
Role string `json:"role"`
|
||||
} `json:"content"`
|
||||
FinishReason string `json:"finishReason"`
|
||||
Index int `json:"index"`
|
||||
} `json:"candidates"`
|
||||
UsageMetadata json.RawMessage `json:"usageMetadata"`
|
||||
}
|
||||
|
||||
func (g *geminiChunk) UnmarshalJSON(data []byte) error {
|
||||
type alias geminiChunk
|
||||
var decoded alias
|
||||
if err := json.Unmarshal(data, &decoded); err != nil {
|
||||
return err
|
||||
}
|
||||
var raw map[string]json.RawMessage
|
||||
if err := json.Unmarshal(data, &raw); err != nil {
|
||||
return err
|
||||
}
|
||||
*g = geminiChunk(decoded)
|
||||
g.Raw = raw
|
||||
if candidatesRaw, ok := raw["candidates"]; ok {
|
||||
var rawCandidates []map[string]json.RawMessage
|
||||
if err := json.Unmarshal(candidatesRaw, &rawCandidates); err != nil {
|
||||
return err
|
||||
}
|
||||
for i := range g.Candidates {
|
||||
if i < len(rawCandidates) {
|
||||
g.Candidates[i].Raw = rawCandidates[i]
|
||||
}
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
type geminiCandidate struct {
|
||||
raw map[string]json.RawMessage
|
||||
index int
|
||||
role string
|
||||
text strings.Builder
|
||||
partRaw map[string]json.RawMessage
|
||||
}
|
||||
|
||||
func assembleGeminiSSE(payloads []string) ([]byte, bool) {
|
||||
candidates := map[int]*geminiCandidate{}
|
||||
order := make([]int, 0, 1)
|
||||
var usage json.RawMessage
|
||||
matched := false
|
||||
|
||||
for _, payload := range payloads {
|
||||
var chunk geminiChunk
|
||||
if err := json.Unmarshal([]byte(payload), &chunk); err != nil {
|
||||
return nil, false
|
||||
}
|
||||
if len(chunk.Candidates) == 0 {
|
||||
return nil, false
|
||||
}
|
||||
matched = true
|
||||
for _, candidate := range chunk.Candidates {
|
||||
assembled := candidates[candidate.Index]
|
||||
if assembled == nil {
|
||||
assembled = &geminiCandidate{index: candidate.Index}
|
||||
candidates[candidate.Index] = assembled
|
||||
order = append(order, candidate.Index)
|
||||
}
|
||||
if candidate.Raw != nil {
|
||||
assembled.raw = cloneRawMap(candidate.Raw)
|
||||
}
|
||||
if candidate.Content.Role != "" {
|
||||
assembled.role = candidate.Content.Role
|
||||
}
|
||||
if len(candidate.Content.Parts) > 0 {
|
||||
for _, part := range candidate.Content.Parts {
|
||||
assembled.text.WriteString(part.Text)
|
||||
}
|
||||
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) == nil && len(contentRaw.Parts) > 0 {
|
||||
assembled.partRaw = cloneRawMap(contentRaw.Parts[0])
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
if len(chunk.UsageMetadata) > 0 {
|
||||
usage = chunk.UsageMetadata
|
||||
}
|
||||
}
|
||||
if !matched {
|
||||
return nil, false
|
||||
}
|
||||
|
||||
assembled := map[string]any{"candidates": make([]any, 0, len(order))}
|
||||
for _, index := range order {
|
||||
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)}
|
||||
}
|
||||
candidateRaw["content"] = mustJSON(contentRaw)
|
||||
assembled["candidates"] = append(assembled["candidates"].([]any), rawMapToMap(candidateRaw))
|
||||
}
|
||||
if len(usage) > 0 {
|
||||
assembled["usageMetadata"] = usage
|
||||
}
|
||||
|
||||
data, err := json.Marshal(assembled)
|
||||
return data, err == nil
|
||||
}
|
||||
|
||||
func cloneRawMap(in map[string]json.RawMessage) map[string]json.RawMessage {
|
||||
out := make(map[string]json.RawMessage, len(in))
|
||||
for k, v := range in {
|
||||
out[k] = append(json.RawMessage(nil), v...)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func mustJSON(v any) json.RawMessage {
|
||||
data, err := json.Marshal(v)
|
||||
if err != nil {
|
||||
return nil
|
||||
}
|
||||
return data
|
||||
}
|
||||
|
||||
func rawMapToMap(in map[string]json.RawMessage) map[string]any {
|
||||
out := make(map[string]any, len(in))
|
||||
for k, v := range in {
|
||||
var decoded any
|
||||
if err := json.Unmarshal(v, &decoded); err != nil {
|
||||
out[k] = string(v)
|
||||
continue
|
||||
}
|
||||
out[k] = decoded
|
||||
}
|
||||
return out
|
||||
}
|
||||
@@ -0,0 +1,65 @@
|
||||
package proxy
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/json"
|
||||
"strings"
|
||||
)
|
||||
|
||||
type sseEventTracker struct {
|
||||
buf []byte
|
||||
terminal bool
|
||||
}
|
||||
|
||||
func (t *sseEventTracker) Write(p []byte) {
|
||||
t.buf = append(t.buf, p...)
|
||||
for {
|
||||
end, separator := completeSSEEvent(t.buf)
|
||||
if end < 0 {
|
||||
return
|
||||
}
|
||||
event := t.buf[:end]
|
||||
t.buf = t.buf[end+separator:]
|
||||
if terminalSSEEvent(event) {
|
||||
t.terminal = true
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (t *sseEventTracker) Complete() bool { return len(t.buf) == 0 }
|
||||
|
||||
func completeSSEEvent(buf []byte) (int, int) {
|
||||
lf := bytes.Index(buf, []byte("\n\n"))
|
||||
crlf := bytes.Index(buf, []byte("\r\n\r\n"))
|
||||
if crlf >= 0 && (lf < 0 || crlf < lf) {
|
||||
return crlf, 4
|
||||
}
|
||||
if lf >= 0 {
|
||||
return lf, 2
|
||||
}
|
||||
return -1, 0
|
||||
}
|
||||
|
||||
func terminalSSEEvent(event []byte) bool {
|
||||
var data strings.Builder
|
||||
for _, line := range strings.Split(strings.ReplaceAll(string(event), "\r\n", "\n"), "\n") {
|
||||
if !strings.HasPrefix(line, "data:") {
|
||||
continue
|
||||
}
|
||||
if data.Len() > 0 {
|
||||
data.WriteByte('\n')
|
||||
}
|
||||
data.WriteString(strings.TrimSpace(strings.TrimPrefix(line, "data:")))
|
||||
}
|
||||
payload := data.String()
|
||||
if payload == "[DONE]" {
|
||||
return true
|
||||
}
|
||||
var envelope struct {
|
||||
Type string `json:"type"`
|
||||
}
|
||||
if json.Unmarshal([]byte(payload), &envelope) != nil {
|
||||
return false
|
||||
}
|
||||
return envelope.Type == "message_stop" || envelope.Type == "response.completed"
|
||||
}
|
||||
@@ -0,0 +1,101 @@
|
||||
package proxy
|
||||
|
||||
import (
|
||||
"io"
|
||||
"net/http"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
)
|
||||
|
||||
type trackingBody struct {
|
||||
io.ReadCloser
|
||||
failed *atomic.Bool
|
||||
}
|
||||
|
||||
func (b *trackingBody) Read(p []byte) (int, error) {
|
||||
n, err := b.ReadCloser.Read(p)
|
||||
if err != nil && err != io.EOF {
|
||||
b.failed.Store(true)
|
||||
}
|
||||
return n, err
|
||||
}
|
||||
|
||||
func responseBodyTimeout(resp *http.Response, opts Options) time.Duration {
|
||||
if resp.StatusCode == http.StatusSwitchingProtocols {
|
||||
return 0
|
||||
}
|
||||
if isStreamResponse(resp.Header) {
|
||||
return opts.SSEIdleTimeout
|
||||
}
|
||||
return opts.ResponseTimeout
|
||||
}
|
||||
|
||||
type timeoutBody struct {
|
||||
body io.ReadCloser
|
||||
idle bool
|
||||
timeout time.Duration
|
||||
timer *time.Timer
|
||||
mu sync.Mutex
|
||||
done bool
|
||||
sequence uint64
|
||||
closeOnce sync.Once
|
||||
closeErr error
|
||||
}
|
||||
|
||||
func newTimeoutBody(body io.ReadCloser, timeout time.Duration, idle bool) *timeoutBody {
|
||||
t := &timeoutBody{body: body, idle: idle, timeout: timeout}
|
||||
t.resetLocked()
|
||||
return t
|
||||
}
|
||||
|
||||
func (b *timeoutBody) Read(p []byte) (int, error) {
|
||||
n, err := b.body.Read(p)
|
||||
b.mu.Lock()
|
||||
if !b.done {
|
||||
if err != nil {
|
||||
b.done = true
|
||||
b.timer.Stop()
|
||||
} else if n > 0 && b.idle {
|
||||
b.resetLocked()
|
||||
}
|
||||
}
|
||||
b.mu.Unlock()
|
||||
return n, err
|
||||
}
|
||||
|
||||
func (b *timeoutBody) Close() error {
|
||||
b.mu.Lock()
|
||||
if !b.done {
|
||||
b.done = true
|
||||
b.sequence++
|
||||
b.timer.Stop()
|
||||
}
|
||||
b.mu.Unlock()
|
||||
return b.closeUnderlying()
|
||||
}
|
||||
|
||||
func (b *timeoutBody) resetLocked() {
|
||||
if b.timer != nil {
|
||||
b.timer.Stop()
|
||||
}
|
||||
b.sequence++
|
||||
sequence := b.sequence
|
||||
b.timer = time.AfterFunc(b.timeout, func() { b.expire(sequence) })
|
||||
}
|
||||
|
||||
func (b *timeoutBody) expire(sequence uint64) {
|
||||
b.mu.Lock()
|
||||
if b.done || sequence != b.sequence {
|
||||
b.mu.Unlock()
|
||||
return
|
||||
}
|
||||
b.done = true
|
||||
b.mu.Unlock()
|
||||
_ = b.closeUnderlying()
|
||||
}
|
||||
|
||||
func (b *timeoutBody) closeUnderlying() error {
|
||||
b.closeOnce.Do(func() { b.closeErr = b.body.Close() })
|
||||
return b.closeErr
|
||||
}
|
||||
+128
@@ -0,0 +1,128 @@
|
||||
package proxy
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"bytes"
|
||||
"errors"
|
||||
"net"
|
||||
"net/http"
|
||||
"sync"
|
||||
)
|
||||
|
||||
// captureWriter 包装 http.ResponseWriter，边转发边缓冲响应体。
|
||||
// 实现 http.Flusher 与 http.Hijacker 以支持 SSE/chunked/WebSocket。
|
||||
type captureWriter struct {
|
||||
http.ResponseWriter
|
||||
buf bytes.Buffer
|
||||
max int64
|
||||
written int64
|
||||
truncated bool
|
||||
status int
|
||||
wroteHeader bool
|
||||
hijacked bool
|
||||
writeFailed bool
|
||||
onHijack func(net.Conn) net.Conn
|
||||
sse sseEventTracker
|
||||
}
|
||||
|
||||
func newCaptureWriter(w http.ResponseWriter, max int64) *captureWriter {
|
||||
return &captureWriter{ResponseWriter: w, max: max, status: http.StatusOK}
|
||||
}
|
||||
|
||||
func (c *captureWriter) WriteHeader(code int) {
|
||||
if c.wroteHeader {
|
||||
return
|
||||
}
|
||||
c.status = code
|
||||
c.wroteHeader = true
|
||||
c.ResponseWriter.WriteHeader(code)
|
||||
}
|
||||
|
||||
func (c *captureWriter) Write(p []byte) (int, error) {
|
||||
if !c.wroteHeader {
|
||||
c.wroteHeader = true
|
||||
}
|
||||
n, err := c.ResponseWriter.Write(p)
|
||||
if err != nil || n != len(p) {
|
||||
c.writeFailed = true
|
||||
}
|
||||
if n > len(p) {
|
||||
n = len(p)
|
||||
}
|
||||
if n > 0 {
|
||||
// 仅缓冲 max 字节以内的内容。
|
||||
remaining := c.max - c.written
|
||||
if remaining > 0 {
|
||||
toBuf := n
|
||||
if int64(toBuf) > remaining {
|
||||
toBuf = int(remaining)
|
||||
c.truncated = true
|
||||
}
|
||||
c.buf.Write(p[:toBuf])
|
||||
} else if c.max > 0 {
|
||||
c.truncated = true
|
||||
}
|
||||
c.written += int64(n)
|
||||
if isStreamResponse(c.Header()) {
|
||||
c.sse.Write(p[:n])
|
||||
}
|
||||
if f, ok := c.ResponseWriter.(http.Flusher); ok {
|
||||
f.Flush()
|
||||
}
|
||||
}
|
||||
return n, err
|
||||
}
|
||||
|
||||
func (c *captureWriter) Flush() {
|
||||
if f, ok := c.ResponseWriter.(http.Flusher); ok {
|
||||
f.Flush()
|
||||
}
|
||||
}
|
||||
|
||||
func (c *captureWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
|
||||
h, ok := c.ResponseWriter.(http.Hijacker)
|
||||
if !ok {
|
||||
return nil, nil, errors.New("hijack not supported")
|
||||
}
|
||||
conn, rw, err := h.Hijack()
|
||||
if err != nil {
|
||||
c.writeFailed = true
|
||||
return nil, nil, err
|
||||
}
|
||||
c.hijacked = true
|
||||
if c.onHijack != nil {
|
||||
conn = c.onHijack(conn)
|
||||
}
|
||||
return conn, rw, nil
|
||||
}
|
||||
|
||||
func (c *captureWriter) Body() []byte { return c.buf.Bytes() }
|
||||
func (c *captureWriter) Truncated() bool { return c.truncated }
|
||||
func (c *captureWriter) Status() int { return c.status }
|
||||
func (c *captureWriter) Hijacked() bool { return c.hijacked }
|
||||
func (c *captureWriter) Complete() bool { return !c.writeFailed }
|
||||
func (c *captureWriter) SSEComplete() bool { return c.sse.Complete() }
|
||||
|
||||
func (c *captureWriter) OnHijack(fn func(net.Conn) net.Conn) { c.onHijack = fn }
|
||||
|
||||
func (c *captureWriter) SetHijackedResponse(status int, header http.Header) {
|
||||
c.status = status
|
||||
for key := range c.Header() {
|
||||
c.Header().Del(key)
|
||||
}
|
||||
for key, values := range header {
|
||||
c.Header()[key] = append([]string(nil), values...)
|
||||
}
|
||||
}
|
||||
|
||||
type trackedConn struct {
|
||||
net.Conn
|
||||
closeOnce sync.Once
|
||||
onClose func()
|
||||
}
|
||||
|
||||
func (c *trackedConn) Close() error {
|
||||
err := c.Conn.Close()
|
||||
c.closeOnce.Do(c.onClose)
|
||||
return err
|
||||
}
|
||||
Reference in New Issue
Block a user