From 941cf10a8ebe14d2b03bf7253731134629fc7f80 Mon Sep 17 00:00:00 2001
From: Patrick Jones <ithuriel@google.com>
Date: Thu, 24 Jun 2021 16:24:14 -0700
Subject: [PATCH] downscope: move validation checks

---
 google/downscope/downscoping.go  | 38 ++++++++++++++++----------------
 google/downscope/example_test.go |  6 +++--
 2 files changed, 23 insertions(+), 21 deletions(-)

diff --git a/google/downscope/downscoping.go b/google/downscope/downscoping.go
index 48ae8f2..2d74c37 100644
--- a/google/downscope/downscoping.go
+++ b/google/downscope/downscoping.go
@@ -92,8 +92,25 @@ type downscopingTokenSource struct {
 }
 
 // NewTokenSource returns an empty downscopingTokenSource.
-func NewTokenSource(ctx context.Context, conf DownscopingConfig) oauth2.TokenSource {
-	return downscopingTokenSource{ctx: ctx, config: conf}
+func NewTokenSource(ctx context.Context, conf DownscopingConfig) (oauth2.TokenSource, error) {
+	if conf.RootSource == nil {
+		return nil, fmt.Errorf("downscope: rootSource cannot be nil")
+	}
+	if len(conf.Rules) == 0 {
+		return nil, fmt.Errorf("downscope: length of AccessBoundaryRules must be at least 1")
+	}
+	if len(conf.Rules) > 10 {
+		return nil, fmt.Errorf("downscope: length of AccessBoundaryRules may not be greater than 10")
+	}
+	for _, val := range conf.Rules {
+		if val.AvailableResource == "" {
+			return nil, fmt.Errorf("downscope: all rules must have a nonempty AvailableResource: %+v", val)
+		}
+		if len(val.AvailablePermissions) == 0 {
+			return nil, fmt.Errorf("downscope: all rules must provide at least one permission: %+v", val)
+		}
+	}
+	return downscopingTokenSource{ctx: ctx, config: conf}, nil
 }
 
 // Token() uses a downscopingTokenSource to generate an oauth2 Token.
@@ -102,23 +119,6 @@ func NewTokenSource(ctx context.Context, conf DownscopingConfig) oauth2.TokenSou
 // TokenSource struct with the Token held by the StaticTokenSource and wrap
 // that TokenSource in an oauth2.ReuseTokenSource.
 func (dts downscopingTokenSource) Token() (*oauth2.Token, error) {
-	if dts.config.RootSource == nil {
-		return nil, fmt.Errorf("downscope: rootSource cannot be nil")
-	}
-	if len(dts.config.Rules) == 0 {
-		return nil, fmt.Errorf("downscope: length of AccessBoundaryRules must be at least 1")
-	}
-	if len(dts.config.Rules) > 10 {
-		return nil, fmt.Errorf("downscope: length of AccessBoundaryRules may not be greater than 10")
-	}
-	for _, val := range dts.config.Rules {
-		if val.AvailableResource == "" {
-			return nil, fmt.Errorf("downscope: all rules must have a nonempty AvailableResource: %+v", val)
-		}
-		if len(val.AvailablePermissions) == 0 {
-			return nil, fmt.Errorf("downscope: all rules must provide at least one permission: %+v", val)
-		}
-	}
 
 	downscopedOptions := struct {
 		Boundary accessBoundary `json:"accessBoundary"`
diff --git a/google/downscope/example_test.go b/google/downscope/example_test.go
index cffd4ff..061cf57 100644
--- a/google/downscope/example_test.go
+++ b/google/downscope/example_test.go
@@ -27,8 +27,10 @@ func ExampleNewTokenSource() {
 
 	// rootSource, err := google.DefaultTokenSource(ctx, "https://www.googleapis.com/auth/cloud-platform")
 
-	dts := downscope.NewTokenSource(ctx, downscope.DownscopingConfig{RootSource: rootSource, Rules: accessBoundary})
-	_ = dts
+	dts, err := downscope.NewTokenSource(ctx, downscope.DownscopingConfig{RootSource: rootSource, Rules: accessBoundary})
+	if err != nil {
+		_ = dts
+	}
 	// You can now use the token held in myTokenSource to make
 	// Google Cloud Storage calls, as follows: