From 0e37cd0c635192dab4eee62dea9f593111d438c5 Mon Sep 17 00:00:00 2001
From: Giannis Katsanos <giannis@clerk.dev>
Date: Wed, 14 Feb 2024 14:42:59 +0200
Subject: [PATCH] feat: Handler to check azp claim

The JWT verification function accepts a handler that can be used to
validate the 'azp' claim of the token.
The WithHeaderAuthorization middleware accepts either a handler or a
list of authorized parties that must contain the 'azp' claim.
---
 http/middleware.go      | 27 +++++++++++++++++++++++----
 http/middleware_test.go | 33 +++++++++++++++++++++++++++++++++
 jwt/jwt.go              | 24 +++++++-----------------
 3 files changed, 63 insertions(+), 21 deletions(-)

diff --git a/http/middleware.go b/http/middleware.go
index 7d0eacb7..8936415f 100644
--- a/http/middleware.go
+++ b/http/middleware.go
@@ -160,15 +160,34 @@ type AuthorizationParams struct {
 // authorization options.
 type AuthorizationOption func(*AuthorizationParams) error
 
-// AuthorizedParty sets the authorized parties that will be checked
-// against the azp JWT claim.
-func AuthorizedParty(parties ...string) AuthorizationOption {
+// AuthorizedParty allows to provide a handler that accepts the
+// 'azp' claim.
+// The handler can be used to perform validations on the azp claim
+// and should return false to indicate that something is wrong.
+func AuthorizedParty(handler func(string) bool) AuthorizationOption {
 	return func(params *AuthorizationParams) error {
-		params.SetAuthorizedParties(parties...)
+		params.AuthorizedPartyHandler = handler
 		return nil
 	}
 }
 
+// AuthorizedPartyMatches registers a handler that checks that the
+// 'azp' claim's value is included in the provided parties.
+func AuthorizedPartyMatches(parties ...string) func(string) bool {
+	authorizedParties := make(map[string]struct{})
+	for _, p := range parties {
+		authorizedParties[p] = struct{}{}
+	}
+
+	return func(azp string) bool {
+		if azp == "" || len(authorizedParties) == 0 {
+			return true
+		}
+		_, ok := authorizedParties[azp]
+		return ok
+	}
+}
+
 // CustomClaims allows to pass a type (e.g. struct), which will be populated with the token claims based on json tags.
 // You must pass a pointer for this option to work.
 func CustomClaims(claims any) AuthorizationOption {
diff --git a/http/middleware_test.go b/http/middleware_test.go
index 76047a3c..b8456c5b 100644
--- a/http/middleware_test.go
+++ b/http/middleware_test.go
@@ -62,3 +62,36 @@ func TestRequireHeaderAuthorization_InvalidAuthorization(t *testing.T) {
 	require.NoError(t, err)
 	require.Equal(t, http.StatusForbidden, res.StatusCode)
 }
+
+func TestAuthorizedPartyFunc(t *testing.T) {
+	t.Parallel()
+	for _, tc := range []struct {
+		azp     string
+		parties []string
+		want    bool
+	}{
+		{
+			azp:     "clerk.com",
+			parties: []string{"clerk.com", "clerk.dev"},
+			want:    true,
+		},
+		{
+			azp:     "clerk.com",
+			parties: []string{"clerk.dev"},
+			want:    false,
+		},
+		{
+			azp:     "",
+			parties: []string{"clerk.com"},
+			want:    true,
+		},
+		{
+			azp:     "clerk.com",
+			parties: []string{},
+			want:    true,
+		},
+	} {
+		fn := AuthorizedPartyMatches(tc.parties...)
+		require.Equal(t, tc.want, fn(tc.azp))
+	}
+}
diff --git a/jwt/jwt.go b/jwt/jwt.go
index f00fd439..b42aff83 100644
--- a/jwt/jwt.go
+++ b/jwt/jwt.go
@@ -12,6 +12,8 @@ import (
 	"github.com/go-jose/go-jose/v3/jwt"
 )
 
+type AuthorizedPartyHandler func(string) bool
+
 type VerifyParams struct {
 	// Token is the JWT that will be verified. Required.
 	Token string
@@ -26,19 +28,9 @@ type VerifyParams struct {
 	IsSatellite bool
 	// ProxyURL is the URL of the server that proxies the Clerk Frontend API.
 	ProxyURL *string
-	// List of values that should match the azp claim.
-	// Use SetAuthorizedParties to set the value.
-	authorizedParties map[string]struct{}
-}
-
-// SetAuthorizedParties accepts a list of authorized parties to be
-// set on the params.
-func (params *VerifyParams) SetAuthorizedParties(parties ...string) {
-	azp := make(map[string]struct{})
-	for _, p := range parties {
-		azp[p] = struct{}{}
-	}
-	params.authorizedParties = azp
+	// AuthorizedPartyHandler can be used to perform validations on the
+	// 'azp' claim.
+	AuthorizedPartyHandler AuthorizedPartyHandler
 }
 
 // Verify verifies a Clerk session JWT and returns the parsed
@@ -81,10 +73,8 @@ func Verify(ctx context.Context, params *VerifyParams) (*clerk.SessionClaims, er
 		return nil, fmt.Errorf("invalid issuer %s", iss)
 	}
 
-	if claims.AuthorizedParty != "" && len(params.authorizedParties) > 0 {
-		if _, ok := params.authorizedParties[claims.AuthorizedParty]; !ok {
-			return nil, fmt.Errorf("invalid authorized party %s", claims.AuthorizedParty)
-		}
+	if params.AuthorizedPartyHandler != nil && !params.AuthorizedPartyHandler(claims.AuthorizedParty) {
+		return nil, fmt.Errorf("invalid authorized party %s", claims.AuthorizedParty)
 	}
 
 	return claims, nil