Multipart upload

attachment/store_s3.go

@@ -4,7 +4,6 @@ import (
 	"context"
 	"fmt"
 	"io"
-	"os"
 	"sync"
 
 	"heckel.io/ntfy/v2/log"
@@ -51,33 +50,31 @@ func (c *s3Store) Write(id string, in io.Reader, limiters ...util.Limiter) (int6
 	}
 	log.Tag(tagS3Store).Field("message_id", id).Debug("Writing attachment to S3")
 
-	// Write through limiters into a temp file. This avoids buffering the full attachment in
-	// memory while still giving us the Content-Length that PutObject requires.
+	// Stream through limiters via an io.Pipe directly to S3. PutObject supports chunked
+	// uploads, so no temp file or Content-Length is needed.
 	limiters = append(limiters, util.NewFixedLimiter(c.Remaining()))
-	tmpFile, err := os.CreateTemp("", "ntfy-s3-upload-*")
-	if err != nil {
-		return 0, fmt.Errorf("s3 store: failed to create temp file: %w", err)
+	pr, pw := io.Pipe()
+	lw := util.NewLimitWriter(pw, limiters...)
+	var size int64
+	var copyErr error
+	done := make(chan struct{})
+	go func() {
+		defer close(done)
+		size, copyErr = io.Copy(lw, in)
+		if copyErr != nil {
+			pw.CloseWithError(copyErr)
+		} else {
+			pw.Close()
+		}
+	}()
+	putErr := c.client.PutObject(context.Background(), id, pr)
+	pr.Close()
+	<-done
+	if copyErr != nil {
+		return 0, copyErr
 	}
-	tmpPath := tmpFile.Name()
-	defer os.Remove(tmpPath)
-	limitWriter := util.NewLimitWriter(tmpFile, limiters...)
-	size, err := io.Copy(limitWriter, in)
-	if err != nil {
-		tmpFile.Close()
-		return 0, err
-	}
-	if err := tmpFile.Close(); err != nil {
-		return 0, err
-	}
-
-	// Re-open the temp file for reading and stream it to S3
-	f, err := os.Open(tmpPath)
-	if err != nil {
-		return 0, err
-	}
-	defer f.Close()
-	if err := c.client.PutObject(context.Background(), id, f, size); err != nil {
-		return 0, err
+	if putErr != nil {
+		return 0, putErr
 	}
 	c.mu.Lock()
 	c.totalSizeCurrent += size

attachment/store_s3_test.go

@@ -167,27 +167,41 @@ func newTestS3Store(t *testing.T, server *httptest.Server, bucket, prefix string
 // ListObjectsV2. Uses path-style addressing: /{bucket}/{key}. Objects are stored in memory.
 
 type mockS3Server struct {
-	objects map[string][]byte // full key (bucket/key) -> body
+	objects map[string][]byte         // full key (bucket/key) -> body
+	uploads map[string]map[int][]byte // uploadID -> partNumber -> data
+	nextID  int                       // counter for generating upload IDs
 	mu      sync.RWMutex
 }
 
 func newMockS3Server() *httptest.Server {
-	m := &mockS3Server{objects: make(map[string][]byte)}
+	m := &mockS3Server{
+		objects: make(map[string][]byte),
+		uploads: make(map[string]map[int][]byte),
+	}
 	return httptest.NewTLSServer(m)
 }
 
 func (m *mockS3Server) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	// Path is /{bucket}[/{key...}]
 	path := strings.TrimPrefix(r.URL.Path, "/")
+	q := r.URL.Query()
 
 	switch {
+	case r.Method == http.MethodPut && q.Has("partNumber"):
+		m.handleUploadPart(w, r, path)
 	case r.Method == http.MethodPut:
 		m.handlePut(w, r, path)
-	case r.Method == http.MethodGet && r.URL.Query().Get("list-type") == "2":
+	case r.Method == http.MethodPost && q.Has("uploads"):
+		m.handleInitiateMultipart(w, r, path)
+	case r.Method == http.MethodPost && q.Has("uploadId"):
+		m.handleCompleteMultipart(w, r, path)
+	case r.Method == http.MethodDelete && q.Has("uploadId"):
+		m.handleAbortMultipart(w, r, path)
+	case r.Method == http.MethodGet && q.Get("list-type") == "2":
 		m.handleList(w, r, path)
 	case r.Method == http.MethodGet:
 		m.handleGet(w, r, path)
-	case r.Method == http.MethodPost && r.URL.Query().Has("delete"):
+	case r.Method == http.MethodPost && q.Has("delete"):
 		m.handleDelete(w, r, path)
 	default:
 		http.Error(w, "not implemented", http.StatusNotImplemented)
@@ -206,6 +220,77 @@ func (m *mockS3Server) handlePut(w http.ResponseWriter, r *http.Request, path st
 	w.WriteHeader(http.StatusOK)
 }
 
+func (m *mockS3Server) handleInitiateMultipart(w http.ResponseWriter, r *http.Request, path string) {
+	m.mu.Lock()
+	m.nextID++
+	uploadID := fmt.Sprintf("upload-%d", m.nextID)
+	m.uploads[uploadID] = make(map[int][]byte)
+	m.mu.Unlock()
+
+	w.Header().Set("Content-Type", "application/xml")
+	w.WriteHeader(http.StatusOK)
+	fmt.Fprintf(w, `<?xml version="1.0" encoding="UTF-8"?><InitiateMultipartUploadResult><UploadId>%s</UploadId></InitiateMultipartUploadResult>`, uploadID)
+}
+
+func (m *mockS3Server) handleUploadPart(w http.ResponseWriter, r *http.Request, path string) {
+	uploadID := r.URL.Query().Get("uploadId")
+	var partNumber int
+	fmt.Sscanf(r.URL.Query().Get("partNumber"), "%d", &partNumber)
+
+	body, err := io.ReadAll(r.Body)
+	if err != nil {
+		http.Error(w, err.Error(), http.StatusInternalServerError)
+		return
+	}
+
+	m.mu.Lock()
+	parts, ok := m.uploads[uploadID]
+	if !ok {
+		m.mu.Unlock()
+		http.Error(w, "NoSuchUpload", http.StatusNotFound)
+		return
+	}
+	parts[partNumber] = body
+	m.mu.Unlock()
+
+	etag := fmt.Sprintf(`"etag-part-%d"`, partNumber)
+	w.Header().Set("ETag", etag)
+	w.WriteHeader(http.StatusOK)
+}
+
+func (m *mockS3Server) handleCompleteMultipart(w http.ResponseWriter, r *http.Request, path string) {
+	uploadID := r.URL.Query().Get("uploadId")
+
+	m.mu.Lock()
+	parts, ok := m.uploads[uploadID]
+	if !ok {
+		m.mu.Unlock()
+		http.Error(w, "NoSuchUpload", http.StatusNotFound)
+		return
+	}
+
+	// Assemble parts in order
+	var assembled []byte
+	for i := 1; i <= len(parts); i++ {
+		assembled = append(assembled, parts[i]...)
+	}
+	m.objects[path] = assembled
+	delete(m.uploads, uploadID)
+	m.mu.Unlock()
+
+	w.Header().Set("Content-Type", "application/xml")
+	w.WriteHeader(http.StatusOK)
+	fmt.Fprintf(w, `<?xml version="1.0" encoding="UTF-8"?><CompleteMultipartUploadResult><Key>%s</Key></CompleteMultipartUploadResult>`, path)
+}
+
+func (m *mockS3Server) handleAbortMultipart(w http.ResponseWriter, r *http.Request, path string) {
+	uploadID := r.URL.Query().Get("uploadId")
+	m.mu.Lock()
+	delete(m.uploads, uploadID)
+	m.mu.Unlock()
+	w.WriteHeader(http.StatusNoContent)
+}
+
 func (m *mockS3Server) handleGet(w http.ResponseWriter, r *http.Request, path string) {
 	m.mu.RLock()
 	body, ok := m.objects[path]