Allow streaming to S3

s3/client.go

@@ -45,25 +45,36 @@ func New(config *Config) *Client {
 }
 
 // PutObject uploads body to the given key. The key is automatically prefixed with the client's
-// configured prefix. The body size does not need to be known in advance.
+// configured prefix.
 //
-// If the entire body fits in a single part (5 MB), it is uploaded with a simple PUT request
-// (https://docs.aws.amazon.com/AmazonS3/latest/API/API_PutObject.html). Otherwise, the body
-// is uploaded using S3 multipart upload, reading one part at a time into memory
-// (https://docs.aws.amazon.com/AmazonS3/latest/API/API_CreateMultipartUpload.html).
-func (c *Client) PutObject(ctx context.Context, key string, body io.Reader) error {
+// If untrustedLength is between 1 and 5 GB, the body is streamed directly to S3 via a
+// single PUT request without buffering. The read is limited to untrustedLength bytes;
+// any extra data in the body is ignored. If the body is shorter than claimed, the upload fails.
+//
+// Otherwise (untrustedLength <= 0 or > 5 GB), the first 5 MB are buffered to decide
+// between a simple PUT and multipart upload.
+//
+// See https://docs.aws.amazon.com/AmazonS3/latest/API/API_PutObject.html
+// and https://docs.aws.amazon.com/AmazonS3/latest/API/API_CreateMultipartUpload.html
+func (c *Client) PutObject(ctx context.Context, key string, body io.Reader, untrustedLength int64) error {
+	if untrustedLength > 0 && untrustedLength <= maxSinglePutSize {
+		// Stream directly: Content-Length is known (but untrusted). LimitReader ensures we send at most
+		// untrustedLength bytes, and any extra data in body is ignored.
+		return c.putObject(ctx, key, io.LimitReader(body, untrustedLength), untrustedLength)
+	}
+	// Buffered path: read first 5 MB to decide simple vs multipart
 	first := make([]byte, partSize)
 	n, err := io.ReadFull(body, first)
 	if errors.Is(err, io.ErrUnexpectedEOF) || err == io.EOF {
-		return c.putObjectSimple(ctx, key, bytes.NewReader(first[:n]), int64(n))
+		return c.putObject(ctx, key, bytes.NewReader(first[:n]), int64(n))
 	} else if err != nil {
 		return fmt.Errorf("error reading object %s from client: %w", key, err)
 	}
 	return c.putObjectMultipart(ctx, key, io.MultiReader(bytes.NewReader(first), body))
 }
 
-// putObjectSimple uploads a body with known size using a simple PUT with UNSIGNED-PAYLOAD.
-func (c *Client) putObjectSimple(ctx context.Context, key string, body io.Reader, size int64) error {
+// putObject uploads a body with known size using a simple PUT with UNSIGNED-PAYLOAD.
+func (c *Client) putObject(ctx context.Context, key string, body io.Reader, size int64) error {
 	log.Tag(tagS3Client).Debug("Uploading object %s (%d bytes)", key, size)
 	req, err := http.NewRequestWithContext(ctx, http.MethodPut, c.config.ObjectURL(key), body)
 	if err != nil {

s3/client_test.go

@@ -419,7 +419,7 @@ func TestClient_PutGetObject(t *testing.T) {
 	ctx := context.Background()
 
 	// Put
-	err := client.PutObject(ctx, "test-key", strings.NewReader("hello world"))
+	err := client.PutObject(ctx, "test-key", strings.NewReader("hello world"), 0)
 	require.Nil(t, err)
 
 	// Get
@@ -439,7 +439,7 @@ func TestClient_PutGetObject_WithPrefix(t *testing.T) {
 
 	ctx := context.Background()
 
-	err := client.PutObject(ctx, "test-key", strings.NewReader("hello"))
+	err := client.PutObject(ctx, "test-key", strings.NewReader("hello"), 0)
 	require.Nil(t, err)
 
 	reader, _, err := client.GetObject(ctx, "test-key")
@@ -471,7 +471,7 @@ func TestClient_DeleteObjects(t *testing.T) {
 
 	// Put several objects
 	for i := 0; i < 5; i++ {
-		err := client.PutObject(ctx, fmt.Sprintf("key-%d", i), bytes.NewReader([]byte("data")))
+		err := client.PutObject(ctx, fmt.Sprintf("key-%d", i), bytes.NewReader([]byte("data")), 0)
 		require.Nil(t, err)
 	}
 	require.Equal(t, 5, mock.objectCount())
@@ -502,13 +502,13 @@ func TestClient_ListObjects(t *testing.T) {
 	// Client with prefix "pfx": list should only return objects under pfx/
 	client := newTestClient(server, "my-bucket", "pfx")
 	for i := 0; i < 3; i++ {
-		err := client.PutObject(ctx, fmt.Sprintf("%d", i), bytes.NewReader([]byte("x")))
+		err := client.PutObject(ctx, fmt.Sprintf("%d", i), bytes.NewReader([]byte("x")), 0)
 		require.Nil(t, err)
 	}
 
 	// Also put an object outside the prefix using a no-prefix client
 	clientNoPrefix := newTestClient(server, "my-bucket", "")
-	err := clientNoPrefix.PutObject(ctx, "other", bytes.NewReader([]byte("y")))
+	err := clientNoPrefix.PutObject(ctx, "other", bytes.NewReader([]byte("y")), 0)
 	require.Nil(t, err)
 
 	// List with prefix client: should only see 3
@@ -532,7 +532,7 @@ func TestClient_ListObjects_Pagination(t *testing.T) {
 
 	// Put 5 objects
 	for i := 0; i < 5; i++ {
-		err := client.PutObject(ctx, fmt.Sprintf("key-%02d", i), bytes.NewReader([]byte("x")))
+		err := client.PutObject(ctx, fmt.Sprintf("key-%02d", i), bytes.NewReader([]byte("x")), 0)
 		require.Nil(t, err)
 	}
 
@@ -564,7 +564,7 @@ func TestClient_ListAllObjects(t *testing.T) {
 	ctx := context.Background()
 
 	for i := 0; i < 10; i++ {
-		err := client.PutObject(ctx, fmt.Sprintf("key-%02d", i), bytes.NewReader([]byte("x")))
+		err := client.PutObject(ctx, fmt.Sprintf("key-%02d", i), bytes.NewReader([]byte("x")), 0)
 		require.Nil(t, err)
 	}
 
@@ -585,7 +585,7 @@ func TestClient_PutObject_LargeBody(t *testing.T) {
 	for i := range data {
 		data[i] = byte(i % 256)
 	}
-	err := client.PutObject(ctx, "large", bytes.NewReader(data))
+	err := client.PutObject(ctx, "large", bytes.NewReader(data), 0)
 	require.Nil(t, err)
 
 	reader, size, err := client.GetObject(ctx, "large")
@@ -609,7 +609,7 @@ func TestClient_PutObject_ChunkedUpload(t *testing.T) {
 	for i := range data {
 		data[i] = byte(i % 256)
 	}
-	err := client.PutObject(ctx, "multipart", bytes.NewReader(data))
+	err := client.PutObject(ctx, "multipart", bytes.NewReader(data), 0)
 	require.Nil(t, err)
 
 	reader, size, err := client.GetObject(ctx, "multipart")
@@ -633,7 +633,7 @@ func TestClient_PutObject_ExactPartSize(t *testing.T) {
 	for i := range data {
 		data[i] = byte(i % 256)
 	}
-	err := client.PutObject(ctx, "exact", bytes.NewReader(data))
+	err := client.PutObject(ctx, "exact", bytes.NewReader(data), 0)
 	require.Nil(t, err)
 
 	reader, size, err := client.GetObject(ctx, "exact")
@@ -645,6 +645,63 @@ func TestClient_PutObject_ExactPartSize(t *testing.T) {
 	require.Equal(t, data, got)
 }
 
+func TestClient_PutObject_StreamingExactLength(t *testing.T) {
+	server, _ := newMockS3Server()
+	defer server.Close()
+	client := newTestClient(server, "my-bucket", "pfx")
+
+	ctx := context.Background()
+
+	// untrustedLength matches body exactly — streams directly via putObject
+	err := client.PutObject(ctx, "stream-exact", strings.NewReader("hello world"), 11)
+	require.Nil(t, err)
+
+	reader, size, err := client.GetObject(ctx, "stream-exact")
+	require.Nil(t, err)
+	require.Equal(t, int64(11), size)
+	got, err := io.ReadAll(reader)
+	reader.Close()
+	require.Nil(t, err)
+	require.Equal(t, "hello world", string(got))
+}
+
+func TestClient_PutObject_StreamingBodyLongerThanClaimed(t *testing.T) {
+	server, _ := newMockS3Server()
+	defer server.Close()
+	client := newTestClient(server, "my-bucket", "pfx")
+
+	ctx := context.Background()
+
+	// Body has 11 bytes, but we claim 5 — only first 5 bytes should be stored
+	err := client.PutObject(ctx, "stream-long", strings.NewReader("hello world"), 5)
+	require.Nil(t, err)
+
+	reader, size, err := client.GetObject(ctx, "stream-long")
+	require.Nil(t, err)
+	require.Equal(t, int64(5), size)
+	got, err := io.ReadAll(reader)
+	reader.Close()
+	require.Nil(t, err)
+	require.Equal(t, "hello", string(got))
+}
+
+func TestClient_PutObject_StreamingBodyShorterThanClaimed(t *testing.T) {
+	server, _ := newMockS3Server()
+	defer server.Close()
+	client := newTestClient(server, "my-bucket", "pfx")
+
+	ctx := context.Background()
+
+	// Body has 5 bytes, but we claim 100 — should fail
+	err := client.PutObject(ctx, "stream-short", strings.NewReader("hello"), 100)
+	require.Error(t, err)
+	require.Contains(t, err.Error(), "ContentLength")
+
+	// Object should not exist
+	_, _, err = client.GetObject(ctx, "stream-short")
+	require.Error(t, err)
+}
+
 func TestClient_PutObject_NestedKey(t *testing.T) {
 	server, _ := newMockS3Server()
 	defer server.Close()
@@ -652,7 +709,7 @@ func TestClient_PutObject_NestedKey(t *testing.T) {
 
 	ctx := context.Background()
 
-	err := client.PutObject(ctx, "deep/nested/prefix/file.txt", strings.NewReader("nested"))
+	err := client.PutObject(ctx, "deep/nested/prefix/file.txt", strings.NewReader("nested"), 0)
 	require.Nil(t, err)
 
 	reader, _, err := client.GetObject(ctx, "deep/nested/prefix/file.txt")
@@ -682,7 +739,7 @@ func TestClient_ListAllObjects_20k(t *testing.T) {
 		for i := 0; i < batchSize; i++ {
 			idx := batch*batchSize + i
 			key := fmt.Sprintf("%08d", idx)
-			err := client.PutObject(ctx, key, bytes.NewReader([]byte("x")))
+			err := client.PutObject(ctx, key, bytes.NewReader([]byte("x")), 0)
 			require.Nil(t, err)
 		}
 	}
@@ -780,7 +837,7 @@ func TestClient_RealBucket(t *testing.T) {
 		content := "hello from ntfy s3 test"
 
 		// Put
-		err := client.PutObject(ctx, key, strings.NewReader(content))
+		err := client.PutObject(ctx, key, strings.NewReader(content), 0)
 		require.Nil(t, err)
 
 		// Get
@@ -818,7 +875,7 @@ func TestClient_RealBucket(t *testing.T) {
 
 		// Put 10 objects
 		for i := 0; i < 10; i++ {
-			err := listClient.PutObject(ctx, fmt.Sprintf("%d", i), strings.NewReader("x"))
+			err := listClient.PutObject(ctx, fmt.Sprintf("%d", i), strings.NewReader("x"), 0)
 			require.Nil(t, err)
 		}
 
@@ -843,7 +900,7 @@ func TestClient_RealBucket(t *testing.T) {
 			data[i] = byte(i % 256)
 		}
 
-		err := client.PutObject(ctx, key, bytes.NewReader(data))
+		err := client.PutObject(ctx, key, bytes.NewReader(data), 0)
 		require.Nil(t, err)
 
 		reader, size, err := client.GetObject(ctx, key)

s3/util.go

@@ -28,6 +28,10 @@ const (
 	// part size of 5 MB for all parts except the last.
 	partSize = 5 * 1024 * 1024
 
+	// maxSinglePutSize is the maximum size for a single PUT upload (5 GB).
+	// Objects larger than this must use multipart upload.
+	maxSinglePutSize = 5 * 1024 * 1024 * 1024
+
 	// maxPages is the max number of pages to iterate through when listing objects
 	maxPages = 500
 )