From 0e2d080a8a86f64e2260ec8a1047b3a1a4916872 Mon Sep 17 00:00:00 2001
From: Damien Mathieu <dmathieu@salesforce.com>
Date: Thu, 9 May 2019 14:07:58 +0200
Subject: [PATCH] append the written bytes to the blob writer's size

Any byte written should append to the size. Otherwise, the full Size is
always zero.

Signed-off-by: Damien Mathieu <dmathieu@salesforce.com>
---
 registry/client/blob_writer.go      |  4 +-
 registry/client/blob_writer_test.go | 88 +++++++++++++++++++++++++++++
 2 files changed, 90 insertions(+), 2 deletions(-)

diff --git a/registry/client/blob_writer.go b/registry/client/blob_writer.go
index 695bf852..cc6e88ca 100644
--- a/registry/client/blob_writer.go
+++ b/registry/client/blob_writer.go
@@ -64,8 +64,8 @@ func (hbu *httpBlobUpload) ReadFrom(r io.Reader) (n int64, err error) {
 		return 0, fmt.Errorf("bad range format: %s", rng)
 	}
 
+	hbu.offset += end - start + 1
 	return (end - start + 1), nil
-
 }
 
 func (hbu *httpBlobUpload) Write(p []byte) (n int, err error) {
@@ -99,8 +99,8 @@ func (hbu *httpBlobUpload) Write(p []byte) (n int, err error) {
 		return 0, fmt.Errorf("bad range format: %s", rng)
 	}
 
+	hbu.offset += int64(end - start + 1)
 	return (end - start + 1), nil
-
 }
 
 func (hbu *httpBlobUpload) Size() int64 {
diff --git a/registry/client/blob_writer_test.go b/registry/client/blob_writer_test.go
index 099dca4f..21b3b3d2 100644
--- a/registry/client/blob_writer_test.go
+++ b/registry/client/blob_writer_test.go
@@ -209,3 +209,91 @@ func TestUploadReadFrom(t *testing.T) {
 		t.Fatalf("Unexpected response status: %s, expected %s", uploadErr.Status, expected)
 	}
 }
+
+func TestUploadSize(t *testing.T) {
+	_, b := newRandomBlob(64)
+	readFromLocationPath := "/v2/test/upload/readfrom/uploads/testid"
+	writeLocationPath := "/v2/test/upload/readfrom/uploads/testid"
+
+	m := testutil.RequestResponseMap([]testutil.RequestResponseMapping{
+		{
+			Request: testutil.Request{
+				Method: "GET",
+				Route:  "/v2/",
+			},
+			Response: testutil.Response{
+				StatusCode: http.StatusOK,
+				Headers: http.Header(map[string][]string{
+					"Docker-Distribution-API-Version": {"registry/2.0"},
+				}),
+			},
+		},
+		{
+			Request: testutil.Request{
+				Method: "PATCH",
+				Route:  readFromLocationPath,
+				Body:   b,
+			},
+			Response: testutil.Response{
+				StatusCode: http.StatusAccepted,
+				Headers: http.Header(map[string][]string{
+					"Docker-Upload-UUID": {"46603072-7a1b-4b41-98f9-fd8a7da89f9b"},
+					"Location":           {readFromLocationPath},
+					"Range":              {"0-63"},
+				}),
+			},
+		},
+		{
+			Request: testutil.Request{
+				Method: "PATCH",
+				Route:  writeLocationPath,
+				Body:   b,
+			},
+			Response: testutil.Response{
+				StatusCode: http.StatusAccepted,
+				Headers: http.Header(map[string][]string{
+					"Docker-Upload-UUID": {"46603072-7a1b-4b41-98f9-fd8a7da89f9b"},
+					"Location":           {writeLocationPath},
+					"Range":              {"0-63"},
+				}),
+			},
+		},
+	})
+
+	e, c := testServer(m)
+	defer c()
+
+	// Writing with ReadFrom
+	blobUpload := &httpBlobUpload{
+		client:   &http.Client{},
+		location: e + readFromLocationPath,
+	}
+
+	if blobUpload.Size() != 0 {
+		t.Fatalf("Wrong size returned from Size: %d, expected 0", blobUpload.Size())
+	}
+
+	_, err := blobUpload.ReadFrom(bytes.NewReader(b))
+	if err != nil {
+		t.Fatalf("Error calling ReadFrom: %s", err)
+	}
+
+	if blobUpload.Size() != 64 {
+		t.Fatalf("Wrong size returned from Size: %d, expected 64", blobUpload.Size())
+	}
+
+	// Writing with Write
+	blobUpload = &httpBlobUpload{
+		client:   &http.Client{},
+		location: e + writeLocationPath,
+	}
+
+	_, err = blobUpload.Write(b)
+	if err != nil {
+		t.Fatalf("Error calling Write: %s", err)
+	}
+
+	if blobUpload.Size() != 64 {
+		t.Fatalf("Wrong size returned from Size: %d, expected 64", blobUpload.Size())
+	}
+}