From d28e90654c6e8fa2c706f8c8a43cd85ece77c918 Mon Sep 17 00:00:00 2001 From: imkp1 Date: Sat, 4 Jul 2026 21:33:46 +0530 Subject: [PATCH] Enforce attachment size limit on the byte stream The 50 MB attachment cap was only checked against the client-declared size (header.Size), while the actual write used an unbounded io.Copy and ParseMultipartForm's argument is an in-memory threshold, not a hard limit. A size-spoofed or chunked upload could therefore spill large temp files to disk during parsing and be written unbounded before the post-hoc size-mismatch check removed it. - handler: wrap the request body with http.MaxBytesReader so an oversized upload is refused up front (413) rather than spilled to disk while parsing the multipart form. The cap has headroom above the file size so a legitimate max-size file still fits. - service: copy through io.LimitReader(maxSize+1) and reject when the stream exceeds the cap, so enforcement no longer trusts the declared size. The cap is a struct field (defaulted to the const) to keep the guard testable without streaming 50 MB. Adds tests for both layers. --- .../internal/handler/attachment_handler.go | 28 +++++++++-- .../attachment_handler_internal_test.go | 48 ++++++++++++++++++ .../internal/service/attachment_service.go | 21 ++++++-- .../service/attachment_service_test.go | 49 +++++++++++++++++++ 4 files changed, 137 insertions(+), 9 deletions(-) create mode 100644 backend/internal/handler/attachment_handler_internal_test.go diff --git a/backend/internal/handler/attachment_handler.go b/backend/internal/handler/attachment_handler.go index a5b8e08..a057963 100644 --- a/backend/internal/handler/attachment_handler.go +++ b/backend/internal/handler/attachment_handler.go @@ -1,6 +1,7 @@ package handler import ( + "errors" "io" "net/http" "strings" @@ -10,26 +11,43 @@ import ( "go.uber.org/zap" ) +const ( + // maxUploadBytes is the attachment file-size cap. + maxUploadBytes = 50 * 1024 * 1024 // 50 MB + // maxRequestBytes caps the whole multipart request a little above the file + // cap — headroom for multipart framing and the small form fields — so a + // legitimate max-size file still fits while an oversized upload is refused + // before it can spill large temporary files to disk during parsing. + maxRequestBytes = maxUploadBytes + 1*1024*1024 +) + // AttachmentHandler serves the file attachment endpoints. // // POST /v1/attachments — upload a file // GET /v1/attachments/{id} — serve a file by ID type AttachmentHandler struct { - svc service.AttachmentService - logger *logger.Logger + svc service.AttachmentService + logger *logger.Logger + maxBody int64 } // NewAttachmentHandler creates a new AttachmentHandler. func NewAttachmentHandler(svc service.AttachmentService, log *logger.Logger) *AttachmentHandler { - return &AttachmentHandler{svc: svc, logger: log} + return &AttachmentHandler{svc: svc, logger: log, maxBody: maxRequestBytes} } -const maxUploadBytes = 50 * 1024 * 1024 // 50 MB - // Upload handles POST /v1/attachments. // Expects multipart/form-data with fields: file (required), conversationId, uploadBucket. func (h *AttachmentHandler) Upload(w http.ResponseWriter, r *http.Request) { + // Hard-cap the request body so an oversized upload is rejected up front + // rather than buffered/spilled to disk while parsing the multipart form. + r.Body = http.MaxBytesReader(w, r.Body, h.maxBody) if err := r.ParseMultipartForm(maxUploadBytes); err != nil { + var maxErr *http.MaxBytesError + if errors.As(err, &maxErr) { + http.Error(w, "upload exceeds max size", http.StatusRequestEntityTooLarge) + return + } http.Error(w, "request too large or not multipart", http.StatusBadRequest) return } diff --git a/backend/internal/handler/attachment_handler_internal_test.go b/backend/internal/handler/attachment_handler_internal_test.go new file mode 100644 index 0000000..f2f45e6 --- /dev/null +++ b/backend/internal/handler/attachment_handler_internal_test.go @@ -0,0 +1,48 @@ +package handler + +import ( + "bytes" + "mime/multipart" + "net/http" + "net/http/httptest" + "testing" + + "go.uber.org/mock/gomock" + + "github.com/dexiask/dexiask/internal/pkg/logger" + svcmocks "github.com/dexiask/dexiask/test/svcmocks" +) + +// TestUpload_RejectsOversizeBody verifies the request body is hard-capped: a +// body larger than maxBody is refused with 413 and the service is never called. +// maxBody is shrunk so the test sends a few hundred bytes instead of 50 MB. +func TestUpload_RejectsOversizeBody(t *testing.T) { + ctrl := gomock.NewController(t) + svc := svcmocks.NewMockAttachmentService(ctrl) // no calls expected → any call fails the test + h := NewAttachmentHandler(svc, logger.NewNop()) + h.maxBody = 64 + + var body bytes.Buffer + mw := multipart.NewWriter(&body) + fw, err := mw.CreateFormFile("file", "big.bin") + if err != nil { + t.Fatalf("create form file: %v", err) + } + if _, err := fw.Write(bytes.Repeat([]byte("a"), 256)); err != nil { + t.Fatalf("write body: %v", err) + } + _ = mw.WriteField("conversationId", "c1") + if err := mw.Close(); err != nil { + t.Fatalf("close writer: %v", err) + } + + req := httptest.NewRequest(http.MethodPost, "/v1/attachments", &body) + req.Header.Set("Content-Type", mw.FormDataContentType()) + rec := httptest.NewRecorder() + + h.Upload(rec, req) + + if rec.Code != http.StatusRequestEntityTooLarge { + t.Fatalf("expected 413, got %d", rec.Code) + } +} diff --git a/backend/internal/service/attachment_service.go b/backend/internal/service/attachment_service.go index 1b7bcde..a2ece3a 100644 --- a/backend/internal/service/attachment_service.go +++ b/backend/internal/service/attachment_service.go @@ -54,20 +54,26 @@ type attachmentService struct { workspaceRoot string attachmentRepo repository.AttachmentRepository logger *logger.Logger + maxSize int64 } // NewAttachmentService creates a new AttachmentService. workspaceRoot is the // host path mounted at /workspace (DEXIASK_WORKSPACE_MOUNT). func NewAttachmentService(workspaceRoot string, repo repository.AttachmentRepository, log *logger.Logger) AttachmentService { - return &attachmentService{workspaceRoot: workspaceRoot, attachmentRepo: repo, logger: log} + return &attachmentService{ + workspaceRoot: workspaceRoot, + attachmentRepo: repo, + logger: log, + maxSize: maxAttachmentSize, + } } func (s *attachmentService) Store(ctx context.Context, in StoreInput) (*model.Attachment, error) { if in.ConversationID == "" && in.UploadBucket == "" { return nil, pkgerrors.InvalidArgument("conversation_id or upload_bucket is required") } - if in.Size > maxAttachmentSize { - return nil, pkgerrors.InvalidArgumentf("file exceeds max size of %d bytes", maxAttachmentSize) + if in.Size > s.maxSize { + return nil, pkgerrors.InvalidArgumentf("file exceeds max size of %d bytes", s.maxSize) } safe := sanitizeFilename(in.Filename) @@ -111,11 +117,18 @@ func (s *attachmentService) Store(ctx context.Context, in StoreInput) (*model.At } defer f.Close() - written, err := io.Copy(f, in.Reader) + // Enforce the cap on the actual byte stream, not just the client-declared + // size: read at most maxSize+1 bytes so an oversized (or size-spoofed) + // upload is rejected rather than written to disk unbounded. + written, err := io.Copy(f, io.LimitReader(in.Reader, s.maxSize+1)) if err != nil { os.Remove(absPath) return nil, pkgerrors.Internal("failed to write attachment", err) } + if written > s.maxSize { + os.Remove(absPath) + return nil, pkgerrors.InvalidArgumentf("file exceeds max size of %d bytes", s.maxSize) + } if in.Size > 0 && written != in.Size { os.Remove(absPath) return nil, pkgerrors.InvalidArgumentf("attachment size mismatch: expected %d, got %d", in.Size, written) diff --git a/backend/internal/service/attachment_service_test.go b/backend/internal/service/attachment_service_test.go index b06d6fb..9ed44ab 100644 --- a/backend/internal/service/attachment_service_test.go +++ b/backend/internal/service/attachment_service_test.go @@ -2,6 +2,7 @@ package service import ( "context" + "io" "os" "path/filepath" "strings" @@ -14,6 +15,25 @@ import ( mocks "github.com/dexiask/dexiask/test/mocks" ) +// filler is an io.Reader that yields n bytes of 'a' without allocating them all, +// so a test can stream more than the size cap cheaply. +type filler struct{ remaining int64 } + +func (r *filler) Read(p []byte) (int, error) { + if r.remaining <= 0 { + return 0, io.EOF + } + n := int64(len(p)) + if n > r.remaining { + n = r.remaining + } + for i := int64(0); i < n; i++ { + p[i] = 'a' + } + r.remaining -= n + return int(n), nil +} + func newAttachmentSvc(t *testing.T) (*attachmentService, *mocks.MockAttachmentRepository, string) { t.Helper() ctrl := gomock.NewController(t) @@ -108,6 +128,35 @@ func TestStore_WritesUnderJail(t *testing.T) { } } +// TestStore_RejectsOversizeStream verifies the size cap is enforced on the byte +// stream itself — not just the client-declared size — and that a rejected +// upload leaves no file behind. The cap is shrunk so the test streams a handful +// of bytes instead of the real 50 MB. +func TestStore_RejectsOversizeStream(t *testing.T) { + svc, _, root := newAttachmentSvc(t) + svc.maxSize = 10 + + // Create must never be reached — the write is rejected before the DB insert + // (the repo mock has no expectations, so any call fails the test). + _, err := svc.Store(context.Background(), StoreInput{ + ConversationID: "c1", + Filename: "big.bin", + MediaType: "application/octet-stream", + Size: 0, // client under-declares / omits the size + Reader: &filler{remaining: 100}, // 100 bytes > 10-byte cap + }) + if err == nil { + t.Fatal("expected Store to reject an over-cap stream") + } + + // No partial file may remain under the workspace root. + attDir := filepath.Join(root, ".dexiask", "conversations", "c1", "attachments") + entries, _ := os.ReadDir(attDir) + if len(entries) != 0 { + t.Fatalf("expected no leftover files after rejection, found %d", len(entries)) + } +} + // TestReconcile_RequiresIDs guards the reconcile preconditions. func TestReconcile_RequiresIDs(t *testing.T) { svc, _, _ := newAttachmentSvc(t)