Repository: google/go-containerregistry
Pull request #834 (2020-11-17 20:21:51)

Plumb contexts through transport.New
Fixes https://github.com/google/go-containerregistry/issues/832

Add transport.NewWithContext.

Use the provided context for the initial ping request and token
exchange.

Subsequent token exchanges (to refresh) will use the context of the
incoming request in RoundTrip.

Drop the hardcoded ping timeout.

Use transport.NewWithContext everywhere that's easily plumbable.

jonjohnsonjr: cc @CodingJzy this will revert https://github.com/google/go-containerregistry/pull/808 and use the context from `remote.WithContext` for most cases, just FYI.
CodingJzy: > cc @CodingJzy，这将还原[＃808](https://github.com/google/go-containerregistry/pull/808)并在`remote.WithContext`大多数情况下使用FYI中的上下文。

thank you，Your way of passing in the context parameter is correct. I just discovered the problem, and then temporarily added a 2s timeout, which did not fundamentally solve the problem.
jonjohnsonjr: Thanks for making us aware of the problem @CodingJzy 😃 

Patch:
diff --git a/pkg/internal/legacy/copy.go b/pkg/internal/legacy/copy.go
index ba8f96ab..557c1d1f 100644
--- a/pkg/internal/legacy/copy.go
+++ b/pkg/internal/legacy/copy.go
@@ -58,6 +58,8 @@ func CopySchema1(desc *remote.Descriptor, srcRef, dstRef name.Reference, srcAuth
 func putManifest(desc *remote.Descriptor, dstRef name.Reference, dstAuth authn.Authenticator) error {
 	reg := dstRef.Context().Registry
 	scopes := []string{dstRef.Scope(transport.PushScope)}
+
+	// TODO(jonjohnsonjr): Use NewWithContext.
 	tr, err := transport.New(reg, dstAuth, http.DefaultTransport, scopes)
 	if err != nil {
 		return err
diff --git a/pkg/v1/google/list.go b/pkg/v1/google/list.go
index 5db3e4a8..cdc7d486 100644
--- a/pkg/v1/google/list.go
+++ b/pkg/v1/google/list.go
@@ -66,7 +66,7 @@ func newLister(repo name.Repository, options ...ListerOption) (*lister, error) {
 	l.transport = transport.NewRetry(l.transport)
 
 	scopes := []string{repo.Scope(transport.PullScope)}
-	tr, err := transport.New(repo.Registry, l.auth, l.transport, scopes)
+	tr, err := transport.NewWithContext(l.ctx, repo.Registry, l.auth, l.transport, scopes)
 	if err != nil {
 		return nil, err
 	}
diff --git a/pkg/v1/remote/catalog.go b/pkg/v1/remote/catalog.go
index e0794dc7..21b5dbba 100644
--- a/pkg/v1/remote/catalog.go
+++ b/pkg/v1/remote/catalog.go
@@ -37,7 +37,7 @@ func CatalogPage(target name.Registry, last string, n int, options ...Option) ([
 	}
 
 	scopes := []string{target.Scope(transport.PullScope)}
-	tr, err := transport.New(target, o.auth, o.transport, scopes)
+	tr, err := transport.NewWithContext(o.context, target, o.auth, o.transport, scopes)
 	if err != nil {
 		return nil, err
 	}
@@ -82,7 +82,7 @@ func Catalog(ctx context.Context, target name.Registry, options ...Option) ([]st
 	}
 
 	scopes := []string{target.Scope(transport.PullScope)}
-	tr, err := transport.New(target, o.auth, o.transport, scopes)
+	tr, err := transport.NewWithContext(o.context, target, o.auth, o.transport, scopes)
 	if err != nil {
 		return nil, err
 	}
diff --git a/pkg/v1/remote/delete.go b/pkg/v1/remote/delete.go
index 1921ae4e..3b902271 100644
--- a/pkg/v1/remote/delete.go
+++ b/pkg/v1/remote/delete.go
@@ -30,7 +30,7 @@ func Delete(ref name.Reference, options ...Option) error {
 		return err
 	}
 	scopes := []string{ref.Scope(transport.DeleteScope)}
-	tr, err := transport.New(ref.Context().Registry, o.auth, o.transport, scopes)
+	tr, err := transport.NewWithContext(o.context, ref.Context().Registry, o.auth, o.transport, scopes)
 	if err != nil {
 		return err
 	}
diff --git a/pkg/v1/remote/descriptor.go b/pkg/v1/remote/descriptor.go
index 653c175d..b77cb9a8 100644
--- a/pkg/v1/remote/descriptor.go
+++ b/pkg/v1/remote/descriptor.go
@@ -218,7 +218,7 @@ type fetcher struct {
 }
 
 func makeFetcher(ref name.Reference, o *options) (*fetcher, error) {
-	tr, err := transport.New(ref.Context().Registry, o.auth, o.transport, []string{ref.Scope(transport.PullScope)})
+	tr, err := transport.NewWithContext(o.context, ref.Context().Registry, o.auth, o.transport, []string{ref.Scope(transport.PullScope)})
 	if err != nil {
 		return nil, err
 	}
diff --git a/pkg/v1/remote/list.go b/pkg/v1/remote/list.go
index ea4cacc5..e4a005aa 100644
--- a/pkg/v1/remote/list.go
+++ b/pkg/v1/remote/list.go
@@ -44,7 +44,7 @@ func ListWithContext(ctx context.Context, repo name.Repository, options ...Optio
 		return nil, err
 	}
 	scopes := []string{repo.Scope(transport.PullScope)}
-	tr, err := transport.New(repo.Registry, o.auth, o.transport, scopes)
+	tr, err := transport.NewWithContext(o.context, repo.Registry, o.auth, o.transport, scopes)
 	if err != nil {
 		return nil, err
 	}
diff --git a/pkg/v1/remote/multi_write.go b/pkg/v1/remote/multi_write.go
index 631ab23f..803c2293 100644
--- a/pkg/v1/remote/multi_write.go
+++ b/pkg/v1/remote/multi_write.go
@@ -81,7 +81,7 @@ func MultiWrite(m map[name.Reference]Taggable, options ...Option) error {
 		ls = append(ls, l)
 	}
 	scopes := scopesForUploadingImage(repo, ls)
-	tr, err := transport.New(repo.Registry, o.auth, o.transport, scopes)
+	tr, err := transport.NewWithContext(o.context, repo.Registry, o.auth, o.transport, scopes)
 	if err != nil {
 		return err
 	}
diff --git a/pkg/v1/remote/transport/bearer.go b/pkg/v1/remote/transport/bearer.go
index e851ec59..e83819f7 100644
--- a/pkg/v1/remote/transport/bearer.go
+++ b/pkg/v1/remote/transport/bearer.go
@@ -15,6 +15,7 @@
 package transport
 
 import (
+	"context"
 	"encoding/json"
 	"fmt"
 	"io/ioutil"
@@ -75,7 +76,7 @@ func (bt *bearerTransport) RoundTrip(in *http.Request) (*http.Response, error) {
 
 	// Perform a token refresh() and retry the request in case the token has expired
 	if res.StatusCode == http.StatusUnauthorized {
-		if err = bt.refresh(); err != nil {
+		if err = bt.refresh(in.Context()); err != nil {
 			return nil, err
 		}
 		return sendRequest()
@@ -88,7 +89,7 @@ func (bt *bearerTransport) RoundTrip(in *http.Request) (*http.Response, error) {
 // so we rely on heuristics and fallbacks to support as many registries as possible.
 // The basic token exchange is attempted first, falling back to the oauth flow.
 // If the IdentityToken is set, this indicates that we should start with the oauth flow.
-func (bt *bearerTransport) refresh() error {
+func (bt *bearerTransport) refresh(ctx context.Context) error {
 	auth, err := bt.basic.Authorization()
 	if err != nil {
 		return err
@@ -104,15 +105,15 @@ func (bt *bearerTransport) refresh() error {
 		// If the secret being stored is an identity token,
 		// the Username should be set to <token>, which indicates
 		// we are using an oauth flow.
-		content, err = bt.refreshOauth()
+		content, err = bt.refreshOauth(ctx)
 		if terr, ok := err.(*Error); ok && terr.StatusCode == http.StatusNotFound {
 			// Note: Not all token servers implement oauth2.
 			// If the request to the endpoint returns 404 using the HTTP POST method,
 			// refer to Token Documentation for using the HTTP GET method supported by all token servers.
-			content, err = bt.refreshBasic()
+			content, err = bt.refreshBasic(ctx)
 		}
 	} else {
-		content, err = bt.refreshBasic()
+		content, err = bt.refreshBasic(ctx)
 	}
 	if err != nil {
 		return err
@@ -186,7 +187,7 @@ func canonicalAddress(host, scheme string) (address string) {
 }
 
 // https://docs.docker.com/registry/spec/auth/oauth/
-func (bt *bearerTransport) refreshOauth() ([]byte, error) {
+func (bt *bearerTransport) refreshOauth(ctx context.Context) ([]byte, error) {
 	auth, err := bt.basic.Authorization()
 	if err != nil {
 		return nil, err
@@ -220,7 +221,7 @@ func (bt *bearerTransport) refreshOauth() ([]byte, error) {
 	req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
 
 	// We don't want to log credentials.
-	ctx := redact.NewContext(req.Context(), "oauth token response contains credentials")
+	ctx = redact.NewContext(ctx, "oauth token response contains credentials")
 
 	resp, err := client.Do(req.WithContext(ctx))
 	if err != nil {
@@ -236,7 +237,7 @@ func (bt *bearerTransport) refreshOauth() ([]byte, error) {
 }
 
 // https://docs.docker.com/registry/spec/auth/token/
-func (bt *bearerTransport) refreshBasic() ([]byte, error) {
+func (bt *bearerTransport) refreshBasic(ctx context.Context) ([]byte, error) {
 	u, err := url.Parse(bt.realm)
 	if err != nil {
 		return nil, err
@@ -259,7 +260,7 @@ func (bt *bearerTransport) refreshBasic() ([]byte, error) {
 	}
 
 	// We don't want to log credentials.
-	ctx := redact.NewContext(req.Context(), "basic token response contains credentials")
+	ctx = redact.NewContext(ctx, "basic token response contains credentials")
 
 	resp, err := client.Do(req.WithContext(ctx))
 	if err != nil {
diff --git a/pkg/v1/remote/transport/ping.go b/pkg/v1/remote/transport/ping.go
index 08b2214c..77e13a1e 100644
--- a/pkg/v1/remote/transport/ping.go
+++ b/pkg/v1/remote/transport/ping.go
@@ -15,12 +15,12 @@
 package transport
 
 import (
+	"context"
 	"fmt"
 	"io"
 	"io/ioutil"
 	"net/http"
 	"strings"
-	"time"
 
 	"github.com/google/go-containerregistry/pkg/name"
 )
@@ -66,8 +66,8 @@ func parseChallenge(suffix string) map[string]string {
 	return kv
 }
 
-func ping(reg name.Registry, t http.RoundTripper) (*pingResp, error) {
-	client := http.Client{Transport: t, Timeout: 120 * time.Second}
+func ping(ctx context.Context, reg name.Registry, t http.RoundTripper) (*pingResp, error) {
+	client := http.Client{Transport: t}
 
 	// This first attempts to use "https" for every request, falling back to http
 	// if the registry matches our localhost heuristic or if it is intentionally
@@ -80,7 +80,11 @@ func ping(reg name.Registry, t http.RoundTripper) (*pingResp, error) {
 	var connErr error
 	for _, scheme := range schemes {
 		url := fmt.Sprintf("%s://%s/v2/", scheme, reg.Name())
-		resp, err := client.Get(url)
+		req, err := http.NewRequest(http.MethodGet, url, nil)
+		if err != nil {
+			return nil, err
+		}
+		resp, err := client.Do(req.WithContext(ctx))
 		if err != nil {
 			connErr = err
 			// Potentially retry with http.
diff --git a/pkg/v1/remote/transport/transport.go b/pkg/v1/remote/transport/transport.go
index 1bea604a..496ba39a 100644
--- a/pkg/v1/remote/transport/transport.go
+++ b/pkg/v1/remote/transport/transport.go
@@ -15,6 +15,7 @@
 package transport
 
 import (
+	"context"
 	"fmt"
 	"net/http"
 
@@ -25,7 +26,16 @@ import (
 // New returns a new RoundTripper based on the provided RoundTripper that has been
 // setup to authenticate with the remote registry "reg", in the capacity
 // laid out by the specified scopes.
+//
+// TODO(jonjohnsonjr): Deprecate this.
 func New(reg name.Registry, auth authn.Authenticator, t http.RoundTripper, scopes []string) (http.RoundTripper, error) {
+	return NewWithContext(context.Background(), reg, auth, t, scopes)
+}
+
+// NewWithContext returns a new RoundTripper based on the provided RoundTripper that has been
+// setup to authenticate with the remote registry "reg", in the capacity
+// laid out by the specified scopes.
+func NewWithContext(ctx context.Context, reg name.Registry, auth authn.Authenticator, t http.RoundTripper, scopes []string) (http.RoundTripper, error) {
 	// The handshake:
 	//  1. Use "t" to ping() the registry for the authentication challenge.
 	//
@@ -40,7 +50,7 @@ func New(reg name.Registry, auth authn.Authenticator, t http.RoundTripper, scope
 
 	// First we ping the registry to determine the parameters of the authentication handshake
 	// (if one is even necessary).
-	pr, err := ping(reg, t)
+	pr, err := ping(ctx, reg, t)
 	if err != nil {
 		return nil, err
 	}
@@ -81,7 +91,7 @@ func New(reg name.Registry, auth authn.Authenticator, t http.RoundTripper, scope
 			scopes:   scopes,
 			scheme:   pr.scheme,
 		}
-		if err := bt.refresh(); err != nil {
+		if err := bt.refresh(ctx); err != nil {
 			return nil, err
 		}
 		return bt, nil
diff --git a/pkg/v1/remote/write.go b/pkg/v1/remote/write.go
index 0aa7a0e4..521b2250 100644
--- a/pkg/v1/remote/write.go
+++ b/pkg/v1/remote/write.go
@@ -54,7 +54,7 @@ func Write(ref name.Reference, img v1.Image, options ...Option) error {
 	}
 
 	scopes := scopesForUploadingImage(ref.Context(), ls)
-	tr, err := transport.New(ref.Context().Registry, o.auth, o.transport, scopes)
+	tr, err := transport.NewWithContext(o.context, ref.Context().Registry, o.auth, o.transport, scopes)
 	if err != nil {
 		return err
 	}
@@ -549,7 +549,7 @@ func WriteIndex(ref name.Reference, ii v1.ImageIndex, options ...Option) error {
 		return err
 	}
 	scopes := []string{ref.Scope(transport.PushScope)}
-	tr, err := transport.New(ref.Context().Registry, o.auth, o.transport, scopes)
+	tr, err := transport.NewWithContext(o.context, ref.Context().Registry, o.auth, o.transport, scopes)
 	if err != nil {
 		return err
 	}
@@ -568,7 +568,7 @@ func WriteLayer(repo name.Repository, layer v1.Layer, options ...Option) error {
 		return err
 	}
 	scopes := scopesForUploadingImage(repo, []v1.Layer{layer})
-	tr, err := transport.New(repo.Registry, o.auth, o.transport, scopes)
+	tr, err := transport.NewWithContext(o.context, repo.Registry, o.auth, o.transport, scopes)
 	if err != nil {
 		return err
 	}
@@ -595,7 +595,7 @@ func Tag(tag name.Tag, t Taggable, options ...Option) error {
 	// * Allow callers to pass in a transport.Transport, typecheck
 	//   it to allow them to reuse the transport across multiple calls.
 	// * WithTag option to do multiple manifest PUTs in commitManifest.
-	tr, err := transport.New(tag.Context().Registry, o.auth, o.transport, scopes)
+	tr, err := transport.NewWithContext(o.context, tag.Context().Registry, o.auth, o.transport, scopes)
 	if err != nil {
 		return err
 	}


Test patch:
diff --git a/pkg/v1/remote/transport/bearer_test.go b/pkg/v1/remote/transport/bearer_test.go
index 9a86f160..c0f707ad 100644
--- a/pkg/v1/remote/transport/bearer_test.go
+++ b/pkg/v1/remote/transport/bearer_test.go
@@ -15,6 +15,7 @@
 package transport
 
 import (
+	"context"
 	"fmt"
 	"net/http"
 	"net/http/httptest"
@@ -79,7 +80,7 @@ func TestBearerRefresh(t *testing.T) {
 				scheme:   "http",
 			}
 
-			if err := bt.refresh(); (err != nil) != tc.wantErr {
+			if err := bt.refresh(context.Background()); (err != nil) != tc.wantErr {
 				t.Errorf("refresh() = %v", err)
 			}
 		})
diff --git a/pkg/v1/remote/transport/ping_test.go b/pkg/v1/remote/transport/ping_test.go
index df8c09d9..4d17e810 100644
--- a/pkg/v1/remote/transport/ping_test.go
+++ b/pkg/v1/remote/transport/ping_test.go
@@ -15,6 +15,7 @@
 package transport
 
 import (
+	"context"
 	"net/http"
 	"net/http/httptest"
 	"net/url"
@@ -83,7 +84,7 @@ func TestPingNoChallenge(t *testing.T) {
 		},
 	}
 
-	pr, err := ping(testRegistry, tprt)
+	pr, err := ping(context.Background(), testRegistry, tprt)
 	if err != nil {
 		t.Errorf("ping() = %v", err)
 	}
@@ -108,7 +109,7 @@ func TestPingBasicChallengeNoParams(t *testing.T) {
 		},
 	}
 
-	pr, err := ping(testRegistry, tprt)
+	pr, err := ping(context.Background(), testRegistry, tprt)
 	if err != nil {
 		t.Errorf("ping() = %v", err)
 	}
@@ -133,7 +134,7 @@ func TestPingBearerChallengeWithParams(t *testing.T) {
 		},
 	}
 
-	pr, err := ping(testRegistry, tprt)
+	pr, err := ping(context.Background(), testRegistry, tprt)
 	if err != nil {
 		t.Errorf("ping() = %v", err)
 	}
@@ -158,7 +159,7 @@ func TestUnsupportedStatus(t *testing.T) {
 		},
 	}
 
-	pr, err := ping(testRegistry, tprt)
+	pr, err := ping(context.Background(), testRegistry, tprt)
 	if err == nil {
 		t.Errorf("ping() = %v", pr)
 	}
@@ -194,7 +195,7 @@ func TestPingHttpFallback(t *testing.T) {
 	}
 
 	for _, test := range tests {
-		pr, err := ping(test.reg, tprt)
+		pr, err := ping(context.Background(), test.reg, tprt)
 		if err == nil {
 			t.Errorf("ping() = %v", pr)
 		}