Merge pull request #3662 from cpuguy83/set_custom_scopes_for_authorizer

crosbymichael · web-flow · commit 6fafc8a724ef · 2019-09-18T15:10:51.000-04:00
Add function to set custom auth scope in context
diff --git a/remotes/docker/scope.go b/remotes/docker/scope.go
@@ -51,21 +51,27 @@ func contextWithRepositoryScope(ctx context.Context, refspec reference.Spec, pus
 	if err != nil {
 		return nil, err
 	}
-	return context.WithValue(ctx, tokenScopesKey{}, []string{s}), nil
+	return WithScope(ctx, s), nil
 }
 
-// contextWithAppendPullRepositoryScope is used to append repository pull
-// scope into existing scopes indexed by the tokenScopesKey{}.
-func contextWithAppendPullRepositoryScope(ctx context.Context, repo string) context.Context {
+// WithScope appends a custom registry auth scope to the context.
+func WithScope(ctx context.Context, scope string) context.Context {
 	var scopes []string
-
 	if v := ctx.Value(tokenScopesKey{}); v != nil {
-		scopes = append(scopes, v.([]string)...)
+		scopes = v.([]string)
+		scopes = append(scopes, scope)
+	} else {
+		scopes = []string{scope}
 	}
-	scopes = append(scopes, fmt.Sprintf("repository:%s:pull", repo))
 	return context.WithValue(ctx, tokenScopesKey{}, scopes)
 }
 
+// contextWithAppendPullRepositoryScope is used to append repository pull
+// scope into existing scopes indexed by the tokenScopesKey{}.
+func contextWithAppendPullRepositoryScope(ctx context.Context, repo string) context.Context {
+	return WithScope(ctx, fmt.Sprintf("repository:%s:pull", repo))
+}
+
 // getTokenScopes returns deduplicated and sorted scopes from ctx.Value(tokenScopesKey{}) and common scopes.
 func getTokenScopes(ctx context.Context, common []string) []string {
 	var scopes []string
diff --git a/remotes/docker/scope_test.go b/remotes/docker/scope_test.go
@@ -22,6 +22,7 @@ import (
 
 	"github.com/containerd/containerd/reference"
 	"gotest.tools/assert"
+	"gotest.tools/assert/cmp"
 )
 
 func TestRepositoryScope(t *testing.T) {
@@ -94,3 +95,14 @@ func TestGetTokenScopes(t *testing.T) {
 		assert.DeepEqual(t, tc.expected, actual)
 	}
 }
+
+func TestCustomScope(t *testing.T) {
+	scope := "whatever:foo/bar:pull"
+	ctx := WithScope(context.Background(), scope)
+	ctx = contextWithAppendPullRepositoryScope(ctx, "foo/bar")
+
+	scopes := getTokenScopes(ctx, []string{})
+	assert.Assert(t, cmp.Len(scopes, 2))
+	assert.Check(t, cmp.Equal(scopes[0], "repository:foo/bar:pull"))
+	assert.Check(t, cmp.Equal(scopes[1], scope))
+}
