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 <[email protected]> Co-authored-by: Kashif Khan <[email protected]>
2025-12-15 13:09:11 +05:00
parent 0fab92f434
commit 3cb2166ddd
2 changed files with 309 additions and 123 deletions
@@ -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/<ns>/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)) // <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=<namespace>&public=true&private=true&next_page=<token>
+//
+// 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> -> 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/<namespace>/<image_name>
 		}

-		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:
+//
+// <https://api.github.com/user/abc/packages?package_type=container&per_page=100&page=2>; rel="next",
+// <https://api.github.com/user/abc/packages?package_type=container&per_page=100&page=5>; 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> -> 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/<namespace>/<image_name>
 		}

-		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()
 	}
 }
@@ -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)
+}