From a4ca1949a378ef7f0913b1e81e494156b0bf19cc Mon Sep 17 00:00:00 2001
From: Shin Fan <shinfan@google.com>
Date: Tue, 15 Jun 2021 11:15:03 -0700
Subject: [PATCH] Address comments

---
 google/jwt.go      | 10 +++++---
 google/jwt_test.go | 64 ++++++++++++++++++++++------------------------
 2 files changed, 38 insertions(+), 36 deletions(-)

diff --git a/google/jwt.go b/google/jwt.go
index efbd46b..67d97b9 100644
--- a/google/jwt.go
+++ b/google/jwt.go
@@ -32,17 +32,21 @@ func JWTAccessTokenSourceFromJSON(jsonKey []byte, audience string) (oauth2.Token
 // key file to read the credentials that authorize and authenticate the
 // requests, and returns a TokenSource that does not use any OAuth2 flow but
 // instead creates a JWT and sends that as the access token.
-// The scopes is typically a list of URLs that specifies the scope of the
+// The scope is typically a list of URLs that specifies the scope of the
 // credentials.
 //
 // Note that this is not a standard OAuth flow, but rather an
 // optimization supported by a few Google services.
 // Unless you know otherwise, you should use JWTConfigFromJSON instead.
-func JWTAccessTokenSourceWithScope(jsonKey []byte, scopes []string) (oauth2.TokenSource, error) {
-	return newJWTSource(jsonKey, "", scopes)
+func JWTAccessTokenSourceWithScope(jsonKey []byte, scope ...string) (oauth2.TokenSource, error) {
+	return newJWTSource(jsonKey, "", scope)
 }
 
 func newJWTSource(jsonKey []byte, audience string, scopes []string) (oauth2.TokenSource, error) {
+	if len(scopes) == 0 && audience == "" {
+		return nil, fmt.Errorf("google: missing scope/audience for JWT access token")
+	}
+
 	cfg, err := JWTConfigFromJSON(jsonKey)
 	if err != nil {
 		return nil, fmt.Errorf("google: could not parse JSON key: %v", err)
diff --git a/google/jwt_test.go b/google/jwt_test.go
index 28be9fa..3bfcb7e 100644
--- a/google/jwt_test.go
+++ b/google/jwt_test.go
@@ -13,29 +13,21 @@ import (
 	"encoding/json"
 	"encoding/pem"
 	"strings"
+	"sync"
 	"testing"
 	"time"
 
 	"golang.org/x/oauth2/jws"
 )
 
-func TestJWTAccessTokenSourceFromJSON(t *testing.T) {
-	// Generate a key we can use in the test data.
-	privateKey, err := rsa.GenerateKey(rand.Reader, 2048)
-	if err != nil {
-		t.Fatal(err)
-	}
+var (
+	privateKey *rsa.PrivateKey
+	jsonKey    []byte
+	once sync.Once
+)
 
-	// Encode the key and substitute into our example JSON.
-	enc := pem.EncodeToMemory(&pem.Block{
-		Type:  "PRIVATE KEY",
-		Bytes: x509.MarshalPKCS1PrivateKey(privateKey),
-	})
-	enc, err = json.Marshal(string(enc))
-	if err != nil {
-		t.Fatalf("json.Marshal: %v", err)
-	}
-	jsonKey := bytes.Replace(jwtJSONKey, []byte(`"super secret key"`), enc, 1)
+func TestJWTAccessTokenSourceFromJSON(t *testing.T) {
+	setupDummyKey(t)
 
 	ts, err := JWTAccessTokenSourceFromJSON(jsonKey, "audience")
 	if err != nil {
@@ -91,24 +83,9 @@ func TestJWTAccessTokenSourceFromJSON(t *testing.T) {
 }
 
 func TestJWTAccessTokenSourceWithScope(t *testing.T) {
-	// Generate a key we can use in the test data.
-	privateKey, err := rsa.GenerateKey(rand.Reader, 2048)
-	if err != nil {
-		t.Fatal(err)
-	}
+	setupDummyKey(t)
 
-	// Encode the key and substitute into our example JSON.
-	enc := pem.EncodeToMemory(&pem.Block{
-		Type:  "PRIVATE KEY",
-		Bytes: x509.MarshalPKCS1PrivateKey(privateKey),
-	})
-	enc, err = json.Marshal(string(enc))
-	if err != nil {
-		t.Fatalf("json.Marshal: %v", err)
-	}
-	jsonKey := bytes.Replace(jwtJSONKey, []byte(`"super secret key"`), enc, 1)
-
-	ts, err := JWTAccessTokenSourceWithScope(jsonKey, []string{"scope1", "scope2"})
+	ts, err := JWTAccessTokenSourceWithScope(jsonKey, "scope1", "scope2")
 	if err != nil {
 		t.Fatalf("JWTAccessTokenSourceWithScope: %v\nJSON: %s", err, string(jsonKey))
 	}
@@ -160,3 +137,24 @@ func TestJWTAccessTokenSourceWithScope(t *testing.T) {
 		t.Errorf("Header KeyID = %q, want %q", got, want)
 	}
 }
+
+func setupDummyKey(t *testing.T) {
+	once.Do(func () {
+		// Generate a key we can use in the test data.
+		pk, err := rsa.GenerateKey(rand.Reader, 2048)
+		if err != nil {
+			t.Fatal(err)
+		}
+		privateKey = pk
+		// Encode the key and substitute into our example JSON.
+		enc := pem.EncodeToMemory(&pem.Block{
+			Type:  "PRIVATE KEY",
+			Bytes: x509.MarshalPKCS1PrivateKey(privateKey),
+		})
+		enc, err = json.Marshal(string(enc))
+		if err != nil {
+			t.Fatalf("json.Marshal: %v", err)
+		}
+		jsonKey = bytes.Replace(jwtJSONKey, []byte(`"super secret key"`), enc, 1)
+  })
+}