Add first cut of KMS implementation and tweak IAM implementation a bit

Author: someone1

config.go

@@ -26,12 +26,10 @@ var (
 )
 
 type iamConfigKey struct{}
+type kmsConfigKey struct{}
 
-// gcpConfig common config elements between signing methods - not to be used on its own
-type gcpConfig struct {
-	// ProjectID is the project id that contains the service account you want to sign with. Defaults to "-" to infer the project from the account
-	ProjectID string
-
+// GCPConfig common config elements between signing methods - not to be used on its own
+type GCPConfig struct {
 	// OAuth2HTTPClient is a user provided oauth2 authenticated *http.Client to use, google.DefaultClient used otherwise
 	// Used for signing requests
 	OAuth2HTTPClient *http.Client
@@ -47,20 +45,31 @@ type gcpConfig struct {
 
 	// 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.
+	// 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
 type IAMConfig struct {
+	// ProjectID is the project id that contains the service account you want to sign with. Defaults to "-" to infer the project from the account
+	ProjectID string
+
 	// Service account can be the email address or the uniqueId of the service account used to sign the JWT with
 	ServiceAccount string
 
 	// 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
 
-	gcpConfig
+	GCPConfig
+}
+
+// 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 string
+
+	GCPConfig
 }
 
 // NewIAMContext returns a new context.Context that carries a provided IAMConfig value
@@ -74,6 +83,17 @@ func IAMFromContext(ctx context.Context) (*IAMConfig, bool) {
 	return val, ok
 }
 
+// NewKMSContext returns a new context.Context that carries a provided KMSConfig value
+func NewKMSContext(parent context.Context, val *KMSConfig) context.Context {
+	return context.WithValue(parent, kmsConfigKey{}, val)
+}
+
+// KMSFromContext extracts a KMSConfig from a context.Context
+func KMSFromContext(ctx context.Context) (*KMSConfig, bool) {
+	val, ok := ctx.Value(kmsConfigKey{}).(*KMSConfig)
+	return val, ok
+}
+
 func getDefaultOauthClient(ctx context.Context) (*http.Client, error) {
 	return google.DefaultClient(ctx, iam.CloudPlatformScope)
 }

iam.go

@@ -11,42 +11,32 @@ import (
 
 type signingMethodIAM struct {
 	alg  string
-	sign func(iamService *iam.Service, config *IAMConfig, signingString string) (string, error)
+	sign func(ctx context.Context, iamService *iam.Service, config *IAMConfig, signingString string) (string, error)
 }
 
-func getConfigFromKey(key interface{}) (context.Context, *IAMConfig, error) {
+func (s *signingMethodIAM) Alg() string {
+	return s.alg
+}
+
+// Sign implements the Sign method from jwt.SigningMethod
+// For this signing method, a valid context.Context must be
+// passed as the key containing a IAMConfig value
+// NOTE: The HEADER IS IGNORED for the signJWT API as the API will add its own
+func (s *signingMethodIAM) Sign(signingString string, key interface{}) (string, error) {
 	var ctx context.Context
 
 	// check to make sure the key is a context.Context
 	switch k := key.(type) {
 	case context.Context:
 		ctx = k
 	default:
-		return nil, nil, jwt.ErrInvalidKey
+		return "", jwt.ErrInvalidKey
 	}
 
 	// Get the IAMConfig from the context
 	config, ok := IAMFromContext(ctx)
 	if !ok {
-		return ctx, nil, ErrMissingConfig
-	}
-
-	return ctx, config, nil
-}
-
-func (s *signingMethodIAM) Alg() string {
-	return s.alg
-}
-
-// Sign implements the Sign method from jwt.SigningMethod
-// For this signing method, a valid context.Context must be
-// passed as the key containing a IAMConfig value
-// NOTE: The HEADER IS IGNORED for the signJWT API as the API will add its own
-func (s *signingMethodIAM) Sign(signingString string, key interface{}) (string, error) {
-	// Process the key
-	ctx, config, err := getConfigFromKey(key)
-	if err != nil {
-		return "", err
+		return "", ErrMissingConfig
 	}
 
 	// Default config.OAuth2HTTPClient is a google.DefaultClient
@@ -70,12 +60,12 @@ func (s *signingMethodIAM) Sign(signingString string, key interface{}) (string,
 		return "", err
 	}
 
-	return s.sign(iamService, config, signingString)
+	return s.sign(ctx, iamService, config, signingString)
 }
 
-// VerfiyKeyfunc is a helper meant that returns a jwt.Keyfunc. It will handle pulling and selecting the certificates
+// IAMVerfiyKeyfunc is a helper meant that returns a jwt.Keyfunc. It will handle pulling and selecting the certificates
 // to verify signatures with, caching when enabled.
-func VerfiyKeyfunc(ctx context.Context, config *IAMConfig) jwt.Keyfunc {
+func IAMVerfiyKeyfunc(ctx context.Context, config *IAMConfig) jwt.Keyfunc {
 	return func(token *jwt.Token) (interface{}, error) {
 		// Make sure we have the proper header alg
 		if _, ok := token.Method.(*signingMethodIAM); !ok {

iam_blob.go

@@ -1,6 +1,7 @@
 package gcpjwt
 
 import (
+	"context"
 	"encoding/base64"
 	"fmt"
 	"net/http"
@@ -34,15 +35,15 @@ func OverrideRS256WithIAMBlob() {
 	})
 }
 
-func signBlob(iamService *iam.Service, config *IAMConfig, signingString string) (string, error) {
+func signBlob(ctx context.Context, iamService *iam.Service, config *IAMConfig, signingString string) (string, error) {
 	// Prepare the call
 	signReq := &iam.SignBlobRequest{
 		BytesToSign: base64.StdEncoding.EncodeToString([]byte(signingString)),
 	}
 	name := fmt.Sprintf("projects/%s/serviceAccounts/%s", config.ProjectID, config.ServiceAccount)
 
 	// Do the call
-	signResp, err := iamService.Projects.ServiceAccounts.SignBlob(name, signReq).Do()
+	signResp, err := iamService.Projects.ServiceAccounts.SignBlob(name, signReq).Context(ctx).Do()
 	if err != nil {
 		return "", err
 	}

iam_jwt.go

@@ -1,6 +1,7 @@
 package gcpjwt
 
 import (
+	"context"
 	"fmt"
 	"net/http"
 	"strings"
@@ -34,7 +35,7 @@ func OverrideRS256WithIAMJWT() {
 	})
 }
 
-func signJwt(iamService *iam.Service, config *IAMConfig, signingString string) (string, error) {
+func signJwt(ctx context.Context, iamService *iam.Service, config *IAMConfig, signingString string) (string, error) {
 	// Prepare the call
 	// First decode the JSON string and discard the header
 	parts := strings.Split(signingString, ".")
@@ -50,7 +51,7 @@ func signJwt(iamService *iam.Service, config *IAMConfig, signingString string) (
 	name := fmt.Sprintf("projects/%s/serviceAccounts/%s", config.ProjectID, config.ServiceAccount)
 
 	// Do the call
-	signResp, err := iamService.Projects.ServiceAccounts.SignJwt(name, signReq).Do()
+	signResp, err := iamService.Projects.ServiceAccounts.SignJwt(name, signReq).Context(ctx).Do()
 	if err != nil {
 		return "", err
 	}

iam_test.go

@@ -100,7 +100,7 @@ func TestIAMSignAndVerify(t *testing.T) {
 			}
 
 			token.Method = method
-			keyFunc := VerfiyKeyfunc(c, config)
+			keyFunc := IAMVerfiyKeyfunc(c, config)
 			key, err := keyFunc(token)
 			if err != nil {
 				t.Fatalf("[%v] Error getting key: %v", data.name, err)

jwtmiddleware/appengine.go

@@ -25,7 +25,7 @@ func NewHandler(_ context.Context, config *gcpjwt.IAMConfig, audience string) fu
 		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 			ctx := gcpjwt.NewIAMContext(appengine.NewContext(r), config)
 
-			keyFunc := gcpjwt.VerfiyKeyfunc(ctx, config)
+			keyFunc := gcpjwt.IAMVerfiyKeyfunc(ctx, config)
 			claims := &jwt.StandardClaims{}
 
 			token, err := request.ParseFromRequest(r, request.AuthorizationHeaderExtractor, keyFunc, request.WithClaims(claims))

jwtmiddleware/middleware.go

@@ -23,7 +23,7 @@ import (
 func NewHandler(ctx context.Context, config *gcpjwt.IAMConfig, audience string) func(http.Handler) http.Handler {
 	ctx = gcpjwt.NewIAMContext(ctx, config)
 
-	keyFunc := gcpjwt.VerfiyKeyfunc(ctx, config)
+	keyFunc := gcpjwt.IAMVerfiyKeyfunc(ctx, config)
 
 	return func(h http.Handler) http.Handler {
 		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {