File size: 3,391 Bytes
8d3471e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 | package requestbody
import (
"bytes"
"errors"
"io"
"mime"
"net/http"
"strings"
"unicode/utf8"
)
var (
ErrInvalidUTF8Body = errors.New("invalid utf-8 request body")
errRequestBodyTooLarge = errors.New("request body too large")
)
const maxJSONUTF8ValidationSize = 100 << 20
// ValidateJSONUTF8 validates complete JSON request bodies before downstream
// decoders can silently replace malformed UTF-8 or stop before trailing bytes.
func ValidateJSONUTF8(next http.Handler) http.Handler {
if next == nil {
return http.HandlerFunc(func(http.ResponseWriter, *http.Request) {})
}
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if shouldValidateJSONBody(r) {
r.Body = validateAndReplayBody(r.Body)
}
next.ServeHTTP(w, r)
})
}
func shouldValidateJSONBody(r *http.Request) bool {
if r == nil || r.Body == nil {
return false
}
path := ""
if r.URL != nil {
path = r.URL.Path
}
return isJSONContentType(r.Header.Get("Content-Type")) || isKnownJSONRequestPath(r.Method, path)
}
func isJSONContentType(raw string) bool {
raw = strings.TrimSpace(raw)
if raw == "" {
return false
}
mediaType, _, err := mime.ParseMediaType(raw)
if err != nil {
mediaType = raw
}
mediaType = strings.ToLower(strings.TrimSpace(mediaType))
return strings.Contains(mediaType, "json")
}
func isKnownJSONRequestPath(method, path string) bool {
switch strings.ToUpper(strings.TrimSpace(method)) {
case http.MethodPost, http.MethodPut, http.MethodPatch, http.MethodDelete:
default:
return false
}
path = strings.TrimSpace(path)
if path == "" {
return false
}
switch {
case path == "/v1/chat/completions" || path == "/chat/completions":
return true
case path == "/v1/responses" || path == "/responses":
return true
case path == "/v1/embeddings" || path == "/embeddings":
return true
case path == "/anthropic/v1/messages" || path == "/v1/messages" || path == "/messages":
return true
case path == "/anthropic/v1/messages/count_tokens" || path == "/v1/messages/count_tokens" || path == "/messages/count_tokens":
return true
case strings.HasPrefix(path, "/v1beta/models/") || strings.HasPrefix(path, "/v1/models/"):
return strings.Contains(path, ":generateContent") || strings.Contains(path, ":streamGenerateContent")
case strings.HasPrefix(path, "/admin/"):
return true
default:
return false
}
}
func validateAndReplayBody(body io.ReadCloser) io.ReadCloser {
if body == nil {
return body
}
raw, err := io.ReadAll(io.LimitReader(body, maxJSONUTF8ValidationSize+1))
if err != nil {
return &errorReadCloser{err: err, closer: body}
}
if len(raw) > maxJSONUTF8ValidationSize {
return &errorReadCloser{err: errRequestBodyTooLarge, closer: body}
}
if !utf8.Valid(raw) {
return &errorReadCloser{err: ErrInvalidUTF8Body, closer: body}
}
return &replayReadCloser{Reader: bytes.NewReader(raw), closer: body}
}
type replayReadCloser struct {
*bytes.Reader
closer io.Closer
}
func (r *replayReadCloser) Close() error {
if r == nil || r.closer == nil {
return nil
}
return r.closer.Close()
}
type errorReadCloser struct {
err error
closer io.Closer
}
func (r *errorReadCloser) Read([]byte) (int, error) {
if r == nil || r.err == nil {
return 0, io.EOF
}
return 0, r.err
}
func (r *errorReadCloser) Close() error {
if r == nil || r.closer == nil {
return nil
}
return r.closer.Close()
}
|