fix(storage): fix panic when Flush called early (#11934)

tritone · web-flow · commit 7d0b8a75ae55 · 2025-04-02T16:50:21.000Z
Writer.Flush was panicking due to blocking content sniffing logic if Writer.Flush was called before 512 bytes.

Re-ordering calls here allows Flush to signal EOF and unblock the content sniffer if a flush happens that early.

Add a test with multiple small Flushes to validate that this works as expected.
diff --git a/storage/client_test.go b/storage/client_test.go
@@ -1235,6 +1235,118 @@ func TestWriterFlushAtCloseEmulated(t *testing.T) {
 	})
 }
 
+// Tests small flush (under 512 bytes) to verify that logic avoiding
+// content type sniffing works as expected in this case.
+func TestWriterSmallFlushEmulated(t *testing.T) {
+	transportClientTest(skipHTTP("appends only supported via gRPC"), t, func(t *testing.T, ctx context.Context, project, bucket string, client storageClient) {
+		// Create test bucket.
+		_, err := client.CreateBucket(ctx, project, bucket, &BucketAttrs{
+			Name: bucket,
+		}, nil)
+		if err != nil {
+			t.Fatalf("client.CreateBucket: %v", err)
+		}
+		prefix := time.Now().Nanosecond()
+		testCases := []struct {
+			initialBytes    []byte
+			wantContentType string
+		}{
+			{
+				initialBytes:    []byte{0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10},
+				wantContentType: "application/octet-stream",
+			},
+			{
+				initialBytes:    []byte("helloworld"),
+				wantContentType: "text/plain; charset=utf-8",
+			},
+		}
+
+		for _, tc := range testCases {
+			t.Run(tc.wantContentType, func(t *testing.T) {
+				objName := fmt.Sprintf("%d-object-%d", prefix, time.Now().Nanosecond())
+
+				vc := &Client{tc: client}
+				w := vc.Bucket(bucket).Object(objName).NewWriter(ctx)
+				w.Append = true
+				w.ChunkSize = MiB
+				var gotOffsets []int64
+				w.ProgressFunc = func(offset int64) {
+					gotOffsets = append(gotOffsets, offset)
+				}
+				wantOffsets := []int64{10, 1010, 1010 + MiB, 1010 + 2*MiB, 3 * MiB}
+
+				// Make content with fixed first 10 bytes which will yield
+				// expected type when sniffed.
+				content := bytes.Clone(randomBytes3MiB)
+				copy(content, tc.initialBytes)
+
+				// Test Flush at a 10 byte offset.
+				n, err := w.Write(content[:10])
+				if err != nil {
+					t.Fatalf("writing data: got %v; want ok", err)
+				}
+				if n != 10 {
+					t.Errorf("writing data: got %v bytes written, want %v", n, 10)
+				}
+				off, err := w.Flush()
+				if err != nil {
+					t.Fatalf("flush: got %v; want ok", err)
+				}
+				if off != 10 {
+					t.Errorf("flushing data: got %v bytes written, want %v", off, 10)
+				}
+				// Write another 1000 bytes and flush again.
+				n, err = w.Write(content[10:1010])
+				if err != nil {
+					t.Fatalf("writing data: got %v; want ok", err)
+				}
+				if n != 1000 {
+					t.Errorf("writing data: got %v bytes written, want %v", n, 1000)
+				}
+				off, err = w.Flush()
+				if err != nil {
+					t.Fatalf("flush: got %v; want ok", err)
+				}
+				if off != 1010 {
+					t.Errorf("flushing data: got %v bytes written, want %v", off, 1010)
+				}
+				// Write the rest of the object
+				_, err = w.Write(content[1010:])
+				if err != nil {
+					t.Fatalf("writing data: got %v; want ok", err)
+				}
+				if err := w.Close(); err != nil {
+					t.Fatalf("closing writer: %v", err)
+				}
+				// Check offsets
+				if !slices.Equal(gotOffsets, wantOffsets) {
+					t.Errorf("progress offsets: got %v, want %v", gotOffsets, wantOffsets)
+				}
+
+				// Download object and check data
+				r, err := veneerClient.Bucket(bucket).Object(objName).NewReader(ctx)
+				defer r.Close()
+				if err != nil {
+					t.Fatalf("opening reading: %v", err)
+				}
+				wantLen := 3 * MiB
+				got, err := io.ReadAll(r)
+				if n := len(got); n != wantLen {
+					t.Fatalf("expected to read %d bytes, but got %d (%v)", wantLen, n, err)
+				}
+				if diff := cmp.Diff(got, content); diff != "" {
+					t.Errorf("checking written content: got(-), want(+):\n%s", diff)
+				}
+				// Check expected content type.
+				if got, want := r.Attrs.ContentType, tc.wantContentType; got != want {
+					t.Errorf("content type: got %v, want %v", got, want)
+				}
+			})
+		}
+
+	})
+}
+
 func TestListNotificationsEmulated(t *testing.T) {
 	transportClientTest(skipGRPC("notifications not implemented"), t, func(t *testing.T, ctx context.Context, project, bucket string, client storageClient) {
 		// Populate test object.
diff --git a/storage/grpc_client.go b/storage/grpc_client.go
@@ -1718,28 +1718,29 @@ func (c *grpcStorageClient) OpenWriter(params *openWriterParams, opts ...storage
 	}
 	s.retry.maxRetryDuration = retryDeadline
 
+	// Set Flush func for use by exported Writer.Flush.
+	var gw *gRPCWriter
+	setFlush(func() (int64, error) {
+		return gw.flush()
+	})
+	gw, err := newGRPCWriter(c, s, params, pr, pw, params.setPipeWriter)
+	if err != nil {
+		errorf(err)
+		pr.CloseWithError(err)
+		close(params.donec)
+		return nil, err
+	}
+
 	// This function reads the data sent to the pipe and sends sets of messages
 	// on the gRPC client-stream as the buffer is filled.
 	go func() {
 		err := func() error {
 			// Unless the user told us the content type, we have to determine it from
 			// the first read.
-			var r io.Reader = pr
 			if params.attrs.ContentType == "" && !params.forceEmptyContentType {
-				r, params.attrs.ContentType = gax.DetermineContentType(r)
-			}
-
-			var gw *gRPCWriter
-			gw, err := newGRPCWriter(c, s, params, r, pw, params.setPipeWriter)
-			if err != nil {
-				return err
+				gw.reader, gw.spec.Resource.ContentType = gax.DetermineContentType(gw.reader)
 			}
 
-			// Set Flush func for use by exported Writer.Flush.
-			setFlush(func() (int64, error) {
-				return gw.flush()
-			})
-
 			// Loop until there is an error or the Object has been finalized.
 			for {
 				// Note: This blocks until either the buffer is full or EOF is read.
diff --git a/storage/writer.go b/storage/writer.go
@@ -186,6 +186,9 @@ func (w *Writer) Write(p []byte) (n int, err error) {
 // Do not call Flush concurrently with Write or Close. A single Writer is not
 // safe for unsynchronized use across threads.
 //
+// Note that calling Flush very early (before 512 bytes) may interfere with
+// automatic content sniffing in the Writer.
+//
 // Flush is supported only on gRPC clients where [Writer.Append] is set
 // to true. This feature is in preview and is not yet available for general use.
 func (w *Writer) Flush() (int64, error) {
