From 25234af3017d16152153b262d40a9dc66f65a320 Mon Sep 17 00:00:00 2001 From: CJACK Date: Mon, 13 Apr 2026 03:55:14 +0800 Subject: [PATCH] feat: enforce request body size limits and restrict inline file count to prevent resource exhaustion --- internal/adapter/openai/embeddings_handler.go | 5 +++++ internal/adapter/openai/file_inline_upload.go | 7 +++++++ internal/adapter/openai/handler_chat.go | 6 ++++++ internal/adapter/openai/handler_files.go | 6 ++++++ internal/adapter/openai/handler_routes.go | 7 +++++++ internal/adapter/openai/responses_handler.go | 5 +++++ 6 files changed, 36 insertions(+) diff --git a/internal/adapter/openai/embeddings_handler.go b/internal/adapter/openai/embeddings_handler.go index ff61be0..48dfdd8 100644 --- a/internal/adapter/openai/embeddings_handler.go +++ b/internal/adapter/openai/embeddings_handler.go @@ -26,8 +26,13 @@ func (h *Handler) Embeddings(w http.ResponseWriter, r *http.Request) { } defer h.Auth.Release(a) + r.Body = http.MaxBytesReader(w, r.Body, openAIGeneralMaxSize) var req map[string]any if err := json.NewDecoder(r.Body).Decode(&req); err != nil { + if strings.Contains(strings.ToLower(err.Error()), "too large") { + writeOpenAIError(w, http.StatusRequestEntityTooLarge, "request body too large") + return + } writeOpenAIError(w, http.StatusBadRequest, "invalid json") return } diff --git a/internal/adapter/openai/file_inline_upload.go b/internal/adapter/openai/file_inline_upload.go index dd73dec..5955e81 100644 --- a/internal/adapter/openai/file_inline_upload.go +++ b/internal/adapter/openai/file_inline_upload.go @@ -15,6 +15,8 @@ import ( "ds2api/internal/deepseek" ) +const maxInlineFilesPerRequest = 50 + type inlineFileUploadError struct { status int message string @@ -39,6 +41,7 @@ type inlineUploadState struct { handler *Handler auth *auth.RequestAuth uploadedByID map[string]string + uploadCount int } type inlineDecodedFile struct { @@ -129,10 +132,14 @@ func (s *inlineUploadState) tryUploadBlock(block map[string]any) (map[string]any if !ok { return nil, false, nil } + if s.uploadCount >= maxInlineFilesPerRequest { + return nil, true, fmt.Errorf("exceeded maximum of %d inline files per request", maxInlineFilesPerRequest) + } fileID, err := s.uploadInlineFile(decoded) if err != nil { return nil, true, &inlineFileUploadError{status: http.StatusInternalServerError, message: "Failed to upload inline file.", err: err} } + s.uploadCount++ replacement := map[string]any{ "type": decoded.ReplacementType, "file_id": fileID, diff --git a/internal/adapter/openai/handler_chat.go b/internal/adapter/openai/handler_chat.go index 9b242bb..5599eec 100644 --- a/internal/adapter/openai/handler_chat.go +++ b/internal/adapter/openai/handler_chat.go @@ -5,6 +5,7 @@ import ( "encoding/json" "io" "net/http" + "strings" "time" "ds2api/internal/auth" @@ -43,8 +44,13 @@ func (h *Handler) ChatCompletions(w http.ResponseWriter, r *http.Request) { r = r.WithContext(auth.WithAuth(r.Context(), a)) + r.Body = http.MaxBytesReader(w, r.Body, openAIGeneralMaxSize) var req map[string]any if err := json.NewDecoder(r.Body).Decode(&req); err != nil { + if strings.Contains(strings.ToLower(err.Error()), "too large") { + writeOpenAIError(w, http.StatusRequestEntityTooLarge, "request body too large") + return + } writeOpenAIError(w, http.StatusBadRequest, "invalid json") return } diff --git a/internal/adapter/openai/handler_files.go b/internal/adapter/openai/handler_files.go index 8253135..f15ea3b 100644 --- a/internal/adapter/openai/handler_files.go +++ b/internal/adapter/openai/handler_files.go @@ -28,7 +28,13 @@ func (h *Handler) UploadFile(w http.ResponseWriter, r *http.Request) { writeOpenAIError(w, http.StatusBadRequest, "content-type must be multipart/form-data") return } + // Enforce a hard cap on the total request body size to prevent OOM + r.Body = http.MaxBytesReader(w, r.Body, openAIUploadMaxSize) if err := r.ParseMultipartForm(openAIUploadMaxMemory); err != nil { + if strings.Contains(strings.ToLower(err.Error()), "too large") { + writeOpenAIError(w, http.StatusRequestEntityTooLarge, "file size exceeds limit") + return + } writeOpenAIError(w, http.StatusBadRequest, "invalid multipart form") return } diff --git a/internal/adapter/openai/handler_routes.go b/internal/adapter/openai/handler_routes.go index 43f7763..5e48953 100644 --- a/internal/adapter/openai/handler_routes.go +++ b/internal/adapter/openai/handler_routes.go @@ -13,6 +13,13 @@ import ( "ds2api/internal/util" ) +const ( + // openAIUploadMaxSize limits total multipart request body size (100 MiB). + openAIUploadMaxSize = 100 << 20 + // openAIGeneralMaxSize limits total JSON request body size (100 MiB). + openAIGeneralMaxSize = 100 << 20 +) + // writeJSON is a package-internal alias kept to avoid mass-renaming across // every call-site in this package. var writeJSON = util.WriteJSON diff --git a/internal/adapter/openai/responses_handler.go b/internal/adapter/openai/responses_handler.go index f77f725..6494157 100644 --- a/internal/adapter/openai/responses_handler.go +++ b/internal/adapter/openai/responses_handler.go @@ -65,8 +65,13 @@ func (h *Handler) Responses(w http.ResponseWriter, r *http.Request) { return } + r.Body = http.MaxBytesReader(w, r.Body, openAIGeneralMaxSize) var req map[string]any if err := json.NewDecoder(r.Body).Decode(&req); err != nil { + if strings.Contains(strings.ToLower(err.Error()), "too large") { + writeOpenAIError(w, http.StatusRequestEntityTooLarge, "request body too large") + return + } writeOpenAIError(w, http.StatusBadRequest, "invalid json") return }