adjust

Author: hf

internal/conf/configuration.go

@@ -2,6 +2,7 @@ package conf
 
 import (
 	"bytes"
+	"context"
 	"encoding/base64"
 	"encoding/json"
 	"errors"
@@ -141,6 +142,8 @@ type JWTConfiguration struct {
 	KeyID            string         `json:"key_id" split_words:"true"`
 	Keys             JwtKeysDecoder `json:"keys"`
 	ValidMethods     []string       `json:"-" split_words:"true"`
+
+	SigningKey func(context.Context) (any, error) `json:"-"`
 }
 
 type MFAFactorTypeConfiguration struct {
@@ -1012,13 +1015,22 @@ func (config *GlobalConfiguration) ApplyDefaults() error {
 		if err := config.applyDefaultsJWT([]byte(config.JWT.Secret)); err != nil {
 			return err
 		}
+	} else {
+		jwk, err := GetSigningJwk(&config.JWT)
+		if err != nil {
+			return err
+		}
+		sk, err := getSigningKey(jwk)
+		if err != nil {
+			return err
+		}
+		config.JWT.SigningKey = sk
 	}
 
 	if config.JWT.ValidMethods == nil {
 		config.JWT.ValidMethods = []string{}
 		for _, key := range config.JWT.Keys {
-			alg := GetSigningAlg(key.PublicKey)
-			config.JWT.ValidMethods = append(config.JWT.ValidMethods, alg.Alg())
+			config.JWT.ValidMethods = append(config.JWT.ValidMethods, key.PublicKey.Algorithm().String())
 		}
 
 	}

internal/conf/jwk.go

@@ -6,7 +6,6 @@ import (
 	"fmt"
 	"slices"
 	"strings"
-	"sync"
 
 	"github.com/aws/aws-sdk-go-v2/config"
 	"github.com/aws/aws-sdk-go-v2/service/kms"
@@ -135,37 +134,26 @@ func GetSigningJwk(config *JWTConfiguration) (jwk.Key, error) {
 	return nil, fmt.Errorf("no signing key found")
 }
 
-var kmsClients sync.Map // map[string]func() (*kms.Client, error)
-
-func getKMSClient(region string) (*kms.Client, error) {
-	fn, _ := kmsClients.LoadOrStore(region, sync.OnceValues(func() (*kms.Client, error) {
-		cfg, err := config.LoadDefaultConfig(
-			context.Background(),
-			config.WithRegion(region),
-		)
-		if err != nil {
-			return nil, err
-		}
-
-		return kms.NewFromConfig(cfg), nil
-	}))
-
-	return fn.(func() (*kms.Client, error))()
-}
-
-func GetSigningKey(ctx context.Context, k jwk.Key) (any, error) {
+func getSigningKey(k jwk.Key) (func(context.Context) (any, error), error) {
 	if value, isKMS := k.Get("aws:kms:arn"); isKMS && k.KeyOps()[0] == jwk.KeyOpSign {
 		kmsARN, ok := value.(string)
 		if !ok {
 			return nil, fmt.Errorf("conf: jwk key has aws:kms:arn but is not a string %v", value)
 		}
 
 		parts := strings.SplitN(kmsARN, ":", 5)
-		kmsClient, err := getKMSClient(parts[3])
+		region := parts[3]
+
+		cfg, err := config.LoadDefaultConfig(
+			context.Background(),
+			config.WithRegion(region),
+		)
 		if err != nil {
 			return nil, err
 		}
 
+		kmsClient := kms.NewFromConfig(cfg)
+
 		pub, err := k.PublicKey()
 		if err != nil {
 			return nil, err
@@ -188,21 +176,26 @@ func GetSigningKey(ctx context.Context, k jwk.Key) (any, error) {
 			return nil, err
 		}
 
-		return &awskmsjwk.RS256Key{
-			Ctx: ctx,
-			KMS: kmsClient,
+		return func(ctx context.Context) (any, error) {
+			return &awskmsjwk.RS256Key{
+				Ctx: ctx,
+				KMS: kmsClient,
 
-			KeyID: kmsARN,
-			Raw:   raw,
+				KeyID: kmsARN,
+				Raw:   raw,
+			}, nil
 		}, nil
+
 	}
 
 	var key any
 	if err := k.Raw(&key); err != nil {
 		return nil, err
 	}
 
-	return key, nil
+	return func(ctx context.Context) (any, error) {
+		return key, nil
+	}, nil
 }
 
 func GetSigningAlg(k jwk.Key) jwt.SigningMethod {
@@ -233,8 +226,8 @@ func GetSigningAlg(k jwk.Key) jwt.SigningMethod {
 
 func FindPublicKeyByKid(ctx context.Context, kid string, config *JWTConfiguration) (any, error) {
 	if k, ok := config.Keys[kid]; ok {
-		key, err := GetSigningKey(ctx, k.PublicKey)
-		if err != nil {
+		var key any
+		if err := k.PublicKey.Raw(&key); err != nil {
 			return nil, err
 		}
 		return key, nil

internal/tokens/service.go

@@ -972,7 +972,7 @@ func SignJWT(ctx context.Context, config *conf.JWTConfiguration, claims jwt.Clai
 	}
 	// this serializes the aud claim to a string
 	jwt.MarshalSingleStringAsArray = false
-	signingKey, err := conf.GetSigningKey(ctx, signingJwk)
+	signingKey, err := config.SigningKey(ctx)
 	if err != nil {
 		return "", err
 	}