Paginate google.List if it's not a google registry (#985)

pkg/v1/google/list.go

@@ -21,6 +21,7 @@ import (
 	"net/http"
 	"net/url"
 	"strconv"
+	"strings"
 	"time"
 
 	"github.com/google/go-containerregistry/pkg/authn"
@@ -83,32 +84,99 @@ func newLister(repo name.Repository, options ...Option) (*lister, error) {
 }
 
 func (l *lister) list(repo name.Repository) (*Tags, error) {
-	uri := url.URL{
+	uri := &url.URL{
 		Scheme: repo.Registry.Scheme(),
 		Host:   repo.Registry.RegistryStr(),
 		Path:   fmt.Sprintf("/v2/%s/tags/list", repo.RepositoryStr()),
+		// ECR returns an error if n > 1000:
+		// https://github.com/google/go-containerregistry/issues/681
+		RawQuery: "n=1000",
 	}
 
-	req, err := http.NewRequestWithContext(l.ctx, http.MethodGet, uri.String(), nil)
-	if err != nil {
-		return nil, err
+	tags := Tags{}
+
+	// get responses until there is no next page
+	for {
+		select {
+		case <-l.ctx.Done():
+			return nil, l.ctx.Err()
+		default:
+		}
+
+		req, err := http.NewRequest("GET", uri.String(), nil)
+		if err != nil {
+			return nil, err
+		}
+		req = req.WithContext(l.ctx)
+
+		resp, err := l.client.Do(req)
+		if err != nil {
+			return nil, err
+		}
+
+		if err := transport.CheckError(resp, http.StatusOK); err != nil {
+			return nil, err
+		}
+
+		parsed := Tags{}
+		if err := json.NewDecoder(resp.Body).Decode(&parsed); err != nil {
+			return nil, err
+		}
+
+		if err := resp.Body.Close(); err != nil {
+			return nil, err
+		}
+
+		if len(parsed.Manifests) != 0 {
+			// We're dealing with GCR, just return directly.
+			return &parsed, nil
+		}
+
+		// This isn't GCR, just append the tags and keep paginating.
+		logs.Warn.Printf("saw non-google tag listing response, falling back to pagination")
+		tags.Tags = append(tags.Tags, parsed.Tags...)
+
+		uri, err = getNextPageURL(resp)
+		if err != nil {
+			return nil, err
+		}
+		// no next page
+		if uri == nil {
+			break
+		}
 	}
-	resp, err := l.client.Do(req)
-	if err != nil {
-		return nil, err
+
+	return &tags, nil
+}
+
+// getNextPageURL checks if there is a Link header in a http.Response which
+// contains a link to the next page. If yes it returns the url.URL of the next
+// page otherwise it returns nil.
+func getNextPageURL(resp *http.Response) (*url.URL, error) {
+	link := resp.Header.Get("Link")
+	if link == "" {
+		return nil, nil
 	}
-	defer resp.Body.Close()
 
-	if err := transport.CheckError(resp, http.StatusOK); err != nil {
-		return nil, err
+	if link[0] != '<' {
+		return nil, fmt.Errorf("failed to parse link header: missing '<' in: %s", link)
 	}
 
-	tags := Tags{}
-	if err := json.NewDecoder(resp.Body).Decode(&tags); err != nil {
-		return nil, err
+	end := strings.Index(link, ">")
+	if end == -1 {
+		return nil, fmt.Errorf("failed to parse link header: missing '>' in: %s", link)
 	}
+	link = link[1:end]
 
-	return &tags, nil
+	linkURL, err := url.Parse(link)
+	if err != nil {
+		return nil, err
+	}
+	if resp.Request == nil || resp.Request.URL == nil {
+		return nil, nil
+	}
+	linkURL = resp.Request.URL.ResolveReference(linkURL)
+	return linkURL, nil
 }
 
 type rawManifestInfo struct {

pkg/v1/google/list_test.go

@@ -257,3 +257,76 @@ func TestWalk(t *testing.T) {
 		})
 	}
 }
+
+// Copied shamelessly from remote.
+func TestCancelledList(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel()
+
+	repoName := "doesnotmatter"
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		switch r.URL.Path {
+		case "/v2/":
+			w.WriteHeader(http.StatusOK)
+		default:
+			t.Fatalf("Unexpected path: %v", r.URL.Path)
+		}
+	}))
+	defer server.Close()
+	u, err := url.Parse(server.URL)
+	if err != nil {
+		t.Fatalf("url.Parse(%v) = %v", server.URL, err)
+	}
+
+	repo, err := name.NewRepository(fmt.Sprintf("%s/%s", u.Host, repoName), name.WeakValidation)
+	if err != nil {
+		t.Fatalf("name.NewRepository(%v) = %v", repoName, err)
+	}
+
+	_, err = List(repo, WithContext(ctx))
+	if !strings.Contains(err.Error(), context.Canceled.Error()) {
+		t.Errorf("wanted %q to contain %q", err.Error(), context.Canceled.Error())
+	}
+}
+
+func makeResp(hdr string) *http.Response {
+	return &http.Response{
+		Header: http.Header{
+			"Link": []string{hdr},
+		},
+	}
+}
+
+func TestGetNextPageURL(t *testing.T) {
+	for _, hdr := range []string{
+		"",
+		"<",
+		"><",
+		"<>",
+		fmt.Sprintf("<%c>", 0x7f), // makes url.Parse fail
+	} {
+		u, err := getNextPageURL(makeResp(hdr))
+		if err == nil && u != nil {
+			t.Errorf("Expected err, got %+v", u)
+		}
+	}
+
+	good := &http.Response{
+		Header: http.Header{
+			"Link": []string{"<example.com>"},
+		},
+		Request: &http.Request{
+			URL: &url.URL{
+				Scheme: "https",
+			},
+		},
+	}
+	u, err := getNextPageURL(good)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if u.Scheme != "https" {
+		t.Errorf("expected scheme to match request, got %s", u.Scheme)
+	}
+}