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,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