From 3cb2166dddc6788ba677e2cc1210067c4567d8fb Mon Sep 17 00:00:00 2001 From: Nabeel Alam Date: Mon, 15 Dec 2025 13:09:11 +0500 Subject: [PATCH] Pagination and Rate-Limit Handling In Docker Registry Namespace API Calls (#4557) * Added pagination and retry logic on rate-limit errors in docker registry list images calls * added rate limit error failure test * added exponential backoff * updated defered response body discards, added more comments, removed some unnecessary logic in Quay.ListImages() * readded quay list images query that the requests private images * added golang.org/x/time/rate to limit registry API calling rate * switched to using retryable http client in docker registry calls * modified rate limiter. created constant for max page size in api request --------- Co-authored-by: Amaan Ullah Co-authored-by: Kashif Khan <70996046+kashifkhan0771@users.noreply.github.com> --- pkg/sources/docker/registries.go | 392 ++++++++++++++++++-------- pkg/sources/docker/registries_test.go | 40 +++ 2 files changed, 309 insertions(+), 123 deletions(-) diff --git a/pkg/sources/docker/registries.go b/pkg/sources/docker/registries.go index 61ae89146..b4402c508 100644 --- a/pkg/sources/docker/registries.go +++ b/pkg/sources/docker/registries.go @@ -9,11 +9,21 @@ import ( "net/url" "path" "strings" - "time" + + "github.com/trufflesecurity/trufflehog/v3/pkg/common" + + "golang.org/x/time/rate" ) // defaultHTTPClient defines a shared HTTP client with timeout for all registry requests. -var defaultHTTPClient = &http.Client{Timeout: 10 * time.Second} +var defaultHTTPClient = common.RetryableHTTPClientTimeout(10) + +// registryRateLimiter limits how quickly we make registry API calls across all registries. +// We allow roughly 1 event every 1.5s, with a burst of 2 as a simple safeguard against overloading upstream APIs. +var registryRateLimiter = rate.NewLimiter(rate.Limit(2.0/3.0), 2) + +// maxRegistryPageSize defines the maximum number of images to request per page from a registry API. +const maxRegistryPageSize = 100 // Image represents a container image or repository entry in a registry API response. type Image struct { @@ -22,10 +32,10 @@ type Image struct { // Registry is an interface for any Docker/OCI registry implementation that can list all images under a given namespace. type Registry interface { - Name() string // return name of the registry - WithRegistryToken(registryToken string) // set token for registry - // TODO: Handle pagination and rate limits for list images API Call + Name() string // return name of the registry + WithRegistryToken(registryToken string) // set token for registry ListImages(ctx context.Context, namespace string) ([]string, error) // list all images + WithClient(client *http.Client) // return the HTTP client to use } // MakeRegistryFromNamespace returns a Registry implementation @@ -45,11 +55,12 @@ func MakeRegistryFromNamespace(namespace string) Registry { return registry } -// === DockerHub registry === +// === Docker Hub Registry === -// DockerHub implements the Registry interface for Docker Hub. +// DockerHub implements the Registry interface for hub.docker.com. type DockerHub struct { - Token string + Token string + Client *http.Client } // dockerhubResp models Docker Hub's /v2/namespaces//repositories API response. @@ -66,67 +77,101 @@ func (d *DockerHub) WithRegistryToken(registryToken string) { d.Token = registryToken } +func (d *DockerHub) WithClient(client *http.Client) { + d.Client = client +} + // ListImages lists all images under a Docker Hub namespace using Docker Hub's API. +// +// We fetch images in fixed-size pages and keep following the "next" link until there +// are no more pages. A shared rate limiter is used so we don't accidentally hammer +// the Docker Hub API. func (d *DockerHub) ListImages(ctx context.Context, namespace string) ([]string, error) { - url := &url.URL{ + baseURL := &url.URL{ Scheme: "https", Host: "hub.docker.com", Path: path.Join("v2", "namespaces", namespace, "repositories"), } - req, err := http.NewRequestWithContext(ctx, http.MethodGet, url.String(), http.NoBody) - if err != nil { - return nil, err - } + query := baseURL.Query() + query.Set("page_size", fmt.Sprint(maxRegistryPageSize)) + baseURL.RawQuery = query.Encode() - if d.Token != "" { - req.Header.Set("Authorization", "Bearer "+d.Token) - } + allImages := []string{} + nextURL := baseURL.String() - resp, err := defaultHTTPClient.Do(req) - if err != nil { - return nil, err - } - - defer func() { - _, _ = io.Copy(io.Discard, resp.Body) - _ = resp.Body.Close() - }() - - responseBodyByte, err := io.ReadAll(resp.Body) - if err != nil { - return nil, err - } - - switch resp.StatusCode { - case http.StatusOK: - var allImages = make([]string, 0) - - var listImagesResp dockerhubResp - if err := json.Unmarshal(responseBodyByte, &listImagesResp); err != nil { + for { + if err := registryRateLimiter.Wait(ctx); err != nil { return nil, err } - for _, image := range listImagesResp.Results { + req, err := http.NewRequestWithContext(ctx, http.MethodGet, nextURL, http.NoBody) + if err != nil { + return nil, err + } + + if d.Token != "" { + req.Header.Set("Authorization", "Bearer "+d.Token) + } + + client := d.Client + if client == nil { + client = defaultHTTPClient + } + resp, err := client.Do(req) + if err != nil { + return nil, err + } + + body, err := io.ReadAll(resp.Body) + discardBody(resp) + if err != nil { + return nil, err + } + + if resp.StatusCode != http.StatusOK { + return nil, fmt.Errorf("failed to list dockerhub images: unexpected status code: %d", resp.StatusCode) + } + + var page dockerhubResp + if err := json.Unmarshal(body, &page); err != nil { + return nil, err + } + + for _, image := range page.Results { allImages = append(allImages, fmt.Sprintf("%s/%s", namespace, image.Name)) // / } - return allImages, nil - default: - return nil, fmt.Errorf("failed to list dockerhub images: unexpected status code: %d", resp.StatusCode) + if page.Next == "" { + break + } + + // Docker Hub sometimes returns an absolute "next" URL and sometimes a + // relative one. ResolveReference cleans that up for us and turns whatever + // they send into a proper URL we can call. + next, err := url.Parse(page.Next) + if err != nil { + return nil, err + } + nextURL = baseURL.ResolveReference(next).String() } + + return allImages, nil } // === Red Hat Quay Registry === // Quay implements the Registry interface for Quay.io. type Quay struct { - Token string + Token string + Client *http.Client } // quayResp models the JSON structure returned by Quay's /api/v1/repository endpoint. type quayResp struct { - Repositories []Image `json:"repositories"` + Repositories []Image `json:"repositories"` + HasAdditional bool `json:"has_additional"` + NextPage string `json:"next_page"` } func (q *Quay) Name() string { @@ -137,69 +182,96 @@ func (q *Quay) WithRegistryToken(registryToken string) { q.Token = registryToken } +func (q *Quay) WithClient(client *http.Client) { + q.Client = client +} + // ListImages lists all images under a Quay namespace. +// API reference: +// +// GET https://quay.io/api/v1/repository?namespace=&public=true&private=true&next_page= +// +// We keep following next_page while has_additional is true. func (q *Quay) ListImages(ctx context.Context, namespace string) ([]string, error) { quayNamespace := path.Base(namespace) // quay.io/ -> namespace - url := &url.URL{ + baseURL := &url.URL{ Scheme: "https", Host: "quay.io", Path: path.Join("api", "v1", "repository"), } - query := url.Query() - query.Set("namespace", quayNamespace) - query.Set("public", "true") - query.Set("private", "true") - url.RawQuery = query.Encode() + allImages := []string{} + nextPageToken := "" - req, err := http.NewRequestWithContext(ctx, http.MethodGet, url.String(), http.NoBody) - if err != nil { - return nil, err - } - - if q.Token != "" { - req.Header.Set("Authorization", "Bearer "+q.Token) - } - - resp, err := defaultHTTPClient.Do(req) - if err != nil { - return nil, err - } - - defer func() { - _, _ = io.Copy(io.Discard, resp.Body) - _ = resp.Body.Close() - }() - - responseBodyByte, err := io.ReadAll(resp.Body) - if err != nil { - return nil, err - } - - switch resp.StatusCode { - case http.StatusOK: - var allImages = make([]string, 0) - - var listImagesResp quayResp - if err := json.Unmarshal(responseBodyByte, &listImagesResp); err != nil { + for { + if err := registryRateLimiter.Wait(ctx); err != nil { return nil, err } - for _, image := range listImagesResp.Repositories { + u := *baseURL + query := u.Query() + query.Set("namespace", quayNamespace) + query.Set("public", "true") + query.Set("private", "true") + + // Quay's API controls page size internally; we just fetch page by page. + if nextPageToken != "" { + query.Set("next_page", nextPageToken) + } + u.RawQuery = query.Encode() + + req, err := http.NewRequestWithContext(ctx, http.MethodGet, u.String(), http.NoBody) + if err != nil { + return nil, err + } + + if q.Token != "" { + req.Header.Set("Authorization", "Bearer "+q.Token) + } + + client := q.Client + if client == nil { + client = defaultHTTPClient + } + resp, err := client.Do(req) + if err != nil { + return nil, err + } + + body, err := io.ReadAll(resp.Body) + discardBody(resp) + if err != nil { + return nil, err + } + + if resp.StatusCode != http.StatusOK { + return nil, fmt.Errorf("failed to list quay images: unexpected status code: %d", resp.StatusCode) + } + + var page quayResp + if err := json.Unmarshal(body, &page); err != nil { + return nil, err + } + + for _, image := range page.Repositories { allImages = append(allImages, fmt.Sprintf("%s/%s", namespace, image.Name)) // quay.io// } - return allImages, nil - default: - return nil, fmt.Errorf("failed to list quay images: unexpected status code: %d", resp.StatusCode) + if !page.HasAdditional || page.NextPage == "" { + break + } + nextPageToken = page.NextPage } + + return allImages, nil } // === GHCR Registry === // GHCR implements the Registry interface for GHCR.io. type GHCR struct { - Token string // https://github.com/github/roadmap/issues/558 + Token string // https://github.com/github/roadmap/issues/558 + Client *http.Client } func (g *GHCR) Name() string { @@ -210,60 +282,134 @@ func (g *GHCR) WithRegistryToken(registryToken string) { g.Token = registryToken } -// ListImages lists all images under a Quay namespace. +func (g *GHCR) WithClient(client *http.Client) { + g.Client = client +} + +// GHCR paginates results and includes pagination links in the HTTP Link header. +// The Link header contains URLs for "next", "prev", "first", and "last" pages. +// Example Link header: +// +// ; rel="next", +// ; rel="last" +func parseNextLinkURL(linkHeader string) string { + if linkHeader == "" { + return "" + } + + parts := strings.Split(linkHeader, ",") + for _, part := range parts { + section := strings.Split(strings.TrimSpace(part), ";") + if len(section) < 2 { + continue + } + + linkPart := strings.TrimSpace(section[0]) + if !strings.HasPrefix(linkPart, "<") || !strings.HasSuffix(linkPart, ">") { + continue + } + urlStr := strings.Trim(linkPart, "<>") + + rel := "" + for _, attr := range section[1:] { + attr = strings.TrimSpace(attr) + if strings.HasPrefix(attr, "rel=") { + rel = strings.Trim(strings.TrimPrefix(attr, "rel="), "\"") + break + } + } + + if rel == "next" { + return urlStr + } + } + + return "" +} + +// ListImages lists all images under a GHCR namespace. +// For GitHub Container Registry the package listing endpoint is: +// +// GET https://api.github.com/users/{namespace}/packages?package_type=container&per_page=100 +// +// The GitHub API is paginated via the Link response header. func (g *GHCR) ListImages(ctx context.Context, namespace string) ([]string, error) { ghcrNamespace := path.Base(namespace) // ghcr.io/ -> namespace - url := &url.URL{ + baseURL := &url.URL{ Scheme: "https", Host: "api.github.com", Path: path.Join("users", ghcrNamespace, "packages"), } - query := url.Query() - query.Set("package_type", "container") - url.RawQuery = query.Encode() - - req, err := http.NewRequestWithContext(ctx, http.MethodGet, url.String(), http.NoBody) - if err != nil { - return nil, err - } - - // https://stackoverflow.com/questions/72732582/using-github-packages-without-personal-access-token - if g.Token != "" { - req.Header.Set("Authorization", "Bearer "+g.Token) - } - - resp, err := defaultHTTPClient.Do(req) - if err != nil { - return nil, err - } - - defer func() { - _, _ = io.Copy(io.Discard, resp.Body) - _ = resp.Body.Close() + allImages := []string{} + nextURL := func() string { + u := *baseURL + q := u.Query() + q.Set("package_type", "container") + q.Set("per_page", fmt.Sprint(maxRegistryPageSize)) // fetch images in batches of 100 per page + u.RawQuery = q.Encode() + return u.String() }() - responseBodyByte, err := io.ReadAll(resp.Body) - if err != nil { - return nil, err - } - - switch resp.StatusCode { - case http.StatusOK: - var allImages = make([]string, 0) - - var listImagesResp []Image - if err := json.Unmarshal(responseBodyByte, &listImagesResp); err != nil { + for nextURL != "" { + if err := registryRateLimiter.Wait(ctx); err != nil { return nil, err } - for _, image := range listImagesResp { + req, err := http.NewRequestWithContext(ctx, http.MethodGet, nextURL, http.NoBody) + if err != nil { + return nil, err + } + + // https://stackoverflow.com/questions/72732582/using-github-packages-without-personal-access-token + if g.Token != "" { + req.Header.Set("Authorization", "Bearer "+g.Token) + } + // GitHub recommends explicitly sending the v3 media type. + req.Header.Set("Accept", "application/vnd.github+json") + + client := g.Client + if client == nil { + client = defaultHTTPClient + } + resp, err := client.Do(req) + if err != nil { + return nil, err + } + + body, err := io.ReadAll(resp.Body) + discardBody(resp) + if err != nil { + return nil, err + } + + if resp.StatusCode != http.StatusOK { + return nil, fmt.Errorf("failed to list ghcr images: unexpected status code: %d", resp.StatusCode) + } + + // The GHCR packages list returns an array of package objects. We only + // care about the "name" field at this layer, so reuse the Image struct. + var page []Image + if err := json.Unmarshal(body, &page); err != nil { + return nil, err + } + + for _, image := range page { allImages = append(allImages, fmt.Sprintf("%s/%s", namespace, image.Name)) // ghcr.io// } - return allImages, nil - default: - return nil, fmt.Errorf("failed to list ghcr images: unexpected status code: %d", resp.StatusCode) + // Determine if there's another page via the Link header. + nextURL = parseNextLinkURL(resp.Header.Get("Link")) + } + + return allImages, nil +} + +// Function to discard response body +func discardBody(resp *http.Response) { + if resp != nil && resp.Body != nil { + _, _ = io.Copy(io.Discard, resp.Body) + _ = resp.Body.Close() } } diff --git a/pkg/sources/docker/registries_test.go b/pkg/sources/docker/registries_test.go index 657cf64a7..f6f3c0404 100644 --- a/pkg/sources/docker/registries_test.go +++ b/pkg/sources/docker/registries_test.go @@ -2,6 +2,7 @@ package docker import ( "fmt" + "net/http" "slices" "testing" @@ -60,3 +61,42 @@ func TestGHCRListImages(t *testing.T) { assert.Equal(t, ghcrImages, []string{"ghcr.io/mongodb/kingfisher"}) } + +func TestDockerHubListImages_RateLimitError(t *testing.T) { + t.Parallel() + + // Dockerhub registry + dockerhub := MakeRegistryFromNamespace("trufflesecurity") // no authentication + + // Cast dockerhub to *DockerHub registry to override the HTTP client + dockerhub.WithClient(common.ConstantResponseHttpClient(http.StatusTooManyRequests, "{}")) + dockerImages, err := dockerhub.ListImages(context.Background(), "trufflesecurity") // namespace without any prefix defaults to dockerhub registry + assert.Error(t, err) + assert.Nil(t, dockerImages) +} + +func TestQuayListImages_RateLimitError(t *testing.T) { + t.Parallel() + + // Quay.io registry + quay := MakeRegistryFromNamespace("quay.io/truffledockerman") // no authentication + // Cast quay to *Quay registry to override the HTTP client + quay.WithClient(common.ConstantResponseHttpClient(http.StatusTooManyRequests, "{}")) + + quayImages, err := quay.ListImages(context.Background(), "quay.io/truffledockerman") + assert.Error(t, err) + assert.Nil(t, quayImages) +} + +func TestGHCRListImages_RateLimitError(t *testing.T) { + t.Parallel() + + // GHCR registry + ghcr := MakeRegistryFromNamespace("ghcr.io/mongodb") // no authentication + // Cast ghcr to *GHCR registry to override the HTTP client + ghcr.WithClient(common.ConstantResponseHttpClient(http.StatusTooManyRequests, "{}")) + + ghcrImages, err := ghcr.ListImages(context.Background(), "ghcr.io/mongodb") + assert.Error(t, err) + assert.Nil(t, ghcrImages) +}