Dynamic jwks based on iss and aud

go.mod

@@ -26,6 +26,7 @@ require (
 	github.com/rs/zerolog v1.29.0
 	github.com/spf13/cobra v0.0.7
 	github.com/stretchr/testify v1.8.2
+	github.com/valyala/fasttemplate v1.2.2
 	github.com/vmihailenco/msgpack/v5 v5.3.5
 	go.uber.org/automaxprocs v1.5.2
 	golang.org/x/crypto v0.7.0

go.sum

@@ -241,6 +241,8 @@ github.com/tmc/grpc-websocket-proxy v0.0.0-20190109142713-0ad062ec5ee5/go.mod h1
 github.com/ugorji/go v1.1.4/go.mod h1:uQMGLiO92mf5W77hV/PUCpI3pbzQx3CRekS0kk+RGrc=
 github.com/valyala/bytebufferpool v1.0.0 h1:GqA5TC/0021Y/b9FG4Oi9Mr3q7XYx6KllzawFIhcdPw=
 github.com/valyala/bytebufferpool v1.0.0/go.mod h1:6bBcMArwyJ5K/AmCkWv1jt77kVWyCJ6HpOuEn7z0Csc=
+github.com/valyala/fasttemplate v1.2.2 h1:lxLXG0uE3Qnshl9QyaK6XJxMXlQZELvChBOCmQD0Loo=
+github.com/valyala/fasttemplate v1.2.2/go.mod h1:KHLXt3tVN2HBp8eijSv/kGJopbvo7S+qRAEEKiv+SiQ=
 github.com/vmihailenco/msgpack/v5 v5.3.5 h1:5gO0H1iULLWGhs2H5tbAHIZTV8/cYafcFOr9znI5mJU=
 github.com/vmihailenco/msgpack/v5 v5.3.5/go.mod h1:7xyJ9e+0+9SaZT0Wt1RGleJXzli6Q/V5KbhBonMG9jc=
 github.com/vmihailenco/tagparser/v2 v2.0.0 h1:y09buUbR+b5aycVFQs/g70pqKVZNBmxwAhO7/IwNM9g=

internal/cli/token.go

@@ -63,7 +63,10 @@ func verify(config jwtverify.VerifierConfig, ruleConfig rule.Config, token strin
 	if err != nil {
 		return jwtverify.ConnectToken{}, err
 	}
-	verifier := jwtverify.NewTokenVerifierJWT(config, ruleContainer)
+	verifier, err := jwtverify.NewTokenVerifierJWT(config, ruleContainer)
+	if err != nil {
+		return jwtverify.ConnectToken{}, err
+	}
 	return verifier.VerifyConnectToken(token)
 }
 
@@ -72,7 +75,10 @@ func verifySub(config jwtverify.VerifierConfig, ruleConfig rule.Config, token st
 	if err != nil {
 		return jwtverify.SubscribeToken{}, err
 	}
-	verifier := jwtverify.NewTokenVerifierJWT(config, ruleContainer)
+	verifier, err := jwtverify.NewTokenVerifierJWT(config, ruleContainer)
+	if err != nil {
+		return jwtverify.SubscribeToken{}, err
+	}
 	return verifier.VerifySubscribeToken(token)
 }
 

internal/jwks/manager.go

@@ -11,6 +11,7 @@ import (
 	"time"
 
 	"github.com/rakutentech/jwk-go/jwk"
+	"github.com/valyala/fasttemplate"
 	"golang.org/x/sync/singleflight"
 )
 
@@ -41,7 +42,7 @@ var (
 
 // Manager fetches and returns JWK from public source.
 type Manager struct {
-	url      string
+	url      *fasttemplate.Template
 	cache    Cache
 	client   *http.Client
 	useCache bool
@@ -60,13 +61,18 @@ func defaultHTTPClient() *http.Client {
 
 // NewManager returns a new instance of Manager.
 func NewManager(rawURL string, opts ...Option) (*Manager, error) {
-	_, err := url.Parse(rawURL)
+	u, err := url.Parse(rawURL)
 	if err != nil {
 		return nil, ErrInvalidURL
 	}
+	if u.Scheme != "http" && u.Scheme != "https" {
+		return nil, fmt.Errorf("endpoint must have http:// or https:// scheme, got: %s", rawURL)
+	}
+	urlTemplate := fasttemplate.New(rawURL, "{{", "}}")
 
 	mng := &Manager{
-		url:      rawURL,
+		url: urlTemplate,
+
 		cache:    NewTTLCache(_defaultTTL),
 		client:   defaultHTTPClient(),
 		useCache: true,
@@ -85,7 +91,7 @@ func NewManager(rawURL string, opts ...Option) (*Manager, error) {
 }
 
 // FetchKey fetches JWKS from public source or cache.
-func (m *Manager) FetchKey(ctx context.Context, kid string) (*JWK, error) {
+func (m *Manager) FetchKey(ctx context.Context, kid string, tokenVars map[string]any) (*JWK, error) {
 	if kid == "" {
 		return nil, ErrKeyIDNotProvided
 	}
@@ -100,7 +106,7 @@ func (m *Manager) FetchKey(ctx context.Context, kid string) (*JWK, error) {
 
 	// Otherwise fetch from public JWKS.
 	v, err, _ := m.group.Do(kid, func() (interface{}, error) {
-		return m.fetchKey(ctx, kid)
+		return m.fetchKey(ctx, kid, tokenVars)
 	})
 	if err != nil {
 		return nil, err
@@ -122,8 +128,9 @@ func (m *Manager) loadData(req *http.Request) ([]byte, error) {
 	return io.ReadAll(resp.Body)
 }
 
-func (m *Manager) fetchKey(ctx context.Context, kid string) (*JWK, error) {
-	req, err := http.NewRequestWithContext(ctx, http.MethodGet, m.url, nil)
+func (m *Manager) fetchKey(ctx context.Context, kid string, tokenVars map[string]any) (*JWK, error) {
+	jwkURL := m.url.ExecuteString(tokenVars)
+	req, err := http.NewRequestWithContext(ctx, http.MethodGet, jwkURL, nil)
 	if err != nil {
 		return nil, err
 	}

internal/jwks/manager_test.go

@@ -61,7 +61,7 @@ func TestManagerFetchKey_UnmarshalError(t *testing.T) {
 	manager, err := NewManager(server.URL + path)
 	require.NoError(t, err)
 
-	_, err = manager.FetchKey(context.Background(), "202101")
+	_, err = manager.FetchKey(context.Background(), "202101", nil)
 	require.ErrorIs(t, err, errUnmarshal)
 }
 
@@ -78,7 +78,7 @@ func TestManagerFetchKey_KeyNotFound(t *testing.T) {
 	manager, err := NewManager(server.URL + path)
 	require.NoError(t, err)
 
-	_, err = manager.FetchKey(context.Background(), "202101")
+	_, err = manager.FetchKey(context.Background(), "202101", nil)
 	require.ErrorIs(t, err, ErrPublicKeyNotFound)
 }
 
@@ -94,7 +94,7 @@ func TestManagerFetchKey_WrongStatusCode(t *testing.T) {
 	manager, err := NewManager(server.URL + path)
 	require.NoError(t, err)
 
-	_, err = manager.FetchKey(context.Background(), "202101")
+	_, err = manager.FetchKey(context.Background(), "202101", nil)
 	require.ErrorIs(t, err, errUnexpectedStatusCode)
 }
 
@@ -131,7 +131,7 @@ func TestManagerInitialFetchKey(t *testing.T) {
 			manager, err := NewManager(ts.URL)
 			r.NoError(err)
 
-			key, err := manager.FetchKey(context.Background(), tc.Kid)
+			key, err := manager.FetchKey(context.Background(), tc.Kid, nil)
 			if tc.Error != nil {
 				r.Error(err)
 				r.ErrorIs(err, tc.Error)
@@ -176,7 +176,7 @@ func TestManagerCachedFetchKey(t *testing.T) {
 			manager, err := NewManager(ts.URL, tc.Options...)
 			r.NoError(err)
 
-			key, err := manager.FetchKey(ctx, kid)
+			key, err := manager.FetchKey(ctx, kid, nil)
 			r.NoError(err)
 			r.Equal(kid, key.Kid)