From 4887ba0fccdda3b34a2607c3a16be95ae64501b9 Mon Sep 17 00:00:00 2001 From: Jose Vega Date: Tue, 1 Sep 2026 22:06:05 -0600 Subject: [PATCH] Make block uploads replayable and retry storage 5xx responses UploadBlock handed an io.Reader straight to resty's multipart helper. A request consumes that reader, so when the shared client retried after a dropped connection the second attempt could not contain the original encrypted block bytes. An HTTP 5xx from the storage endpoint was also not a retry condition, although a transient 502 from /storage/blocks is what started the production stall this was found in. Serialize the multipart form once into a byte buffer, keep the generated Content-Type (with its boundary) and give resty the byte slice as the body, so every retry resends the identical complete body. Add a request scoped retry condition for 500-599. The client's bounded retry count is unchanged. Cost: one extra in-memory copy of each in-flight block (4 MiB each, bounded by the caller's upload semaphore). Add regression tests against a local HTTP test server: one drops the connection on the first attempt and asserts the replayed body is byte-identical; one returns 502 first and asserts the retry carries the same body. Both fail before this change. --- block.go | 28 ++++++++++- block_test.go | 130 ++++++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 157 insertions(+), 1 deletion(-) create mode 100644 block_test.go diff --git a/block.go b/block.go index 74e89a62..744272fd 100644 --- a/block.go +++ b/block.go @@ -1,8 +1,11 @@ package proton import ( + "bytes" "context" "io" + "mime/multipart" + "net/http" "github.com/go-resty/resty/v2" ) @@ -33,10 +36,33 @@ func (c *Client) RequestBlockUpload(ctx context.Context, req BlockUploadReq) ([] } func (c *Client) UploadBlock(ctx context.Context, bareURL, token string, block io.Reader) error { + var body bytes.Buffer + writer := multipart.NewWriter(&body) + part, err := writer.CreateFormFile("Block", "blob") + if err != nil { + return err + } + if _, err := io.Copy(part, block); err != nil { + return err + } + if err := writer.Close(); err != nil { + return err + } + contentType := writer.FormDataContentType() + payload := body.Bytes() + return c.do(ctx, func(r *resty.Request) (*resty.Response, error) { return r. SetHeader("pm-storage-token", token). - SetMultipartField("Block", "blob", "application/octet-stream", block). + SetHeader("Content-Type", contentType). + SetBody(payload). + AddRetryCondition(func(res *resty.Response, _ error) bool { + if res == nil { + return false + } + return res.StatusCode() >= http.StatusInternalServerError && + res.StatusCode() < 600 + }). Post(bareURL) }) } diff --git a/block_test.go b/block_test.go new file mode 100644 index 00000000..11941d96 --- /dev/null +++ b/block_test.go @@ -0,0 +1,130 @@ +package proton_test + +import ( + "bytes" + "context" + "io" + "net/http" + "net/http/httptest" + "sync" + "testing" + "time" + + "github.com/rclone/go-proton-api" + "github.com/stretchr/testify/require" +) + +func TestUploadBlockReplaysMultipartBodyAfterConnectionDrop(t *testing.T) { + payload := []byte("encrypted Proton block payload") + + var ( + mu sync.Mutex + attempts [][]byte + ) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + var received []byte + file, _, err := r.FormFile("Block") + if err == nil { + received, err = io.ReadAll(file) + require.NoError(t, err) + require.NoError(t, file.Close()) + } + + mu.Lock() + attempts = append(attempts, received) + attempt := len(attempts) + mu.Unlock() + + if attempt == 1 { + hijacker, ok := w.(http.Hijacker) + require.True(t, ok) + conn, _, err := hijacker.Hijack() + require.NoError(t, err) + require.NoError(t, conn.Close()) + return + } + + w.Header().Set("Date", time.Now().UTC().Format(http.TimeFormat)) + w.Header().Set("Content-Type", "application/json") + _, err = io.WriteString(w, `{"Code":1000}`) + require.NoError(t, err) + })) + defer server.Close() + + manager := proton.New( + proton.WithHostURL(server.URL), + proton.WithRetryCount(1), + ) + defer manager.Close() + client := manager.NewClient("", "", "") + defer client.Close() + + err := client.UploadBlock( + context.Background(), + server.URL+"/storage/blocks", + "test-token", + bytes.NewReader(payload), + ) + require.NoError(t, err) + + mu.Lock() + defer mu.Unlock() + require.Len(t, attempts, 2) + require.Equal(t, payload, attempts[0]) + require.Equal(t, payload, attempts[1]) +} + +func TestUploadBlockRetriesBadGatewayWithSameBody(t *testing.T) { + payload := []byte("encrypted block survives 502 retry") + + var ( + mu sync.Mutex + attempts [][]byte + ) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + file, _, err := r.FormFile("Block") + require.NoError(t, err) + received, err := io.ReadAll(file) + require.NoError(t, err) + require.NoError(t, file.Close()) + + mu.Lock() + attempts = append(attempts, received) + attempt := len(attempts) + mu.Unlock() + + w.Header().Set("Date", time.Now().UTC().Format(http.TimeFormat)) + w.Header().Set("Content-Type", "application/json") + if attempt == 1 { + w.WriteHeader(http.StatusBadGateway) + _, err = io.WriteString(w, `{"Code":0,"Error":"simulated bad gateway"}`) + require.NoError(t, err) + return + } + _, err = io.WriteString(w, `{"Code":1000}`) + require.NoError(t, err) + })) + defer server.Close() + + manager := proton.New( + proton.WithHostURL(server.URL), + proton.WithRetryCount(1), + ) + defer manager.Close() + client := manager.NewClient("", "", "") + defer client.Close() + + err := client.UploadBlock( + context.Background(), + server.URL+"/storage/blocks", + "test-token", + bytes.NewReader(payload), + ) + require.NoError(t, err) + + mu.Lock() + defer mu.Unlock() + require.Len(t, attempts, 2) + require.Equal(t, payload, attempts[0]) + require.Equal(t, payload, attempts[1]) +}