all: More refactoring, bug fixes, and added tests

Author: someone1

README.md

@@ -35,9 +35,9 @@ import (
 
 func init() {
     // Unless we want to keep the original RS256 implementation alive, override it (recommended)
-    gcpjwt.OverrideRS256WithIAMJWT() // For signJwt
-
-    gcpjwt.OverrideRS256WithIAMBlob() // For signBlob
+    gcpjwt.SigningMethodIAMJWT.Override() // For signJwt
+    // OR
+    gcpjwt.SigningMethodIAMBlob.Override() // For signBlob
 }
 ```
 
@@ -95,10 +95,10 @@ func validateToken(tokenString string) {
     }
     config.EnableCache = true // Enable certificates cache
 
-    // To Verify (if we called OverrideRS256WithIAMJWT() or OverrideRS256WithIAMBlob())
+    // To Verify (if we called Override() for our method type prior)
     token, err := jwt.Parse(tokenString, gcpjwt.VerfiyKeyfunc(context.Background(), config))
 
-    // If we DID NOT call a OverrideRS256 function
+    // If we DID NOT call the Override() function
     // This is basically copying the https://github.com/dgrijalva/jwt-go/blob/master/parser.go#L23 ParseWithClaims function here but forcing our own method vs getting one based on the Alg field
     // Or Try and parse, Ignore the result and try with the proper method:
     token, _ := jwt.Parse(tokenString, func(token *jwt.Token) (interface{}, error) {

appengine.go

@@ -1,6 +1,7 @@
 package gcpjwt
 
 import (
+	"context"
 	"crypto/rsa"
 	"encoding/json"
 	"io/ioutil"
@@ -15,21 +16,28 @@ const (
 	certificateURL = "https://www.googleapis.com/robot/v1/metadata/x509/"
 )
 
-type certResponse struct {
-	certs   certificates
-	expires time.Time
-}
-
 // certificates is a map of key id -> public keys
 type certificates map[string]*rsa.PublicKey
 
-func getCertificatesForAccount(hc *http.Client, account string) (*certResponse, error) {
-	req, err := http.NewRequest(http.MethodGet, certificateURL+account, nil)
+func getCertificates(ctx context.Context, config *IAMConfig) (certificates, error) {
+	if config.EnableCache {
+		if certsResp, ok := getCertsFromCache(config.ServiceAccount); ok {
+			return certsResp, nil
+		}
+	}
+
+	// Default config.Client is a http.DefaultClient
+	client := config.Client
+	if client == nil {
+		client = getDefaultClient(ctx)
+	}
+
+	req, err := http.NewRequest(http.MethodGet, certificateURL+config.ServiceAccount, nil)
 	if err != nil {
 		return nil, err
 	}
 
-	resp, err := hc.Do(req)
+	resp, err := client.Do(req)
 	if err != nil {
 		return nil, err
 	}
@@ -41,19 +49,17 @@ func getCertificatesForAccount(hc *http.Client, account string) (*certResponse,
 	}
 
 	certsRaw := make(map[string]string)
-
 	err = json.Unmarshal(b, &certsRaw)
 	if err != nil {
 		return nil, err
 	}
-	_, expires, err := cachecontrol.CachableResponse(req, resp, cachecontrol.Options{PrivateCache: true})
 
-	if err != nil {
-		return nil, err
+	_, expires, err := cachecontrol.CachableResponse(req, resp, cachecontrol.Options{PrivateCache: true})
+	if err != nil && config.CacheExpiration > 0 {
+		expires = time.Now().Add(config.CacheExpiration)
 	}
 
 	certs := make(certificates)
-
 	for key, cert := range certsRaw {
 		rsaKey, err := jwt.ParseRSAPublicKeyFromPEM([]byte(cert))
 		if err != nil {
@@ -62,5 +68,9 @@ func getCertificatesForAccount(hc *http.Client, account string) (*certResponse,
 		certs[key] = rsaKey
 	}
 
-	return &certResponse{certs, expires}, err
+	if config.EnableCache && !expires.IsZero() {
+		updateCache(config.ServiceAccount, certs, expires)
+	}
+
+	return certs, nil
 }

certs.go

@@ -2,8 +2,11 @@ package gcpjwt
 
 import (
 	"context"
+	"crypto/md5"
 	"errors"
+	"fmt"
 	"net/http"
+	"time"
 
 	"golang.org/x/oauth2/google"
 	"google.golang.org/api/iam/v1"
@@ -37,16 +40,6 @@ type GCPConfig struct {
 	// Client is a user provided *http.Client to use, http.DefaultClient is used otherwise (AppEngine URL Fetch Supported)
 	// Used for verify requests
 	Client *http.Client
-
-	// EnableCache will enable the in-memory caching of public certificates.
-	// The cache will expire certificates when an expiration is known or provided and will refresh the cache if
-	// it is unable to verify a signature from any of the certificates cached.
-	EnableCache bool
-
-	// InjectKeyID will overwrite the provided header with one that contains the Key ID of the key used to sign the JWT.
-	// Note that the IAM JWT signing method does this on its own and this is only applicable for the IAM Blob and Cloud KMS
-	// signing methods. For CloudKMS, this will be a hash of the KeyPath configured.
-	InjectKeyID bool
 }
 
 // IAMConfig is relevant for both the signBlob and signJWT IAM API use-cases
@@ -57,21 +50,45 @@ type IAMConfig struct {
 	// Service account can be the email address or the uniqueId of the service account used to sign the JWT with
 	ServiceAccount string
 
+	// EnableCache will enable the in-memory caching of public certificates.
+	// The cache will expire certificates when an expiration is known or fallback to the configured CacheExpiration
+	EnableCache bool
+
+	// CacheExpiration is the default time to keep the certificates in cache if no expiration time is provided
+	// Use a value of 0 to disable the expiration time fallback. Max reccomneded value is 24 hours.
+	// https://cloud.google.com/iam/docs/understanding-service-accounts#managing_service_account_keys
+	CacheExpiration time.Duration
+
 	// IAMType is a helper used to help clarify which IAM signing method this config is meant for.
 	// Used for the jwtmiddleware and oauth2 packages.
 	IAMType iamType
 
+	lastKeyID string
+
 	GCPConfig
 }
 
+// KeyID will return the last used KeyID to sign the JWT - though it should be noted the signJwt method will always
+// add its own token header which is not parsed back to the token.
+// Helper function for adding the kid header to your token.
+func (i *IAMConfig) KeyID() string {
+	return i.lastKeyID
+}
+
 // KMSConfig is used to sign/verify JWTs with Google Cloud KMS
 type KMSConfig struct {
-	// KeyPath is the name of the key to use in the format of '/projects/-/locations/...'
+	// KeyPath is the name of the key to use in the format of:
+	// "name=projects/*/locations/*/keyRings/*/cryptoKeys/*/cryptoKeyVersions/*"
 	KeyPath string
 
 	GCPConfig
 }
 
+// KeyID will return the MD5 hash of the configured KeyPath. Helper function for adding the kid header to your token.
+func (k *KMSConfig) KeyID() string {
+	return fmt.Sprintf("%x", md5.Sum([]byte(k.KeyPath)))
+}
+
 // NewIAMContext returns a new context.Context that carries a provided IAMConfig value
 func NewIAMContext(parent context.Context, val *IAMConfig) context.Context {
 	return context.WithValue(parent, iamConfigKey{}, val)