refactor: stream base64 attachment decoding and enforce request body size limits

Author: maatini

go.mod

@@ -7,6 +7,7 @@ toolchain go1.25.9
 require (
 	github.com/go-chi/chi/v5 v5.2.5
 	github.com/go-playground/validator/v10 v10.30.2
+	github.com/golang-jwt/jwt/v5 v5.3.1
 	github.com/google/uuid v1.6.0
 	github.com/graph-gophers/graphql-go v1.9.0
 	github.com/nats-io/nats.go v1.51.0
@@ -24,7 +25,6 @@ require (
 	github.com/go-playground/locales v0.14.1 // indirect
 	github.com/go-playground/universal-translator v0.18.1 // indirect
 	github.com/go-viper/mapstructure/v2 v2.4.0 // indirect
-	github.com/golang-jwt/jwt/v5 v5.3.1 // indirect
 	github.com/hako/durafmt v0.0.0-20210608085754-5c1018a4e16b // indirect
 	github.com/inconshreveable/mousetrap v1.1.0 // indirect
 	github.com/klauspost/compress v1.18.5 // indirect

internal/gateway/attachstore.go

@@ -1,10 +1,11 @@
 package gateway
 
 import (
-	"bytes"
 	"context"
+	"encoding/base64"
 	"fmt"
 	"strconv"
+	"strings"
 
 	"github.com/nats-io/nats.go"
 
@@ -21,17 +22,18 @@ func NewAttachmentStore(store nats.ObjectStore) *AttachmentStore {
 	return &AttachmentStore{store: store}
 }
 
-func (a *AttachmentStore) Upload(ctx context.Context, traceID string, attachments []domain.AttachmentDO) ([]domain.AttachmentDO, error) {
+func (a *AttachmentStore) Upload(ctx context.Context, traceID string, attachments []domain.Attachment) ([]domain.AttachmentDO, error) {
 	result := make([]domain.AttachmentDO, len(attachments))
 	for i, att := range attachments {
 		key := traceID + "/" + strconv.Itoa(i)
-		_, err := a.store.Put(&nats.ObjectMeta{Name: key}, bytes.NewReader(att.Content))
+		r := base64.NewDecoder(base64.StdEncoding, strings.NewReader(att.Content))
+		_, err := a.store.Put(&nats.ObjectMeta{Name: key}, r)
 		if err != nil {
 			return nil, fmt.Errorf("object store put %s: %w", key, err)
 		}
 		result[i] = domain.AttachmentDO{
 			Name:        att.Name,
-			ContentType: att.ContentType,
+			ContentType: att.MimeType,
 			ObjectKey:   key,
 		}
 	}

internal/gateway/handler.go

@@ -36,7 +36,7 @@ type natsPublisher interface {
 }
 
 type attachmentUploader interface {
-	Upload(ctx context.Context, traceID string, attachments []domain.AttachmentDO) ([]domain.AttachmentDO, error)
+	Upload(ctx context.Context, traceID string, attachments []domain.Attachment) ([]domain.AttachmentDO, error)
 }
 
 const (
@@ -72,9 +72,16 @@ func (h *Handler) handleSend(w http.ResponseWriter, r *http.Request) {
 	traceID := uuid.New().String()
 	ctx := r.Context()
 
+	r.Body = http.MaxBytesReader(w, r.Body, h.cfg.MaxBodySize)
 	var req domain.MailRequest
 	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
-		writeError(w, http.StatusBadRequest, domain.ErrJSONParseError, "invalid JSON body", traceID)
+		var maxBytesErr *http.MaxBytesError
+		if errors.As(err, &maxBytesErr) {
+			writeError(w, http.StatusRequestEntityTooLarge, domain.ErrBodyTooLarge,
+				fmt.Sprintf("request body exceeds limit of %d bytes", h.cfg.MaxBodySize), traceID)
+		} else {
+			writeError(w, http.StatusBadRequest, domain.ErrJSONParseError, "invalid JSON body", traceID)
+		}
 		return
 	}
 
@@ -119,14 +126,10 @@ func (h *Handler) handleSend(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 
-	// Decode attachments and upload to Object Store
-	attachments, err := decodeAttachments(req.Attachments)
-	if err != nil {
-		h.writeValidationError(w, err, traceID)
-		return
-	}
-	if len(attachments) > 0 {
-		attachments, err = h.attStore.Upload(ctx, traceID, attachments)
+	// Stage 6: upload attachments to Object Store (streaming base64 decode)
+	var attachmentDOs []domain.AttachmentDO
+	if len(req.Attachments) > 0 {
+		attachmentDOs, err = h.attStore.Upload(ctx, traceID, req.Attachments)
 		if err != nil {
 			slog.ErrorContext(ctx, "attachment upload failed",
 				slog.String("traceId", traceID),
@@ -148,7 +151,7 @@ func (h *Handler) handleSend(w http.ResponseWriter, r *http.Request) {
 		Subject:         req.Subject,
 		BodyContent:     req.BodyContent,
 		HtmlBodyContent: req.HtmlBodyContent,
-		Attachments:     attachments,
+		Attachments:     attachmentDOs,
 		TraceContext:    req.TraceContext,
 		Test:            sender.Test,
 	}

internal/gateway/handler_test.go

@@ -7,6 +7,7 @@ import (
 	"errors"
 	"net/http"
 	"net/http/httptest"
+	"strings"
 	"testing"
 
 	"dispatch/internal/config"
@@ -42,13 +43,17 @@ func (s *stubPublisher) Publish(_ context.Context, _ *domain.MailRequestDO) erro
 
 type stubAttStore struct{}
 
-func (s *stubAttStore) Upload(_ context.Context, _ string, atts []domain.AttachmentDO) ([]domain.AttachmentDO, error) {
-	return atts, nil
+func (s *stubAttStore) Upload(_ context.Context, _ string, atts []domain.Attachment) ([]domain.AttachmentDO, error) {
+	result := make([]domain.AttachmentDO, len(atts))
+	for i, a := range atts {
+		result[i] = domain.AttachmentDO{Name: a.Name, ContentType: a.MimeType}
+	}
+	return result, nil
 }
 
 type failAttStore struct{}
 
-func (s *failAttStore) Upload(_ context.Context, _ string, _ []domain.AttachmentDO) ([]domain.AttachmentDO, error) {
+func (s *failAttStore) Upload(_ context.Context, _ string, _ []domain.Attachment) ([]domain.AttachmentDO, error) {
 	return nil, errors.New("object store unavailable")
 }
 
@@ -197,6 +202,20 @@ func TestHandleSend_AttachmentUploadError(t *testing.T) {
 	}
 }
 
+func TestHandleSend_BodyTooLarge(t *testing.T) {
+	// MaxBodySize=10: any body > 10 bytes triggers MaxBytesError.
+	// Body must start with valid JSON chars so the scanner doesn't fail first.
+	cfg := config.Config{MaxBodySize: 10, MimeWhitelist: []string{}, MaxTotalAttachmentMB: 20}
+	h := NewHandler(cfg, &stubSenders{sender: defaultSender()}, &stubQuota{}, &stubSpam{}, &stubPublisher{}, &stubAttStore{})
+	body := `{"appTag":"test","bodyContent":"` + strings.Repeat("x", 100) + `"}`
+	req := httptest.NewRequest(http.MethodPost, "/dispatch/api/v1/mail/send", strings.NewReader(body))
+	rr := httptest.NewRecorder()
+	h.Router().ServeHTTP(rr, req)
+	if rr.Code != http.StatusRequestEntityTooLarge {
+		t.Fatalf("expected 413, got %d: %s", rr.Code, rr.Body.String())
+	}
+}
+
 func TestHandleHealth(t *testing.T) {
 	h := buildHandler(&stubSenders{sender: defaultSender()}, &stubQuota{}, &stubSpam{}, &stubPublisher{})
 	req := httptest.NewRequest(http.MethodGet, "/health", nil)

internal/gateway/validation.go

@@ -3,6 +3,7 @@ package gateway
 import (
 	"encoding/base64"
 	"fmt"
+	"io"
 	"strings"
 
 	"github.com/go-playground/validator/v10"
@@ -42,16 +43,22 @@ func validateRequest(req *domain.MailRequest, maxBodySize int64, mimeWhitelist [
 					Message: fmt.Sprintf("MIME type not allowed: %s", a.MimeType),
 				}
 			}
-			// base64 length × 3/4 ≈ raw bytes
-			rawBytes := int64(len(a.Content)) * 3 / 4
-			totalBytes += rawBytes
+			r := base64.NewDecoder(base64.StdEncoding, strings.NewReader(a.Content))
+			n, ioErr := io.Copy(io.Discard, r)
+			if ioErr != nil {
+				return &domain.ValidationError{
+					Code:    domain.ErrInvalidAttachmentType,
+					Message: fmt.Sprintf("attachment %q: invalid base64", a.Name),
+				}
+			}
+			totalBytes += n
 		}
 
 		maxBytes := int64(maxAttachMB) * 1024 * 1024
 		if totalBytes > maxBytes {
 			return &domain.ValidationError{
 				Code:    domain.ErrAttachmentTooLarge,
-				Message: fmt.Sprintf("total attachment size ~%d bytes exceeds limit %d MB", totalBytes, maxAttachMB),
+				Message: fmt.Sprintf("total attachment size %d bytes exceeds limit %d MB", totalBytes, maxAttachMB),
 			}
 		}
 	}
@@ -87,23 +94,3 @@ func checkDomains(sender domain.Sender, req *domain.MailRequest) error {
 	}
 	return nil
 }
-
-// decodeAttachments converts base64 content strings to raw bytes.
-func decodeAttachments(atts []domain.Attachment) ([]domain.AttachmentDO, error) {
-	result := make([]domain.AttachmentDO, 0, len(atts))
-	for _, a := range atts {
-		raw, err := base64.StdEncoding.DecodeString(a.Content)
-		if err != nil {
-			return nil, &domain.ValidationError{
-				Code:    domain.ErrInvalidAttachmentType,
-				Message: fmt.Sprintf("attachment %q: invalid base64", a.Name),
-			}
-		}
-		result = append(result, domain.AttachmentDO{
-			Name:        a.Name,
-			ContentType: a.MimeType,
-			Content:     raw,
-		})
-	}
-	return result, nil
-}

internal/gateway/validation_test.go

@@ -119,44 +119,28 @@ func TestCheckDomains_CCBlocked(t *testing.T) {
 	}
 }
 
-func TestDecodeAttachments_Empty(t *testing.T) {
-	result, err := decodeAttachments(nil)
-	if err != nil || len(result) != 0 {
-		t.Fatalf("expected empty result, nil; got %v, %v", result, err)
-	}
-}
-
-func TestDecodeAttachments_Valid(t *testing.T) {
-	atts := []domain.Attachment{
-		{Name: "f.pdf", MimeType: validationTestMime, Content: "aGVsbG8="}, // "hello"
-	}
-	result, err := decodeAttachments(atts)
-	if err != nil {
-		t.Fatalf("unexpected error: %v", err)
-	}
-	if len(result) != 1 || string(result[0].Content) != "hello" {
-		t.Errorf("unexpected result: %+v", result)
-	}
-}
-
-func TestDecodeAttachments_InvalidBase64(t *testing.T) {
-	atts := []domain.Attachment{
-		{Name: "bad.pdf", MimeType: validationTestMime, Content: "not-valid-base64!!!"},
+func TestValidateRequest_InvalidBase64Attachment(t *testing.T) {
+	req := &domain.MailRequest{
+		AppTag:     "t",
+		Recipients: []string{validationTestEmail},
+		Attachments: []domain.Attachment{
+			{Name: "bad.pdf", MimeType: validationTestMime, Content: "not-valid-base64!!!"},
+		},
 	}
-	_, err := decodeAttachments(atts)
+	err := validateRequest(req, 10_000_000, []string{validationTestMime}, 20)
 	var ve *domain.ValidationError
 	if !errors.As(err, &ve) || ve.Code != domain.ErrInvalidAttachmentType {
 		t.Fatalf("expected ErrInvalidAttachmentType for bad base64, got %v", err)
 	}
 }
 
 func TestValidateRequest_AttachmentTotalSizeExceeded(t *testing.T) {
-	// 2 MB content string → ~1.5 MB raw, exceeds 1 MB limit
+	// 1.6 MB of valid base64 ("AAAA"×400000) decodes to 1.2 MB — exceeds 1 MB limit
 	req := &domain.MailRequest{
 		AppTag:     "t",
 		Recipients: []string{validationTestEmail},
 		Attachments: []domain.Attachment{
-			{Name: "big.pdf", MimeType: validationTestMime, Content: string(make([]byte, 2*1024*1024))},
+			{Name: "big.pdf", MimeType: validationTestMime, Content: strings.Repeat("AAAA", 400_000)},
 		},
 	}
 	err := validateRequest(req, 10_000_000, []string{validationTestMime}, 1)