From 2c3e8638afc6702dcba732a1aa07ccb33eb9304b Mon Sep 17 00:00:00 2001
From: cojenco <cathyo@google.com>
Date: Tue, 20 Sep 2022 17:15:34 -0700
Subject: [PATCH] fix(gensupport): allow initial request for resumable uploads
 to retry w/ non-nil getBody (#1690)

This allows retries when initiating a resumable upload session.

Add support in media.go to return a getBody function so the request body can be recreated during retry attempts.

Tests pass locally in google-cloud-go/storage

Updates googleapis/google-cloud-go#6652

Co-authored-by: Chris Cotter <cjcotter@google.com>
---
 internal/gensupport/media.go      | 29 ++++++++++++++++++++---------
 internal/gensupport/media_test.go |  5 ++---
 2 files changed, 22 insertions(+), 12 deletions(-)

diff --git a/internal/gensupport/media.go b/internal/gensupport/media.go
index d14a22470c1..66caf24cb13 100644
--- a/internal/gensupport/media.go
+++ b/internal/gensupport/media.go
@@ -289,13 +289,12 @@ func (mi *MediaInfo) UploadRequest(reqHeaders http.Header, body io.Reader) (newB
 		// be retried because the data is stored in the MediaBuffer.
 		media, _, _, _ = mi.buffer.Chunk()
 	}
+	toCleanup := []io.Closer{}
 	if media != nil {
 		fb := readerFunc(body)
 		fm := readerFunc(media)
 		combined, ctype := CombineBodyMedia(body, "application/json", media, mi.mType)
-		toCleanup := []io.Closer{
-			combined,
-		}
+		toCleanup = append(toCleanup, combined)
 		if fb != nil && fm != nil {
 			getBody = func() (io.ReadCloser, error) {
 				rb := ioutil.NopCloser(fb())
@@ -309,18 +308,30 @@ func (mi *MediaInfo) UploadRequest(reqHeaders http.Header, body io.Reader) (newB
 				return r, nil
 			}
 		}
-		cleanup = func() {
-			for _, closer := range toCleanup {
-				_ = closer.Close()
-			}
-
-		}
 		reqHeaders.Set("Content-Type", ctype)
 		body = combined
 	}
 	if mi.buffer != nil && mi.mType != "" && !mi.singleChunk {
+		// This happens when initiating a resumable upload session.
+		// The initial request contains a JSON body rather than media.
+		// It can be retried with a getBody function that re-creates the request body.
+		fb := readerFunc(body)
+		if fb != nil {
+			getBody = func() (io.ReadCloser, error) {
+				rb := ioutil.NopCloser(fb())
+				toCleanup = append(toCleanup, rb)
+				return rb, nil
+			}
+		}
 		reqHeaders.Set("X-Upload-Content-Type", mi.mType)
 	}
+	// Ensure that any bodies created in getBody are cleaned up.
+	cleanup = func() {
+		for _, closer := range toCleanup {
+			_ = closer.Close()
+		}
+
+	}
 	return body, getBody, cleanup
 }
 
diff --git a/internal/gensupport/media_test.go b/internal/gensupport/media_test.go
index c05cc0cd739..7127bddf5a5 100644
--- a/internal/gensupport/media_test.go
+++ b/internal/gensupport/media_test.go
@@ -309,12 +309,11 @@ func TestUploadRequestGetBody(t *testing.T) {
 			wantGetBody: true,
 		},
 		{
-			desc: "chunk size < data size: MediaBuffer, >1 chunk, no getBody",
-			// No getBody here, because the initial request contains no media data
+			desc: "chunk size < data size: MediaBuffer, >1 chunk, getBody",
 			// Note that ChunkSize = 1 is rounded up to googleapi.MinUploadChunkSize.
 			r:           &nullReader{2 * googleapi.MinUploadChunkSize},
 			chunkSize:   1,
-			wantGetBody: false,
+			wantGetBody: true,
 		},
 	} {
 		cryptorand.Reader = mathrand.New(mathrand.NewSource(int64(i)))