Return all keys when JWT has no key ID header (#140)

Author: MicahParks

keyfunc.go

@@ -25,6 +25,7 @@ type Keyfunc interface {
 	Keyfunc(token *jwt.Token) (any, error)
 	KeyfuncCtx(ctx context.Context) jwt.Keyfunc
 	Storage() jwkset.Storage
+	VerificationKeySet(ctx context.Context) (jwt.VerificationKeySet, error)
 }
 
 // Options are used to create a new Keyfunc.
@@ -198,7 +199,7 @@ func (k keyfunc) KeyfuncCtx(ctx context.Context) jwt.Keyfunc {
 	return func(token *jwt.Token) (any, error) {
 		kidInter, ok := token.Header[jwkset.HeaderKID]
 		if !ok {
-			return nil, fmt.Errorf("%w: could not find kid in JWT header", ErrKeyfunc)
+			return k.VerificationKeySet(ctx)
 		}
 		kid, ok := kidInter.(string)
 		if !ok {
@@ -236,10 +237,6 @@ func (k keyfunc) KeyfuncCtx(ctx context.Context) jwt.Keyfunc {
 			}
 		}
 
-		type publicKeyer interface {
-			Public() crypto.PublicKey
-		}
-
 		key := jwk.Key()
 		pk, ok := key.(publicKeyer)
 		if ok {
@@ -256,3 +253,23 @@ func (k keyfunc) Keyfunc(token *jwt.Token) (any, error) {
 func (k keyfunc) Storage() jwkset.Storage {
 	return k.storage
 }
+func (k keyfunc) VerificationKeySet(ctx context.Context) (jwt.VerificationKeySet, error) {
+	jwk, err := k.storage.KeyReadAll(ctx)
+	if err != nil {
+		return jwt.VerificationKeySet{}, fmt.Errorf("failed to read all JWK from storage: %w", errors.Join(err, ErrKeyfunc))
+	}
+	var allKeys jwt.VerificationKeySet
+	for _, j := range jwk {
+		key := j.Key()
+		pk, ok := key.(publicKeyer)
+		if ok {
+			key = pk.Public()
+		}
+		allKeys.Keys = append(allKeys.Keys, key)
+	}
+	return allKeys, nil
+}
+
+type publicKeyer interface {
+	Public() crypto.PublicKey
+}

keyfunc_test.go

@@ -175,3 +175,126 @@ func TestNewJWKSetJSON(t *testing.T) {
 		t.Fatalf("The token is not valid.")
 	}
 }
+
+func TestVerificationKeySet(t *testing.T) {
+	ctx := context.Background()
+	_, priv, err := ed25519.GenerateKey(rand.Reader)
+	if err != nil {
+		t.Fatalf("Failed to generate ED25519 key pair: %v", err)
+	}
+	jwk, err := jwkset.NewJWKFromKey(priv, jwkset.JWKOptions{})
+	if err != nil {
+		t.Fatalf("Failed to create JWK: %v", err)
+	}
+	store := jwkset.NewMemoryStorage()
+	err = store.KeyWrite(ctx, jwk)
+	if err != nil {
+		t.Fatalf("Failed to write JWK: %v", err)
+	}
+	k, err := New(Options{Ctx: ctx, Storage: store})
+	if err != nil {
+		t.Fatalf("Failed to create Keyfunc: %v", err)
+	}
+	vks, err := k.VerificationKeySet(ctx)
+	if err != nil {
+		t.Fatalf("VerificationKeySet failed: %v", err)
+	}
+	if len(vks.Keys) != 1 {
+		t.Fatalf("Expected 1 key, got %d", len(vks.Keys))
+	}
+}
+
+func TestNoKIDHeaderCallsVerificationKeySet(t *testing.T) {
+	ctx := context.Background()
+
+	// Generate two key pairs.
+	_, priv1, err := ed25519.GenerateKey(rand.Reader)
+	if err != nil {
+		t.Fatalf("Failed to generate ED25519 key pair 1: %v", err)
+	}
+	_, priv2, err := ed25519.GenerateKey(rand.Reader)
+	if err != nil {
+		t.Fatalf("Failed to generate ED25519 key pair 2: %v", err)
+	}
+
+	jwk1, err := jwkset.NewJWKFromKey(priv1, jwkset.JWKOptions{})
+	if err != nil {
+		t.Fatalf("Failed to create JWK 1: %v", err)
+	}
+	jwk2, err := jwkset.NewJWKFromKey(priv2, jwkset.JWKOptions{})
+	if err != nil {
+		t.Fatalf("Failed to create JWK 2: %v", err)
+	}
+
+	orders := [][]jwkset.JWK{
+		{jwk1, jwk2},
+		{jwk2, jwk1},
+	}
+	privs := []ed25519.PrivateKey{priv1, priv2}
+
+	for i, order := range orders {
+		store := jwkset.NewMemoryStorage()
+		for _, jwk := range order {
+			err = store.KeyWrite(ctx, jwk)
+			if err != nil {
+				t.Fatalf("Failed to write JWK: %v", err)
+			}
+		}
+		k, err := New(Options{Ctx: ctx, Storage: store})
+		if err != nil {
+			t.Fatalf("Failed to create Keyfunc: %v", err)
+		}
+		// Sign a token with the corresponding private key (no KID header)
+		token := jwt.New(jwt.SigningMethodEdDSA)
+		tokenString, err := token.SignedString(privs[i])
+		if err != nil {
+			t.Fatalf("Failed to sign token: %v", err)
+		}
+		parsedToken, err := jwt.Parse(tokenString, k.KeyfuncCtx(ctx))
+		if err != nil {
+			t.Fatalf("Parse failed (order %d): %v", i+1, err)
+		}
+		if !parsedToken.Valid {
+			t.Fatalf("Expected token to be valid (order %d)", i+1)
+		}
+	}
+}
+
+func TestNoKIDHeaderNoMatchingJWK(t *testing.T) {
+	ctx := context.Background()
+
+	_, missingFromSet, err := ed25519.GenerateKey(rand.Reader)
+	if err != nil {
+		t.Fatalf("Failed to generate ED25519 key pair: %v", err)
+	}
+
+	_, presentInSet, err := ed25519.GenerateKey(rand.Reader)
+	if err != nil {
+		t.Fatalf("Failed to generate other ED25519 key pair: %v", err)
+	}
+	jwk, err := jwkset.NewJWKFromKey(presentInSet, jwkset.JWKOptions{})
+	if err != nil {
+		t.Fatalf("Failed to create JWK: %v", err)
+	}
+	store := jwkset.NewMemoryStorage()
+	err = store.KeyWrite(ctx, jwk)
+	if err != nil {
+		t.Fatalf("Failed to write JWK: %v", err)
+	}
+
+	k, err := New(Options{Ctx: ctx, Storage: store})
+	if err != nil {
+		t.Fatalf("Failed to create Keyfunc: %v", err)
+	}
+
+	token := jwt.New(jwt.SigningMethodEdDSA)
+	tokenString, err := token.SignedString(missingFromSet)
+	if err != nil {
+		t.Fatalf("Failed to sign token: %v", err)
+	}
+
+	_, err = jwt.Parse(tokenString, k.KeyfuncCtx(ctx))
+	if err == nil {
+		t.Fatalf("Expected error due to no matching JWK, but got none")
+	}
+}