From 95dc43f71be52179ec58a6011604e0995f46da05 Mon Sep 17 00:00:00 2001
From: Joakim Antman <antmanj@gmail.com>
Date: Sun, 29 Dec 2024 17:55:47 +0200
Subject: [PATCH] Better support for JWK keyfinder and EncodedToken

---
 CHANGELOG.md                   |  1 +
 README.md                      | 29 +++++++++++++++++++++++++++++
 lib/jwt/jwk/key_finder.rb      | 15 ++++++++++++++-
 spec/jwt/encoded_token_spec.rb | 14 ++++++++++++++
 4 files changed, 58 insertions(+), 1 deletion(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index de4c0b66..469fb8fd 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -19,6 +19,7 @@ Take a look at the [upgrade guide](UPGRADING.md) for more details.
 **Features:**
 - JWT::EncodedToken#verify! method that bundles signature and claim validation [#647](https://github.com/jwt/ruby-jwt/pull/647) ([@anakinj](https://github.com/anakinj))
 - Do not override the alg header if already given [#659](https://github.com/jwt/ruby-jwt/pull/659) ([@anakinj](https://github.com/anakinj))
+- Make `JWK::KeyFinder` compatible with `JWT::EncodedToken` [#663](https://github.com/jwt/ruby-jwt/pull/663) ([@anakinj](https://github.com/anakinj))
 - Your contribution here
 
 **Fixes and enhancements:**
diff --git a/README.md b/README.md
index 90ffbf4d..3a2b1eb0 100644
--- a/README.md
+++ b/README.md
@@ -296,6 +296,35 @@ encoded_token.payload # => { 'exp'=>1234, 'jti'=>'1234", 'sub'=>'my-subject' }
 encoded_token.header # {'kid'=>'hmac', 'alg'=>'HS256'}
 ```
 
+A keyfinder can be used to verify a signature. A keyfinder is an object responding to the `#call` method. The method expects to receive one argument, which is the token to be verified.
+
+An example on using the built-in JWK keyfinder:
+```ruby
+# Create and sign a token
+jwk = JWT::JWK.new(OpenSSL::PKey::RSA.generate(2048))
+token = JWT::Token.new(payload: { pay: 'load' }, header: { kid: jwk.kid })
+token.sign!(algorithm: 'RS256', key: jwk.signing_key)
+
+# Create keyfinder object, verify and decode token
+key_finder = JWT::JWK::KeyFinder.new(jwks: JWT::JWK::Set.new(jwk))
+encoded_token = JWT::EncodedToken.new(token.jwt)
+encoded_token.verify!(signature: { algorithm: 'RS256', key_finder: key_finder})
+encoded_token.payload # => { 'pay' => 'load' }
+```
+
+Using a custom keyfinder proc:
+```ruby
+# Create and sign a token
+key = OpenSSL::PKey::RSA.generate(2048)
+token = JWT::Token.new(payload: { pay: 'load' })
+token.sign!(algorithm: 'RS256', key: key)
+
+# Verify and decode token
+encoded_token = JWT::EncodedToken.new(token.jwt)
+encoded_token.verify!(signature: { algorithm: 'RS256', key_finder: ->(_token){ key.public_key }})
+encoded_token.payload # => { 'pay' => 'load' }
+```
+
 ### Detached payload
 
 The `::JWT::Token#detach_payload!` method can be use to detach the payload from the JWT.
diff --git a/lib/jwt/jwk/key_finder.rb b/lib/jwt/jwk/key_finder.rb
index 5498800c..80a2e7fe 100644
--- a/lib/jwt/jwk/key_finder.rb
+++ b/lib/jwt/jwk/key_finder.rb
@@ -2,8 +2,13 @@
 
 module JWT
   module JWK
-    # @api private
+    # JSON Web Key keyfinder
+    # To find the key for a given kid
     class KeyFinder
+      # Initializes a new KeyFinder instance.
+      # @param [Hash] options the options to create a KeyFinder with
+      # @option options [Proc, JWT::JWK::Set] :jwks the jwks or a loader proc
+      # @option options [Boolean] :allow_nil_kid whether to allow nil kid
       def initialize(options)
         @allow_nil_kid = options[:allow_nil_kid]
         jwks_or_loader = options[:jwks]
@@ -15,6 +20,8 @@ def initialize(options)
                        end
       end
 
+      # Returns the verification key for the given kid
+      # @param [String] kid the key id
       def key_for(kid)
         raise ::JWT::DecodeError, 'No key id (kid) found from token headers' unless kid || @allow_nil_kid
         raise ::JWT::DecodeError, 'Invalid type for kid header parameter' unless kid.nil? || kid.is_a?(String)
@@ -27,6 +34,12 @@ def key_for(kid)
         jwk.verify_key
       end
 
+      # Returns the key for the given token
+      # @param [JWT::EncodedToken] token the token
+      def call(token)
+        key_for(token.header['kid'])
+      end
+
       private
 
       def resolve_key(kid)
diff --git a/spec/jwt/encoded_token_spec.rb b/spec/jwt/encoded_token_spec.rb
index d03f05f0..fd48fa85 100644
--- a/spec/jwt/encoded_token_spec.rb
+++ b/spec/jwt/encoded_token_spec.rb
@@ -161,6 +161,20 @@
         expect(token.verify_signature!(algorithm: 'HS256', key: key)).to eq(nil)
       end
     end
+
+    context 'when JWT::KeyFinder is used as a key_finder' do
+      let(:jwk) { JWT::JWK.new(test_pkey('rsa-2048-private.pem')) }
+      let(:encoded_token) do
+        JWT::Token.new(payload: payload, header: { kid: jwk.kid })
+                  .tap { |t| t.sign!(algorithm: 'RS256', key: jwk.signing_key) }
+                  .jwt
+      end
+
+      it 'uses the keys provided by the JWK key finder' do
+        key_finder = JWT::JWK::KeyFinder.new(jwks: JWT::JWK::Set.new(jwk))
+        expect(token.verify_signature!(algorithm: 'RS256', key_finder: key_finder)).to eq(nil)
+      end
+    end
   end
 
   describe '#verify_claims!' do