From ef8ea64e14ba86cbc51ce0b028023b5ab13f0712 Mon Sep 17 00:00:00 2001
From: Lantao Liu <lantaol@google.com>
Date: Wed, 30 May 2018 16:56:59 -0700
Subject: [PATCH] Port docker resolver fix #2364.

Signed-off-by: Lantao Liu <lantaol@google.com>
---
 pkg/containerd/resolver/fetcher.go      | 30 +++++++-
 pkg/containerd/resolver/fetcher_test.go | 94 +++++++++++++++++++++++++
 2 files changed, 122 insertions(+), 2 deletions(-)
 create mode 100644 pkg/containerd/resolver/fetcher_test.go

diff --git a/pkg/containerd/resolver/fetcher.go b/pkg/containerd/resolver/fetcher.go
index a6aa48431..4e01c3c4b 100644
--- a/pkg/containerd/resolver/fetcher.go
+++ b/pkg/containerd/resolver/fetcher.go
@@ -24,6 +24,7 @@ import (
 	"context"
 	"fmt"
 	"io"
+	"io/ioutil"
 	"net/http"
 	"path"
 	"strings"
@@ -92,8 +93,9 @@ func (r dockerFetcher) open(ctx context.Context, u, mediatype string, offset int
 	req.Header.Set("Accept", strings.Join([]string{mediatype, `*`}, ", "))
 
 	if offset > 0 {
-		// TODO(stevvooe): Only set this header in response to the
-		// "Accept-Ranges: bytes" header.
+		// Note: "Accept-Ranges: bytes" cannot be trusted as some endpoints
+		// will return the header without supporting the range. The content
+		// range must always be checked.
 		req.Header.Set("Range", fmt.Sprintf("bytes=%d-", offset))
 	}
 
@@ -114,6 +116,30 @@ func (r dockerFetcher) open(ctx context.Context, u, mediatype string, offset int
 		}
 		return nil, errors.Errorf("unexpected status code %v: %v", u, resp.Status)
 	}
+	if offset > 0 {
+		cr := resp.Header.Get("content-range")
+		if cr != "" {
+			if !strings.HasPrefix(cr, fmt.Sprintf("bytes %d-", offset)) {
+				return nil, errors.Errorf("unhandled content range in response: %v", cr)
+
+			}
+		} else {
+			// TODO: Should any cases where use of content range
+			// without the proper header be considerd?
+			// 206 responses?
+
+			// Discard up to offset
+			// Could use buffer pool here but this case should be rare
+			n, err := io.Copy(ioutil.Discard, io.LimitReader(resp.Body, offset))
+			if err != nil {
+				return nil, errors.Wrap(err, "failed to discard to offset")
+			}
+			if n != offset {
+				return nil, errors.Errorf("unable to discard to offset")
+			}
+
+		}
+	}
 
 	return resp.Body, nil
 }
diff --git a/pkg/containerd/resolver/fetcher_test.go b/pkg/containerd/resolver/fetcher_test.go
new file mode 100644
index 000000000..8c71522a8
--- /dev/null
+++ b/pkg/containerd/resolver/fetcher_test.go
@@ -0,0 +1,94 @@
+/*
+Copyright 2018 The Containerd Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package resolver
+
+import (
+	"context"
+	"fmt"
+	"io/ioutil"
+	"math/rand"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+)
+
+func TestFetcherOpen(t *testing.T) {
+	content := make([]byte, 128)
+	rand.New(rand.NewSource(1)).Read(content)
+	start := 0
+
+	s := httptest.NewServer(http.HandlerFunc(func(rw http.ResponseWriter, r *http.Request) {
+		if start > 0 {
+			rw.Header().Set("content-range", fmt.Sprintf("bytes %d-127/128", start))
+		}
+		rw.Header().Set("content-length", fmt.Sprintf("%d", len(content[start:])))
+		rw.Write(content[start:])
+	}))
+	defer s.Close()
+
+	f := dockerFetcher{&dockerBase{
+		client: s.Client(),
+	}}
+	ctx := context.Background()
+
+	checkReader := func(o int64) {
+		t.Helper()
+		rc, err := f.open(ctx, s.URL, "", o)
+		if err != nil {
+			t.Fatalf("failed to open: %+v", err)
+		}
+		b, err := ioutil.ReadAll(rc)
+		if err != nil {
+			t.Fatal(err)
+		}
+		expected := content[o:]
+		if len(b) != len(expected) {
+			t.Errorf("unexpected length %d, expected %d", len(b), len(expected))
+			return
+		}
+		for i, c := range expected {
+			if b[i] != c {
+				t.Errorf("unexpected byte %x at %d, expected %x", b[i], i, c)
+				return
+			}
+		}
+
+	}
+
+	checkReader(0)
+
+	// Test server ignores content range
+	checkReader(25)
+
+	// Use content range on server
+	start = 20
+	checkReader(20)
+
+	// Check returning just last byte and no bytes
+	start = 127
+	checkReader(127)
+	start = 128
+	checkReader(128)
+
+	// Check that server returning a different content range
+	// then requested errors
+	start = 30
+	_, err := f.open(ctx, s.URL, "", 20)
+	if err == nil {
+		t.Fatal("expected error opening with invalid server response")
+	}
+}