You've already forked agentic-coding-workflow
feat(chat): SSE writer + multipart upload endpoint with mime whitelist
This commit is contained in:
@@ -0,0 +1,83 @@
|
||||
package chat
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"strconv"
|
||||
|
||||
"github.com/go-chi/chi/v5"
|
||||
|
||||
"github.com/yan1h/agent-coding-workflow/internal/infra/errs"
|
||||
)
|
||||
|
||||
// streamMessages implements GET /api/v1/conversations/{id}/stream.
|
||||
// It upgrades the connection to a Server-Sent Events stream and forwards
|
||||
// SSEEvents produced by MessageService.Stream until the channel closes or
|
||||
// the client disconnects.
|
||||
func (h *Handler) streamMessages(w http.ResponseWriter, r *http.Request) {
|
||||
uid, err := h.userID(r)
|
||||
if err != nil {
|
||||
writeErr(w, r, err)
|
||||
return
|
||||
}
|
||||
convID, err := parseUUID(chi.URLParam(r, "id"))
|
||||
if err != nil {
|
||||
writeErr(w, r, err)
|
||||
return
|
||||
}
|
||||
var since int64
|
||||
if s := r.URL.Query().Get("since"); s != "" {
|
||||
since, _ = strconv.ParseInt(s, 10, 64)
|
||||
}
|
||||
|
||||
flusher, ok := w.(http.Flusher)
|
||||
if !ok {
|
||||
http.Error(w, "streaming unsupported", http.StatusInternalServerError)
|
||||
return
|
||||
}
|
||||
|
||||
w.Header().Set("Content-Type", "text/event-stream")
|
||||
w.Header().Set("Cache-Control", "no-cache")
|
||||
w.Header().Set("Connection", "keep-alive")
|
||||
w.Header().Set("X-Accel-Buffering", "no")
|
||||
w.WriteHeader(http.StatusOK)
|
||||
|
||||
ch, err := h.msgSvc.Stream(r.Context(), uid, convID, since)
|
||||
if err != nil {
|
||||
writeSSEError(w, flusher, err)
|
||||
return
|
||||
}
|
||||
for {
|
||||
select {
|
||||
case <-r.Context().Done():
|
||||
return
|
||||
case ev, ok := <-ch:
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
writeSSE(w, flusher, ev)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// writeSSE writes a single SSE frame to w and flushes.
|
||||
func writeSSE(w http.ResponseWriter, f http.Flusher, ev SSEEvent) {
|
||||
data, _ := json.Marshal(ev.Data)
|
||||
fmt.Fprintf(w, "id: %d\nevent: %s\ndata: %s\n\n", ev.ID, ev.Event, data)
|
||||
f.Flush()
|
||||
}
|
||||
|
||||
// writeSSEError writes an error event to the SSE stream and flushes.
|
||||
func writeSSEError(w http.ResponseWriter, f http.Flusher, err error) {
|
||||
code := errs.CodeInternal
|
||||
msg := err.Error()
|
||||
if ae, ok := errs.As(err); ok {
|
||||
code = ae.Code
|
||||
msg = ae.Message
|
||||
}
|
||||
payload := map[string]string{"code": string(code), "message": msg}
|
||||
data, _ := json.Marshal(payload)
|
||||
fmt.Fprintf(w, "event: error\ndata: %s\n\n", data)
|
||||
f.Flush()
|
||||
}
|
||||
@@ -0,0 +1,206 @@
|
||||
package chat
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/go-chi/chi/v5"
|
||||
"github.com/google/uuid"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
"github.com/yan1h/agent-coding-workflow/internal/infra/errs"
|
||||
"github.com/yan1h/agent-coding-workflow/internal/transport/http/middleware"
|
||||
)
|
||||
|
||||
// ===== SSE-capable response writer =====
|
||||
|
||||
// flushRecorder is an httptest.ResponseRecorder that also implements http.Flusher.
|
||||
type flushRecorder struct {
|
||||
*httptest.ResponseRecorder
|
||||
flushed int
|
||||
}
|
||||
|
||||
func newFlushRecorder() *flushRecorder {
|
||||
return &flushRecorder{ResponseRecorder: httptest.NewRecorder()}
|
||||
}
|
||||
|
||||
func (f *flushRecorder) Flush() {
|
||||
f.flushed++
|
||||
}
|
||||
|
||||
// ===== Fake MessageService for SSE tests =====
|
||||
|
||||
// sseMessageService allows controlling the Stream channel.
|
||||
type sseMessageService struct {
|
||||
streamCh chan SSEEvent
|
||||
streamErr error
|
||||
}
|
||||
|
||||
func (s *sseMessageService) Send(_ context.Context, _, _ uuid.UUID, _ string, _ []uuid.UUID) (int64, int64, error) {
|
||||
return 0, 0, nil
|
||||
}
|
||||
func (s *sseMessageService) Stream(_ context.Context, _ uuid.UUID, _ uuid.UUID, _ int64) (<-chan SSEEvent, error) {
|
||||
if s.streamErr != nil {
|
||||
return nil, s.streamErr
|
||||
}
|
||||
return s.streamCh, nil
|
||||
}
|
||||
func (s *sseMessageService) Cancel(_ context.Context, _ uuid.UUID, _ int64) error { return nil }
|
||||
func (s *sseMessageService) Retry(_ context.Context, _ uuid.UUID, _ int64) (int64, error) {
|
||||
return 0, nil
|
||||
}
|
||||
func (s *sseMessageService) ListMessages(_ context.Context, _, _ uuid.UUID, _ *int64, _ int) ([]Message, error) {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
// ===== SSE helper unit tests =====
|
||||
|
||||
// TestSSE_WritesEventInProperFormat verifies that writeSSE produces the
|
||||
// canonical id/event/data/\n\n format.
|
||||
func TestSSE_WritesEventInProperFormat(t *testing.T) {
|
||||
w := newFlushRecorder()
|
||||
ev := SSEEvent{
|
||||
ID: 42,
|
||||
Event: "text_delta",
|
||||
Data: map[string]string{"delta": "hello"},
|
||||
}
|
||||
writeSSE(w, w, ev)
|
||||
|
||||
body := w.Body.String()
|
||||
assert.True(t, strings.HasPrefix(body, "id: 42\n"), "must start with id line")
|
||||
assert.Contains(t, body, "event: text_delta\n")
|
||||
assert.Contains(t, body, "data: ")
|
||||
assert.True(t, strings.HasSuffix(body, "\n\n"), "must end with double newline")
|
||||
assert.Equal(t, 1, w.flushed)
|
||||
|
||||
// Verify the data line is valid JSON.
|
||||
lines := strings.Split(strings.TrimRight(body, "\n"), "\n")
|
||||
var dataLine string
|
||||
for _, l := range lines {
|
||||
if strings.HasPrefix(l, "data: ") {
|
||||
dataLine = strings.TrimPrefix(l, "data: ")
|
||||
}
|
||||
}
|
||||
require.NotEmpty(t, dataLine)
|
||||
var payload map[string]string
|
||||
require.NoError(t, json.Unmarshal([]byte(dataLine), &payload))
|
||||
assert.Equal(t, "hello", payload["delta"])
|
||||
}
|
||||
|
||||
// TestSSE_WriteSSEError verifies that writeSSEError produces an error event
|
||||
// with the expected code/message fields from an AppError.
|
||||
func TestSSE_WriteSSEError(t *testing.T) {
|
||||
w := newFlushRecorder()
|
||||
appErr := errs.New(errs.CodeChatConversationNotFound, "conversation gone")
|
||||
writeSSEError(w, w, appErr)
|
||||
|
||||
body := w.Body.String()
|
||||
assert.Contains(t, body, "event: error\n")
|
||||
assert.True(t, strings.HasSuffix(body, "\n\n"))
|
||||
assert.Equal(t, 1, w.flushed)
|
||||
|
||||
lines := strings.Split(body, "\n")
|
||||
var dataLine string
|
||||
for _, l := range lines {
|
||||
if strings.HasPrefix(l, "data: ") {
|
||||
dataLine = strings.TrimPrefix(l, "data: ")
|
||||
}
|
||||
}
|
||||
require.NotEmpty(t, dataLine)
|
||||
var payload map[string]string
|
||||
require.NoError(t, json.Unmarshal([]byte(dataLine), &payload))
|
||||
assert.Equal(t, string(errs.CodeChatConversationNotFound), payload["code"])
|
||||
assert.Equal(t, "conversation gone", payload["message"])
|
||||
}
|
||||
|
||||
// TestStreamMessages_ContentTypeHeader verifies that the SSE handler sets
|
||||
// the required Content-Type and X-Accel-Buffering headers.
|
||||
func TestStreamMessages_ContentTypeHeader(t *testing.T) {
|
||||
uid := uuid.New()
|
||||
convID := uuid.New()
|
||||
|
||||
ch := make(chan SSEEvent)
|
||||
close(ch) // immediately closed — no events, handler returns quickly.
|
||||
|
||||
msgSvc := &sseMessageService{streamCh: ch}
|
||||
resolver := &fakeSessionResolver{uid: uid}
|
||||
upload := NewUploader(newFakeRepo(), nil, nil, 1<<20)
|
||||
h := NewHandler(&fakeConvSvc{}, msgSvc, &fakeTplSvc{}, &fakeEpSvc{}, &fakeUsageSvc{}, upload, resolver, newFakeAdminLookup())
|
||||
|
||||
r := chi.NewRouter()
|
||||
r.Use(middleware.RequestID)
|
||||
h.Mount(r)
|
||||
|
||||
req := httptest.NewRequest(http.MethodGet, fmt.Sprintf("/api/v1/conversations/%s/stream", convID), nil)
|
||||
req.Header.Set("Authorization", "Bearer test-token")
|
||||
w := newFlushRecorder()
|
||||
r.ServeHTTP(w, req)
|
||||
|
||||
assert.Equal(t, "text/event-stream", w.Header().Get("Content-Type"))
|
||||
assert.Equal(t, "no-cache", w.Header().Get("Cache-Control"))
|
||||
assert.Equal(t, "no", w.Header().Get("X-Accel-Buffering"))
|
||||
assert.Equal(t, http.StatusOK, w.Code)
|
||||
}
|
||||
|
||||
// TestStreamMessages_StopsWhenClientClosesContext verifies that when the
|
||||
// request context is cancelled mid-stream, the handler stops consuming events
|
||||
// and the channel is not blocked.
|
||||
func TestStreamMessages_StopsWhenClientClosesContext(t *testing.T) {
|
||||
uid := uuid.New()
|
||||
convID := uuid.New()
|
||||
|
||||
// Channel that never closes on its own — cancelling the ctx should stop the handler.
|
||||
ch := make(chan SSEEvent, 1)
|
||||
|
||||
msgSvc := &sseMessageService{streamCh: ch}
|
||||
resolver := &fakeSessionResolver{uid: uid}
|
||||
upload := NewUploader(newFakeRepo(), nil, nil, 1<<20)
|
||||
h := NewHandler(&fakeConvSvc{}, msgSvc, &fakeTplSvc{}, &fakeEpSvc{}, &fakeUsageSvc{}, upload, resolver, newFakeAdminLookup())
|
||||
|
||||
// Build a router identical to newTestHandler but without chi's context copying issues.
|
||||
r := chi.NewRouter()
|
||||
r.Use(middleware.RequestID)
|
||||
h.Mount(r)
|
||||
|
||||
// Create a cancellable request.
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
req := httptest.NewRequest(http.MethodGet, fmt.Sprintf("/api/v1/conversations/%s/stream", convID), nil)
|
||||
req.Header.Set("Authorization", "Bearer test-token")
|
||||
req = req.WithContext(ctx)
|
||||
|
||||
w := newFlushRecorder()
|
||||
|
||||
done := make(chan struct{})
|
||||
go func() {
|
||||
defer close(done)
|
||||
r.ServeHTTP(w, req)
|
||||
}()
|
||||
|
||||
// Send one event, then cancel.
|
||||
ch <- SSEEvent{ID: 1, Event: "text_delta", Data: "hi"}
|
||||
cancel()
|
||||
|
||||
// Handler must return promptly after context cancellation.
|
||||
select {
|
||||
case <-done:
|
||||
// success
|
||||
case <-waitTimeout(t, 2*time.Second):
|
||||
t.Fatal("handler did not stop after context cancellation")
|
||||
}
|
||||
|
||||
// The one event before cancel should have been written.
|
||||
assert.Contains(t, w.Body.String(), "text_delta")
|
||||
}
|
||||
|
||||
// waitTimeout returns a channel that closes after the given duration.
|
||||
func waitTimeout(t *testing.T, d time.Duration) <-chan time.Time {
|
||||
t.Helper()
|
||||
return time.After(d)
|
||||
}
|
||||
@@ -0,0 +1,125 @@
|
||||
package chat
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/google/uuid"
|
||||
|
||||
"github.com/yan1h/agent-coding-workflow/internal/infra/errs"
|
||||
"github.com/yan1h/agent-coding-workflow/internal/infra/storage"
|
||||
"github.com/yan1h/agent-coding-workflow/internal/transport/http/middleware"
|
||||
)
|
||||
|
||||
// Uploader handles multipart file uploads for chat attachments.
|
||||
// It validates the file size and MIME type, stores the file via Storage,
|
||||
// and records the attachment in the repository. If storage succeeds but DB
|
||||
// insert fails, the stored blob is deleted to prevent orphaned objects.
|
||||
type Uploader struct {
|
||||
repo Repository
|
||||
storage storage.Storage
|
||||
mimeWhite map[string]bool
|
||||
maxFile int64
|
||||
}
|
||||
|
||||
// NewUploader constructs an Uploader.
|
||||
// mimeList is a whitelist of allowed MIME types; "text/*" is a wildcard prefix.
|
||||
// maxFile is the maximum allowed file size in bytes.
|
||||
func NewUploader(repo Repository, st storage.Storage, mimeList []string, maxFile int64) *Uploader {
|
||||
mw := make(map[string]bool, len(mimeList))
|
||||
for _, m := range mimeList {
|
||||
mw[m] = true
|
||||
}
|
||||
return &Uploader{repo: repo, storage: st, mimeWhite: mw, maxFile: maxFile}
|
||||
}
|
||||
|
||||
// Handle implements POST /api/v1/uploads.
|
||||
func (u *Uploader) Handle(w http.ResponseWriter, r *http.Request) {
|
||||
uid, ok := middleware.UserIDFromContext(r.Context())
|
||||
if !ok {
|
||||
writeErr(w, r, errs.New(errs.CodeUnauthorized, "unauthorized"))
|
||||
return
|
||||
}
|
||||
|
||||
// Parse the multipart form; allow a 1 MiB overhead beyond maxFile for
|
||||
// form boundary / metadata.
|
||||
if err := r.ParseMultipartForm(u.maxFile + 1<<20); err != nil {
|
||||
writeErr(w, r, errs.Wrap(err, errs.CodeInvalidInput, "parse multipart"))
|
||||
return
|
||||
}
|
||||
|
||||
file, header, err := r.FormFile("file")
|
||||
if err != nil {
|
||||
writeErr(w, r, errs.Wrap(err, errs.CodeInvalidInput, "no file field"))
|
||||
return
|
||||
}
|
||||
defer file.Close()
|
||||
|
||||
if header.Size > u.maxFile {
|
||||
writeErr(w, r, errs.New(errs.CodeChatAttachmentTooLarge,
|
||||
fmt.Sprintf("size %d exceeds limit %d", header.Size, u.maxFile)))
|
||||
return
|
||||
}
|
||||
|
||||
mt := header.Header.Get("Content-Type")
|
||||
if !u.mimeAllowed(mt) {
|
||||
writeErr(w, r, errs.New(errs.CodeChatAttachmentUnsupportedMime, mt))
|
||||
return
|
||||
}
|
||||
|
||||
now := time.Now().UTC()
|
||||
id, err := uuid.NewV7()
|
||||
if err != nil {
|
||||
writeErr(w, r, errs.Wrap(err, errs.CodeInternal, "generate uuid"))
|
||||
return
|
||||
}
|
||||
ext := filepath.Ext(header.Filename)
|
||||
key := fmt.Sprintf("%04d/%02d/%s%s", now.Year(), int(now.Month()), id.String(), ext)
|
||||
|
||||
res, err := u.storage.Put(r.Context(), key, file, mt)
|
||||
if err != nil {
|
||||
writeErr(w, r, errs.Wrap(err, errs.CodeStoragePutFailed, "storage put"))
|
||||
return
|
||||
}
|
||||
|
||||
att, err := u.repo.InsertAttachment(r.Context(), InsertAttachmentParams{
|
||||
ID: id,
|
||||
UserID: uid,
|
||||
Filename: header.Filename,
|
||||
MimeType: mt,
|
||||
SizeBytes: res.Size,
|
||||
SHA256: res.SHA256,
|
||||
StoragePath: res.StoragePath,
|
||||
})
|
||||
if err != nil {
|
||||
// Best-effort blob cleanup to prevent orphaned storage objects.
|
||||
_ = u.storage.Delete(context.Background(), res.StoragePath)
|
||||
writeErr(w, r, err)
|
||||
return
|
||||
}
|
||||
|
||||
writeJSON(w, http.StatusOK, AttachmentDTO{
|
||||
ID: att.ID,
|
||||
Filename: att.Filename,
|
||||
MimeType: att.MimeType,
|
||||
SizeBytes: att.SizeBytes,
|
||||
SHA256: att.SHA256,
|
||||
CreatedAt: att.CreatedAt,
|
||||
})
|
||||
}
|
||||
|
||||
// mimeAllowed reports whether the given MIME type is in the whitelist.
|
||||
// "text/*" in the whitelist acts as a wildcard for any text/ subtype.
|
||||
func (u *Uploader) mimeAllowed(mt string) bool {
|
||||
if u.mimeWhite[mt] {
|
||||
return true
|
||||
}
|
||||
if u.mimeWhite["text/*"] && strings.HasPrefix(mt, "text/") {
|
||||
return true
|
||||
}
|
||||
return false
|
||||
}
|
||||
@@ -0,0 +1,231 @@
|
||||
package chat
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"crypto/sha256"
|
||||
"encoding/hex"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"io"
|
||||
"mime/multipart"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"net/textproto"
|
||||
"testing"
|
||||
|
||||
"github.com/go-chi/chi/v5"
|
||||
"github.com/google/uuid"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
"github.com/yan1h/agent-coding-workflow/internal/infra/errs"
|
||||
"github.com/yan1h/agent-coding-workflow/internal/infra/storage"
|
||||
"github.com/yan1h/agent-coding-workflow/internal/transport/http/middleware"
|
||||
)
|
||||
|
||||
// ===== upload-specific fakes =====
|
||||
|
||||
// uploadRepo wraps fakeRepo and overrides InsertAttachment to capture or fail.
|
||||
type uploadRepo struct {
|
||||
*fakeRepo
|
||||
inserted []InsertAttachmentParams
|
||||
insertErr error
|
||||
storedAtt *Attachment
|
||||
}
|
||||
|
||||
func (r *uploadRepo) InsertAttachment(_ context.Context, p InsertAttachmentParams) (*Attachment, error) {
|
||||
r.inserted = append(r.inserted, p)
|
||||
if r.insertErr != nil {
|
||||
return nil, r.insertErr
|
||||
}
|
||||
att := &Attachment{
|
||||
ID: p.ID, UserID: p.UserID, Filename: p.Filename, MimeType: p.MimeType,
|
||||
SizeBytes: p.SizeBytes, SHA256: p.SHA256, StoragePath: p.StoragePath,
|
||||
}
|
||||
r.storedAtt = att
|
||||
return att, nil
|
||||
}
|
||||
|
||||
// uploadStorage records put/delete and supports a forced putErr.
|
||||
type uploadStorage struct {
|
||||
blobs map[string][]byte
|
||||
putErr error
|
||||
deleteCalls []string
|
||||
}
|
||||
|
||||
func newUploadStorage() *uploadStorage {
|
||||
return &uploadStorage{blobs: make(map[string][]byte)}
|
||||
}
|
||||
|
||||
func (s *uploadStorage) Put(_ context.Context, key string, r io.Reader, _ string) (storage.PutResult, error) {
|
||||
if s.putErr != nil {
|
||||
return storage.PutResult{}, s.putErr
|
||||
}
|
||||
data, err := io.ReadAll(r)
|
||||
if err != nil {
|
||||
return storage.PutResult{}, err
|
||||
}
|
||||
s.blobs[key] = data
|
||||
h := sha256.Sum256(data)
|
||||
return storage.PutResult{StoragePath: key, SHA256: hex.EncodeToString(h[:]), Size: int64(len(data))}, nil
|
||||
}
|
||||
|
||||
func (s *uploadStorage) Get(_ context.Context, key string) (io.ReadCloser, error) {
|
||||
data, ok := s.blobs[key]
|
||||
if !ok {
|
||||
return nil, errors.New("not found: " + key)
|
||||
}
|
||||
return io.NopCloser(bytes.NewReader(data)), nil
|
||||
}
|
||||
|
||||
func (s *uploadStorage) Delete(_ context.Context, key string) error {
|
||||
s.deleteCalls = append(s.deleteCalls, key)
|
||||
delete(s.blobs, key)
|
||||
return nil
|
||||
}
|
||||
|
||||
// ===== test helpers =====
|
||||
|
||||
// uploadHarness builds a chi router with Auth middleware that resolves any
|
||||
// non-empty bearer token to uid, then mounts the Uploader under POST /uploads.
|
||||
type uploadHarness struct {
|
||||
uid uuid.UUID
|
||||
repo *uploadRepo
|
||||
stor *uploadStorage
|
||||
up *Uploader
|
||||
r chi.Router
|
||||
}
|
||||
|
||||
func newUploadHarness(t *testing.T, mimes []string, maxFile int64) *uploadHarness {
|
||||
t.Helper()
|
||||
uid := uuid.Must(uuid.NewV7())
|
||||
repo := &uploadRepo{fakeRepo: newFakeRepo()}
|
||||
stor := newUploadStorage()
|
||||
up := NewUploader(repo, stor, mimes, maxFile)
|
||||
resolver := &fakeSessionResolver{uid: uid}
|
||||
r := chi.NewRouter()
|
||||
r.Use(middleware.RequestID)
|
||||
r.Use(middleware.Auth(resolver, middleware.AuthOptions{}))
|
||||
r.Post("/uploads", up.Handle)
|
||||
return &uploadHarness{uid: uid, repo: repo, stor: stor, up: up, r: r}
|
||||
}
|
||||
|
||||
// buildMultipart constructs a multipart/form-data body with one "file" part.
|
||||
// declaredMime is what the part's Content-Type header advertises.
|
||||
func buildMultipart(t *testing.T, filename, declaredMime string, data []byte) (body *bytes.Buffer, contentType string) {
|
||||
t.Helper()
|
||||
body = &bytes.Buffer{}
|
||||
mw := multipart.NewWriter(body)
|
||||
hdr := textproto.MIMEHeader{}
|
||||
hdr.Set("Content-Disposition", `form-data; name="file"; filename="`+filename+`"`)
|
||||
hdr.Set("Content-Type", declaredMime)
|
||||
part, err := mw.CreatePart(hdr)
|
||||
require.NoError(t, err)
|
||||
_, err = part.Write(data)
|
||||
require.NoError(t, err)
|
||||
require.NoError(t, mw.Close())
|
||||
return body, mw.FormDataContentType()
|
||||
}
|
||||
|
||||
func newUploadRequest(t *testing.T, body *bytes.Buffer, contentType string, withAuth bool) *http.Request {
|
||||
t.Helper()
|
||||
req := httptest.NewRequest(http.MethodPost, "/uploads", body)
|
||||
req.Header.Set("Content-Type", contentType)
|
||||
if withAuth {
|
||||
req.Header.Set("Authorization", "Bearer test-token")
|
||||
}
|
||||
return req
|
||||
}
|
||||
|
||||
// ===== tests =====
|
||||
|
||||
func TestUploader_HappyPath_ReturnsAttachmentMetadata(t *testing.T) {
|
||||
h := newUploadHarness(t, []string{"text/plain"}, 1<<20)
|
||||
data := []byte("hello world")
|
||||
body, ct := buildMultipart(t, "greeting.txt", "text/plain", data)
|
||||
w := httptest.NewRecorder()
|
||||
h.r.ServeHTTP(w, newUploadRequest(t, body, ct, true))
|
||||
|
||||
require.Equal(t, http.StatusOK, w.Code, "body: %s", w.Body.String())
|
||||
var resp AttachmentDTO
|
||||
require.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp))
|
||||
assert.Equal(t, "greeting.txt", resp.Filename)
|
||||
assert.Equal(t, "text/plain", resp.MimeType)
|
||||
assert.Equal(t, int64(len(data)), resp.SizeBytes)
|
||||
|
||||
expectedSHA := sha256.Sum256(data)
|
||||
assert.Equal(t, hex.EncodeToString(expectedSHA[:]), resp.SHA256)
|
||||
|
||||
require.Len(t, h.repo.inserted, 1)
|
||||
got := h.repo.inserted[0]
|
||||
assert.Equal(t, h.uid, got.UserID)
|
||||
assert.Equal(t, "greeting.txt", got.Filename)
|
||||
}
|
||||
|
||||
func TestUploader_FileTooLarge_413(t *testing.T) {
|
||||
h := newUploadHarness(t, []string{"text/plain"}, 5)
|
||||
body, ct := buildMultipart(t, "big.txt", "text/plain", []byte("more than five bytes"))
|
||||
w := httptest.NewRecorder()
|
||||
h.r.ServeHTTP(w, newUploadRequest(t, body, ct, true))
|
||||
|
||||
require.Equal(t, http.StatusRequestEntityTooLarge, w.Code, "body: %s", w.Body.String())
|
||||
assert.Empty(t, h.repo.inserted, "no DB row should be inserted")
|
||||
}
|
||||
|
||||
func TestUploader_UnsupportedMime_415(t *testing.T) {
|
||||
h := newUploadHarness(t, []string{"text/plain"}, 1<<20)
|
||||
body, ct := buildMultipart(t, "evil.exe", "application/octet-stream", []byte("MZ..."))
|
||||
w := httptest.NewRecorder()
|
||||
h.r.ServeHTTP(w, newUploadRequest(t, body, ct, true))
|
||||
|
||||
require.Equal(t, http.StatusUnsupportedMediaType, w.Code)
|
||||
assert.Empty(t, h.repo.inserted)
|
||||
assert.Empty(t, h.stor.blobs, "blob must not be stored when mime is rejected")
|
||||
}
|
||||
|
||||
func TestUploader_StoragePutFails_NoOrphanRow(t *testing.T) {
|
||||
h := newUploadHarness(t, []string{"text/plain"}, 1<<20)
|
||||
h.stor.putErr = errors.New("disk full")
|
||||
body, ct := buildMultipart(t, "x.txt", "text/plain", []byte("hi"))
|
||||
w := httptest.NewRecorder()
|
||||
h.r.ServeHTTP(w, newUploadRequest(t, body, ct, true))
|
||||
|
||||
require.Equal(t, http.StatusBadGateway, w.Code)
|
||||
assert.Empty(t, h.repo.inserted, "no DB insert when storage fails")
|
||||
assert.Empty(t, h.stor.deleteCalls, "no delete when nothing was stored")
|
||||
}
|
||||
|
||||
func TestUploader_DBInsertFails_DeletesBlob(t *testing.T) {
|
||||
h := newUploadHarness(t, []string{"text/plain"}, 1<<20)
|
||||
h.repo.insertErr = errs.New(errs.CodeInternal, "db down")
|
||||
body, ct := buildMultipart(t, "x.txt", "text/plain", []byte("hi"))
|
||||
w := httptest.NewRecorder()
|
||||
h.r.ServeHTTP(w, newUploadRequest(t, body, ct, true))
|
||||
|
||||
require.Equal(t, http.StatusInternalServerError, w.Code)
|
||||
require.Len(t, h.repo.inserted, 1)
|
||||
require.Len(t, h.stor.deleteCalls, 1, "blob must be cleaned up after DB failure")
|
||||
assert.Equal(t, h.repo.inserted[0].StoragePath, h.stor.deleteCalls[0])
|
||||
}
|
||||
|
||||
func TestUploader_NoAuth_401(t *testing.T) {
|
||||
h := newUploadHarness(t, []string{"text/plain"}, 1<<20)
|
||||
body, ct := buildMultipart(t, "x.txt", "text/plain", []byte("hi"))
|
||||
w := httptest.NewRecorder()
|
||||
h.r.ServeHTTP(w, newUploadRequest(t, body, ct, false))
|
||||
|
||||
require.Equal(t, http.StatusUnauthorized, w.Code)
|
||||
assert.Empty(t, h.repo.inserted)
|
||||
}
|
||||
|
||||
func TestUploader_TextWildcard_Allows(t *testing.T) {
|
||||
h := newUploadHarness(t, []string{"text/*"}, 1<<20)
|
||||
body, ct := buildMultipart(t, "page.html", "text/html", []byte("<p>ok</p>"))
|
||||
w := httptest.NewRecorder()
|
||||
h.r.ServeHTTP(w, newUploadRequest(t, body, ct, true))
|
||||
|
||||
require.Equal(t, http.StatusOK, w.Code, "body: %s", w.Body.String())
|
||||
require.Len(t, h.repo.inserted, 1)
|
||||
assert.Equal(t, "text/html", h.repo.inserted[0].MimeType)
|
||||
}
|
||||
Reference in New Issue
Block a user