🛡️ fix: prevent OOM on large/decompressed requests; skip heavy prompt meta when token count is disabled

Clamp request body size (including post-decompression) to avoid memory exhaustion caused by huge payloads/zip bombs, especially with large-context Claude requests. Add a configurable `MAX_REQUEST_BODY_MB` (default `32`) and document it.

- Enforce max request body size after gzip/br decompression via `http.MaxBytesReader`
- Add a secondary size guard in `common.GetRequestBody` and cache-safe handling
- Return **413 Request Entity Too Large** on oversized bodies in relay entry
- Avoid building large `TokenCountMeta.CombineText` when both token counting and sensitive check are disabled (use lightweight meta for pricing)
- Update READMEs (CN/EN/FR/JA) with `MAX_REQUEST_BODY_MB`
- Fix a handful of vet/formatting issues encountered during the change
- `go test ./...` passes
This commit is contained in:
t0ng7u
2025-12-16 17:00:19 +08:00
parent 11593bd3da
commit 8e3f9b1faa
21 changed files with 149 additions and 35 deletions

View File

@@ -18,18 +18,47 @@ import (
const KeyRequestBody = "key_request_body"
func GetRequestBody(c *gin.Context) ([]byte, error) {
requestBody, _ := c.Get(KeyRequestBody)
if requestBody != nil {
return requestBody.([]byte), nil
var ErrRequestBodyTooLarge = errors.New("request body too large")
func IsRequestBodyTooLargeError(err error) bool {
if err == nil {
return false
}
requestBody, err := io.ReadAll(c.Request.Body)
if errors.Is(err, ErrRequestBodyTooLarge) {
return true
}
var mbe *http.MaxBytesError
return errors.As(err, &mbe)
}
func GetRequestBody(c *gin.Context) ([]byte, error) {
cached, exists := c.Get(KeyRequestBody)
if exists && cached != nil {
if b, ok := cached.([]byte); ok {
return b, nil
}
}
maxMB := constant.MaxRequestBodyMB
if maxMB <= 0 {
maxMB = 64
}
maxBytes := int64(maxMB) << 20
limited := io.LimitReader(c.Request.Body, maxBytes+1)
body, err := io.ReadAll(limited)
if err != nil {
_ = c.Request.Body.Close()
if IsRequestBodyTooLargeError(err) {
return nil, ErrRequestBodyTooLarge
}
return nil, err
}
_ = c.Request.Body.Close()
c.Set(KeyRequestBody, requestBody)
return requestBody.([]byte), nil
if int64(len(body)) > maxBytes {
return nil, ErrRequestBodyTooLarge
}
c.Set(KeyRequestBody, body)
return body, nil
}
func UnmarshalBodyReusable(c *gin.Context, v any) error {

View File

@@ -117,6 +117,8 @@ func initConstantEnv() {
constant.DifyDebug = GetEnvOrDefaultBool("DIFY_DEBUG", true)
constant.MaxFileDownloadMB = GetEnvOrDefault("MAX_FILE_DOWNLOAD_MB", 20)
constant.StreamScannerMaxBufferMB = GetEnvOrDefault("STREAM_SCANNER_MAX_BUFFER_MB", 64)
// MaxRequestBodyMB 请求体最大大小(解压后),用于防止超大请求/zip bomb导致内存暴涨
constant.MaxRequestBodyMB = GetEnvOrDefault("MAX_REQUEST_BODY_MB", 32)
// ForceStreamOption 覆盖请求参数强制返回usage信息
constant.ForceStreamOption = GetEnvOrDefaultBool("FORCE_STREAM_OPTION", true)
constant.CountToken = GetEnvOrDefaultBool("CountToken", true)