Commit 0b85a8da authored by shaw's avatar shaw
Browse files

fix: add io.LimitReader bounds to prevent OOM in image handling

Limit image download and multipart upload reads to 20MB to prevent
unbounded memory allocation from abnormal upstream responses.
parent 327da8e2
......@@ -48,6 +48,8 @@ const (
openAIImageBackendUserAgent = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/131.0.0.0 Safari/537.36"
openAIImageRequirementsDiff = "0fffff"
openAIImageLifecycleTimeout = 2 * time.Minute
openAIImageMaxDownloadBytes = 20 << 20 // 20MB per image download
openAIImageMaxUploadPartSize = 20 << 20 // 20MB per multipart upload part
)
type OpenAIImagesCapability string
......@@ -218,7 +220,7 @@ func parseOpenAIImagesMultipartRequest(body []byte, contentType string, req *Ope
continue
}
data, err := io.ReadAll(part)
data, err := io.ReadAll(io.LimitReader(part, openAIImageMaxUploadPartSize))
_ = part.Close()
if err != nil {
return fmt.Errorf("read multipart field %s: %w", name, err)
......@@ -1952,7 +1954,7 @@ func downloadOpenAIImageBytes(ctx context.Context, client *req.Client, headers h
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
return nil, newOpenAIImageStatusError(resp, "download image bytes failed")
}
return io.ReadAll(resp.Body)
return io.ReadAll(io.LimitReader(resp.Body, openAIImageMaxDownloadBytes))
}
func handleOpenAIImageBackendError(resp *req.Response) error {
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment