Sub2API v1.0 - AI API 网关（二开初始版本，基于上游 Wei-Shaw/sub2api）

backend/internal/pkg/httputil/body.go

@@ -0,0 +1,178 @@
+package httputil
+
+import (
+	"bytes"
+	"compress/gzip"
+	"compress/zlib"
+	"errors"
+	"fmt"
+	"io"
+	"net/http"
+	"strings"
+
+	"github.com/klauspost/compress/zstd"
+)
+
+const (
+	requestBodyReadInitCap    = 512
+	requestBodyReadMaxInitCap = 1 << 20
+	jsonUTF8BOMLen            = 3
+	// maxDecompressedBodySize limits the decompressed request body to 64 MB
+	// to prevent decompression bomb attacks.
+	maxDecompressedBodySize = 64 << 20
+)
+
+// ReadRequestBodyWithPrealloc reads request body with preallocated buffer based
+// on content length, transparently decoding any Content-Encoding the upstream
+// client used to compress the body (zstd, gzip, deflate).
+func ReadRequestBodyWithPrealloc(req *http.Request) ([]byte, error) {
+	if req == nil || req.Body == nil {
+		return nil, nil
+	}
+
+	capHint := requestBodyReadInitCap
+	if req.ContentLength > 0 {
+		switch {
+		case req.ContentLength < int64(requestBodyReadInitCap):
+			capHint = requestBodyReadInitCap
+		case req.ContentLength > int64(requestBodyReadMaxInitCap):
+			capHint = requestBodyReadMaxInitCap
+		default:
+			capHint = int(req.ContentLength)
+		}
+	}
+
+	buf := bytes.NewBuffer(make([]byte, 0, capHint))
+	if _, err := io.Copy(buf, req.Body); err != nil {
+		return nil, err
+	}
+	raw := buf.Bytes()
+
+	enc := strings.ToLower(strings.TrimSpace(req.Header.Get("Content-Encoding")))
+	if enc == "" || enc == "identity" {
+		return raw, nil
+	}
+
+	decoded, err := decompressRequestBody(enc, raw)
+	if err != nil {
+		return nil, fmt.Errorf("decode Content-Encoding %q: %w", enc, err)
+	}
+
+	req.Header.Del("Content-Encoding")
+	req.Header.Del("Content-Length")
+	req.ContentLength = int64(len(decoded))
+
+	return decoded, nil
+}
+
+// ReadLenientJSONRequestBodyWithPrealloc reads a request body and normalizes
+// JSON string control bytes before strict validation.
+func ReadLenientJSONRequestBodyWithPrealloc(req *http.Request, maxNormalizedBytes int64) ([]byte, error) {
+	body, err := ReadRequestBodyWithPrealloc(req)
+	if err != nil {
+		return nil, err
+	}
+	return NormalizeLenientJSONRequestBody(body, maxNormalizedBytes)
+}
+
+func decompressRequestBody(encoding string, raw []byte) ([]byte, error) {
+	switch encoding {
+	case "zstd":
+		dec, err := zstd.NewReader(bytes.NewReader(raw))
+		if err != nil {
+			return nil, err
+		}
+		defer dec.Close()
+		return io.ReadAll(io.LimitReader(dec, maxDecompressedBodySize))
+	case "gzip", "x-gzip":
+		gr, err := gzip.NewReader(bytes.NewReader(raw))
+		if err != nil {
+			return nil, err
+		}
+		defer func() { _ = gr.Close() }()
+		return io.ReadAll(io.LimitReader(gr, maxDecompressedBodySize))
+	case "deflate":
+		zr, err := zlib.NewReader(bytes.NewReader(raw))
+		if err != nil {
+			return nil, err
+		}
+		defer func() { _ = zr.Close() }()
+		return io.ReadAll(io.LimitReader(zr, maxDecompressedBodySize))
+	default:
+		return nil, errors.New("unsupported Content-Encoding")
+	}
+}
+
+// NormalizeLenientJSONRequestBody escapes raw control bytes that broken
+// OpenAI-compatible clients sometimes place inside JSON strings.
+func NormalizeLenientJSONRequestBody(body []byte, maxNormalizedBytes int64) ([]byte, error) {
+	if maxNormalizedBytes <= 0 {
+		maxNormalizedBytes = maxDecompressedBodySize
+	}
+
+	body = trimUTF8BOM(body)
+	if len(body) == 0 {
+		return body, nil
+	}
+	if int64(len(body)) > maxNormalizedBytes {
+		return nil, &http.MaxBytesError{Limit: maxNormalizedBytes}
+	}
+
+	var out []byte
+	inString := false
+	escaped := false
+	for i, b := range body {
+		if inString && isJSONControlByte(b) {
+			if out == nil {
+				capHint := len(body) + 6
+				if int64(capHint) > maxNormalizedBytes {
+					capHint = int(maxNormalizedBytes)
+				}
+				out = make([]byte, 0, capHint)
+				out = append(out, body[:i]...)
+			}
+			if int64(len(out)+6) > maxNormalizedBytes {
+				return nil, &http.MaxBytesError{Limit: maxNormalizedBytes}
+			}
+			out = appendJSONUnicodeEscape(out, b)
+			escaped = false
+			continue
+		}
+
+		switch {
+		case escaped:
+			escaped = false
+		case inString && b == '\\':
+			escaped = true
+		case b == '"':
+			inString = !inString
+		}
+
+		if out != nil {
+			if int64(len(out)+1) > maxNormalizedBytes {
+				return nil, &http.MaxBytesError{Limit: maxNormalizedBytes}
+			}
+			out = append(out, b)
+		}
+	}
+	if out != nil {
+		return out, nil
+	}
+	return body, nil
+}
+
+func trimUTF8BOM(body []byte) []byte {
+	if len(body) >= jsonUTF8BOMLen && body[0] == 0xef && body[1] == 0xbb && body[2] == 0xbf {
+		return body[jsonUTF8BOMLen:]
+	}
+	return body
+}
+
+func isJSONControlByte(b byte) bool {
+	return b < 0x20 || b == 0x7f
+}
+
+func appendJSONUnicodeEscape(dst []byte, b byte) []byte {
+	const hex = "0123456789abcdef"
+	return append(dst, '\\', 'u', '0', '0', hex[b>>4], hex[b&0x0f])
+}