Proposal of simple JWK support

AUTHORS

@@ -70,3 +70,4 @@ Julio Lopez
 Katelyn Kasperowicz
 Lowell Kirsh
 Lucas Mazza
+Joakim Antman

README.md

@@ -439,6 +439,31 @@ rescue JWT::InvalidSubError
 end
 ```
 
+### JSON Web Key (JWK)
+
+JWK is a JSON structure representing a cryptographic key. Currently only supports RSA public keys.
+
+```ruby
+jwk = JWT::JWK.new(OpenSSL::PKey::RSA.new(2048))
+payload, headers = { data: 'data' }, { kid: jwk.kid }
+
+token = JWT.encode(payload, jwk.keypair, 'RS512', headers)
+
+# The jwk loader would fetch the set of JWKs from a trusted source
+jwk_loader = ->(options) do
+  @cached_keys = nil if options[:invalidate] # need to reload the keys
+  @cached_keys ||= { keys: [jwk.export] }
+end
+
+begin
+  JWT.decode(token, nil, true, { algorithms: ['RS512'], jwks: jwk_loader})
+rescue JWT::JWKError
+  # Handle problems with the provided JWKs
+rescue JWT::DecodeError
+  # Handle other decode related issues e.g. no kid in header, no matching public key found etc. 
+end
+```
+
 # Development and Tests
 
 We depend on [Bundler](http://rubygems.org/gems/bundler) for defining gemspec and performing releases to rubygems.org, which can be done with

lib/jwt.rb

@@ -5,6 +5,7 @@
 require 'jwt/default_options'
 require 'jwt/encode'
 require 'jwt/error'
+require 'jwt/jwk'
 
 # JSON Web Token implementation
 #

lib/jwt/decode.rb

@@ -39,6 +39,7 @@ def decode_segments
 
     def verify_signature
       @key = find_key(&@keyfinder) if @keyfinder
+      @key = ::JWT::JWK::KeyFinder.new(jwks: @options[:jwks]).key_for(header['kid']) if @options[:jwks]
 
       raise(JWT::IncorrectAlgorithm, 'An algorithm must be specified') if allowed_algorithms.empty?
       raise(JWT::IncorrectAlgorithm, 'Expected a different algorithm') unless options_includes_algo_in_header?

lib/jwt/error.rb

@@ -15,4 +15,6 @@ module JWT
   InvalidSubError    = Class.new(DecodeError)
   InvalidJtiError    = Class.new(DecodeError)
   InvalidPayload     = Class.new(DecodeError)
+
+  JWKError           = Class.new(DecodeError)
 end

lib/jwt/jwk.rb

@@ -0,0 +1,31 @@
+# frozen_string_literal: true
+
+require_relative 'jwk/rsa'
+require_relative 'jwk/key_finder'
+
+module JWT
+  module JWK
+    MAPPINGS = {
+      'RSA' => ::JWT::JWK::RSA,
+      OpenSSL::PKey::RSA => ::JWT::JWK::RSA
+    }.freeze
+
+    class << self
+      def import(jwk_data)
+        raise JWT::JWKError, 'Key type (kty) not provided' unless jwk_data[:kty]
+
+        MAPPINGS.fetch(jwk_data[:kty].to_s) do |kty|
+          raise JWT::JWKError, "Key type #{kty} not supported"
+        end.import(jwk_data)
+      end
+
+      def create_from(keypair)
+        MAPPINGS.fetch(keypair.class) do |klass|
+          raise JWT::JWKError, "Cannot create JWK from a #{klass.name}"
+        end.new(keypair)
+      end
+
+      alias new create_from
+    end
+  end
+end

lib/jwt/jwk/key_finder.rb

@@ -0,0 +1,57 @@
+# frozen_string_literal: true
+
+module JWT
+  module JWK
+    class KeyFinder
+      def initialize(options)
+        jwks_or_loader = options[:jwks]
+        @jwks          = jwks_or_loader if jwks_or_loader.is_a?(Hash)
+        @jwk_loader    = jwks_or_loader if jwks_or_loader.respond_to?(:call)
+      end
+
+      def key_for(kid)
+        raise ::JWT::DecodeError, 'No key id (kid) found from token headers' unless kid
+
+        jwk = resolve_key(kid)
+
+        raise ::JWT::DecodeError, "Could not find public key for kid #{kid}" unless jwk
+
+        ::JWT::JWK.import(jwk).keypair
+      end
+
+      private
+
+      def resolve_key(kid)
+        jwk = find_key(kid)
+
+        return jwk if jwk
+
+        if reloadable?
+          load_keys(invalidate: true)
+          return find_key(kid)
+        end
+
+        nil
+      end
+
+      def jwks
+        return @jwks if @jwks
+
+        load_keys
+        @jwks
+      end
+
+      def load_keys(opts = {})
+        @jwks = @jwk_loader.call(opts)
+      end
+
+      def find_key(kid)
+        Array(jwks[:keys]).find { |key| key[:kid] == kid }
+      end
+
+      def reloadable?
+        @jwk_loader
+      end
+    end
+  end
+end

lib/jwt/jwk/rsa.rb

@@ -0,0 +1,45 @@
+# frozen_string_literal: true
+
+module JWT
+  module JWK
+    class RSA
+      extend Forwardable
+
+      attr_reader :keypair
+
+      def_delegators :keypair, :private?, :public_key
+
+      BINARY = 2
+      KTY    = 'RSA'.freeze
+
+      def initialize(keypair)
+        raise ArgumentError, 'keypair must be of type OpenSSL::PKey::RSA' unless keypair.is_a?(OpenSSL::PKey::RSA)
+
+        @keypair = keypair
+      end
+
+      def kid
+        sequence = OpenSSL::ASN1::Sequence([OpenSSL::ASN1::Integer.new(public_key.n),
+                                            OpenSSL::ASN1::Integer.new(public_key.e)])
+        OpenSSL::Digest::SHA256.hexdigest(sequence.to_der)
+      end
+
+      def export
+        {
+          kty: KTY,
+          n: Base64.urlsafe_encode64(public_key.n.to_s(BINARY), padding: false),
+          e: Base64.urlsafe_encode64(public_key.e.to_s(BINARY), padding: false),
+          kid: kid
+        }
+      end
+
+      def self.import(jwk_data)
+        imported_key = OpenSSL::PKey::RSA.new
+        imported_key.set_key(OpenSSL::BN.new(Base64.urlsafe_decode64(jwk_data[:n]), BINARY),
+          OpenSSL::BN.new(Base64.urlsafe_decode64(jwk_data[:e]), BINARY),
+          nil)
+        self.new(imported_key)
+      end
+    end
+  end
+end

spec/integration/readme_examples_spec.rb

@@ -217,5 +217,23 @@
         JWT.decode token, hmac_secret, true, 'sub' => sub, :verify_sub => true, :algorithm => 'HS256'
       end.not_to raise_error
     end
+
+
+    it 'JWK' do
+      jwk = JWT::JWK.new(OpenSSL::PKey::RSA.new(2048))
+      payload, headers = { data: 'data' }, { kid: jwk.kid }
+
+      token = JWT.encode(payload, jwk.keypair, 'RS512', headers)
+
+      # The jwk loader would fetch the set of JWKs from a trusted source
+      jwk_loader = ->(options) do
+        @cached_keys = nil if options[:invalidate] # need to reload the keys
+        @cached_keys ||= { keys: [jwk.export] }
+      end
+
+      expect do
+        JWT.decode(token, nil, true, { algorithms: ['RS512'], jwks: jwk_loader})
+      end.not_to raise_error
+    end
   end
 end

spec/jwk/decode_with_jwk_spec.rb

@@ -0,0 +1,68 @@
+# frozen_string_literal: true
+
+require_relative '../spec_helper'
+require 'jwt'
+
+describe JWT do
+  describe '.decode for JWK usecase' do
+    let(:keypair)       { OpenSSL::PKey::RSA.new(2048) }
+    let(:jwk)           { JWT::JWK.new(keypair) }
+    let(:public_jwks) { { keys: [jwk.export, { kid: 'not_the_correct_one' }] } }
+    let(:token_payload) { {'data' => 'something'} }
+    let(:token_headers) { { kid: jwk.kid } }
+    let(:signed_token)  { described_class.encode(token_payload, jwk.keypair, 'RS512', token_headers) }
+
+    context 'when JWK features are used manually' do
+      it 'is able to decode the token' do
+        payload, _header = described_class.decode(signed_token, nil, true, { algorithms: ['RS512'] }) do |header, _payload|
+          JWT::JWK.import(public_jwks[:keys].find { |key| key[:kid] == header['kid'] }).keypair
+        end
+        expect(payload).to eq(token_payload)
+      end
+    end
+
+    context 'when jwk keys are given as an array' do
+      context 'and kid is in the set' do
+        it 'is able to decode the token' do
+          payload, _header = described_class.decode(signed_token, nil, true, { algorithms: ['RS512'], jwks: public_jwks})
+          expect(payload).to eq(token_payload)
+        end
+      end
+
+      context 'and kid is not in the set' do
+        before do
+          public_jwks[:keys].first[:kid] = 'NOT_A_MATCH'
+        end
+        it 'raises an exception' do
+          expect { described_class.decode(signed_token, nil, true, { algorithms: ['RS512'], jwks: public_jwks}) }.to raise_error(
+            JWT::DecodeError, /Could not find public key for kid .*/
+          )
+        end
+      end
+
+      context 'token does not know the kid' do
+        let(:token_headers) { {} }
+        it 'raises an exception' do
+          expect { described_class.decode(signed_token, nil, true, { algorithms: ['RS512'], jwks: public_jwks}) }.to raise_error(
+            JWT::DecodeError, 'No key id (kid) found from token headers'
+          )
+        end
+      end
+    end
+
+    context 'when jwk keys are loaded using a proc/lambda' do
+      it 'decodes the token' do
+        payload, _header = described_class.decode(signed_token, nil, true, { algorithms: ['RS512'], jwks: lambda { |_opts| public_jwks }})
+        expect(payload).to eq(token_payload)
+      end
+    end
+
+    context 'when jwk keys are rotated' do
+      it 'decodes the token' do
+        key_loader = ->(options) { options[:invalidate] ? public_jwks : { keys: [] } }
+        payload, _header = described_class.decode(signed_token, nil, true, { algorithms: ['RS512'], jwks: key_loader})
+        expect(payload).to eq(token_payload)
+      end
+    end
+  end
+end

spec/jwk/rsa_spec.rb

@@ -0,0 +1,57 @@
+# frozen_string_literal: true
+
+require_relative '../spec_helper'
+require 'jwt'
+
+describe JWT::JWK::RSA do
+  let(:rsa_key) { OpenSSL::PKey::RSA.new(2048) }
+
+  describe '.new' do
+    subject { described_class.new(keypair) }
+
+    context 'when a keypair with both keys given' do
+      let(:keypair) { rsa_key }
+      it 'creates an instance of the class' do
+        expect(subject).to be_a described_class
+        expect(subject.private?).to eq true
+      end
+    end
+
+    context 'when a keypair with only public key is given' do
+      let(:keypair) { rsa_key.public_key }
+      it 'creates an instance of the class' do
+        expect(subject).to be_a described_class
+        expect(subject.private?).to eq false
+      end
+    end
+  end
+
+  describe '#export' do
+    subject { described_class.new(keypair).export }
+
+    context 'when keypair with private key is exported' do
+      let(:keypair) { rsa_key }
+      it 'returns a hash with the public parts of the key' do
+        expect(subject).to be_a Hash
+        expect(subject).to include(:kty, :n, :e, :kid)
+        expect(subject).not_to include(:d)
+      end
+    end
+
+    context 'when keypair with public key is exported' do
+      let(:keypair) { rsa_key.public_key }
+      it 'returns a hash with the public parts of the key' do
+        expect(subject).to be_a Hash
+        expect(subject).to include(:kty, :n, :e, :kid)
+        expect(subject).not_to include(:d)
+      end
+    end
+
+    context 'when unsupported keypair is given' do
+      let(:keypair) { 'key' }
+      it 'raises an error' do
+        expect { subject }.to raise_error(ArgumentError, 'keypair must be of type OpenSSL::PKey::RSA')
+      end
+    end
+  end
+end

spec/jwk_spec.rb

@@ -0,0 +1,44 @@
+# frozen_string_literal: true
+
+require 'spec_helper'
+require 'jwt'
+
+describe JWT::JWK do
+  let(:rsa_key) { OpenSSL::PKey::RSA.new(2048) }
+
+  describe '.import' do
+    let(:keypair) { rsa_key.public_key }
+    let(:params)  { described_class.new(keypair).export }
+
+    subject { described_class.import(params) }
+
+    it 'creates a ::JWT::JWK::RSA instance' do
+      expect(subject).to be_a ::JWT::JWK::RSA
+      expect(subject.export).to eq(params)
+    end
+
+    context 'when keytype is not supported' do
+      let(:params) { { kty: 'unsupported' } }
+
+      it 'raises an error' do
+        expect { subject }.to raise_error(JWT::JWKError)
+      end
+    end
+  end
+
+  describe '.to_jwk' do
+    subject { described_class.new(keypair) }
+
+    context 'when RSA key is given' do
+      let(:keypair) { rsa_key }
+      it { is_expected.to be_a ::JWT::JWK::RSA }
+    end
+
+    context 'when unsupported key is given' do
+      let(:keypair) { 'key' }
+      it 'raises an error' do
+        expect { subject }.to raise_error(::JWT::JWKError, 'Cannot create JWK from a String')
+      end
+    end
+  end
+end