Merge pull request #983 from ktock/header

remote: Support per-registry request headers

docs/overview.md

@@ -176,6 +176,17 @@ host = "exampleregistry.io"
 insecure = true
 ```
 
+`header` field allows to set headers to send to the server.
+
+```toml
+[[resolver.host."registry2:5000".mirrors]]
+  host = "registry2:5000"
+  [resolver.host."registry2:5000".mirrors.header]
+    x-custom-2 = ["value3", "value4"]
+```
+
+> NOTE: Headers aren't passed to the redirected location.
+
 The config file can be passed to stargz snapshotter using `containerd-stargz-grpc`'s `--config` option.
 
 ## Make your remote snapshotter

fs/remote/resolver.go

@@ -238,7 +238,7 @@ func newHTTPFetcher(ctx context.Context, fc *fetcherConfig) (*httpFetcher, int64
 			path.Join(host.Host, host.Path),
 			strings.TrimPrefix(fc.refspec.Locator, fc.refspec.Hostname()+"/"),
 			digest)
-		url, err := redirect(ctx, blobURL, tr, timeout)
+		url, header, err := redirect(ctx, blobURL, tr, timeout, host.Header)
 		if err != nil {
 			rErr = fmt.Errorf("failed to redirect (host %q, ref:%q, digest:%q): %v: %w", host.Host, fc.refspec, digest, err, rErr)
 			continue // Try another
@@ -247,7 +247,7 @@ func newHTTPFetcher(ctx context.Context, fc *fetcherConfig) (*httpFetcher, int64
 		// Get size information
 		// TODO: we should try to use the Size field in the descriptor here.
 		start := time.Now() // start time before getting layer header
-		size, err := getSize(ctx, url, tr, timeout)
+		size, err := getSize(ctx, url, tr, timeout, header)
 		commonmetrics.MeasureLatencyInMilliseconds(commonmetrics.StargzHeaderGet, digest, start) // time to get layer header
 		if err != nil {
 			rErr = fmt.Errorf("failed to get size (host %q, ref:%q, digest:%q): %v: %w", host.Host, fc.refspec, digest, err, rErr)
@@ -256,11 +256,13 @@ func newHTTPFetcher(ctx context.Context, fc *fetcherConfig) (*httpFetcher, int64
 
 		// Hit one destination
 		return &httpFetcher{
-			url:     url,
-			tr:      tr,
-			blobURL: blobURL,
-			digest:  digest,
-			timeout: timeout,
+			url:       url,
+			tr:        tr,
+			blobURL:   blobURL,
+			digest:    digest,
+			timeout:   timeout,
+			header:    header,
+			orgHeader: host.Header,
 		}, size, nil
 	}
 
@@ -309,7 +311,7 @@ func (tr *transport) RoundTrip(req *http.Request) (*http.Response, error) {
 	return resp, nil
 }
 
-func redirect(ctx context.Context, blobURL string, tr http.RoundTripper, timeout time.Duration) (url string, err error) {
+func redirect(ctx context.Context, blobURL string, tr http.RoundTripper, timeout time.Duration, header http.Header) (url string, withHeader http.Header, err error) {
 	if timeout > 0 {
 		var cancel context.CancelFunc
 		ctx, cancel = context.WithTimeout(ctx, timeout)
@@ -320,13 +322,17 @@ func redirect(ctx context.Context, blobURL string, tr http.RoundTripper, timeout
 	// ghcr.io returns 200 on HEAD without Location header (2020).
 	req, err := http.NewRequestWithContext(ctx, "GET", blobURL, nil)
 	if err != nil {
-		return "", fmt.Errorf("failed to make request to the registry: %w", err)
+		return "", nil, fmt.Errorf("failed to make request to the registry: %w", err)
+	}
+	req.Header = http.Header{}
+	for k, v := range header {
+		req.Header[k] = v
 	}
 	req.Close = false
 	req.Header.Set("Range", "bytes=0-1")
 	res, err := tr.RoundTrip(req)
 	if err != nil {
-		return "", fmt.Errorf("failed to request: %w", err)
+		return "", nil, fmt.Errorf("failed to request: %w", err)
 	}
 	defer func() {
 		io.Copy(io.Discard, res.Body)
@@ -335,17 +341,19 @@ func redirect(ctx context.Context, blobURL string, tr http.RoundTripper, timeout
 
 	if res.StatusCode/100 == 2 {
 		url = blobURL
+		withHeader = header
 	} else if redir := res.Header.Get("Location"); redir != "" && res.StatusCode/100 == 3 {
 		// TODO: Support nested redirection
 		url = redir
+		// Do not pass headers to the redirected location.
 	} else {
-		return "", fmt.Errorf("failed to access to the registry with code %v", res.StatusCode)
+		return "", nil, fmt.Errorf("failed to access to the registry with code %v", res.StatusCode)
 	}
 
 	return
 }
 
-func getSize(ctx context.Context, url string, tr http.RoundTripper, timeout time.Duration) (int64, error) {
+func getSize(ctx context.Context, url string, tr http.RoundTripper, timeout time.Duration, header http.Header) (int64, error) {
 	if timeout > 0 {
 		var cancel context.CancelFunc
 		ctx, cancel = context.WithTimeout(ctx, timeout)
@@ -355,6 +363,10 @@ func getSize(ctx context.Context, url string, tr http.RoundTripper, timeout time
 	if err != nil {
 		return 0, err
 	}
+	req.Header = http.Header{}
+	for k, v := range header {
+		req.Header[k] = v
+	}
 	req.Close = false
 	res, err := tr.RoundTrip(req)
 	if err != nil {
@@ -373,6 +385,10 @@ func getSize(ctx context.Context, url string, tr http.RoundTripper, timeout time
 	if err != nil {
 		return 0, fmt.Errorf("failed to make request to the registry: %w", err)
 	}
+	req.Header = http.Header{}
+	for k, v := range header {
+		req.Header[k] = v
+	}
 	req.Close = false
 	req.Header.Set("Range", "bytes=0-1")
 	res, err = tr.RoundTrip(req)
@@ -404,6 +420,8 @@ type httpFetcher struct {
 	singleRange   bool
 	singleRangeMu sync.Mutex
 	timeout       time.Duration
+	header        http.Header
+	orgHeader     http.Header
 }
 
 type multipartReadCloser interface {
@@ -443,6 +461,10 @@ func (f *httpFetcher) fetch(ctx context.Context, rs []region, retry bool) (multi
 	if err != nil {
 		return nil, err
 	}
+	req.Header = http.Header{}
+	for k, v := range f.header {
+		req.Header[k] = v
+	}
 	var ranges string
 	for _, reg := range requests {
 		ranges += fmt.Sprintf("%d-%d,", reg.b, reg.e)
@@ -514,6 +536,10 @@ func (f *httpFetcher) check() error {
 	if err != nil {
 		return fmt.Errorf("check failed: failed to make request: %w", err)
 	}
+	req.Header = http.Header{}
+	for k, v := range f.header {
+		req.Header[k] = v
+	}
 	req.Close = false
 	req.Header.Set("Range", "bytes=0-1")
 	res, err := f.tr.RoundTrip(req)
@@ -544,12 +570,13 @@ func (f *httpFetcher) check() error {
 }
 
 func (f *httpFetcher) refreshURL(ctx context.Context) error {
-	newURL, err := redirect(ctx, f.blobURL, f.tr, f.timeout)
+	newURL, headers, err := redirect(ctx, f.blobURL, f.tr, f.timeout, f.orgHeader)
 	if err != nil {
 		return err
 	}
 	f.urlMu.Lock()
 	f.url = newURL
+	f.header = headers
 	f.urlMu.Unlock()
 	return nil
 }