Enforce attachment size limit on the byte stream

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
 	}

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)
+	}
+}

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)

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)