First idea of a extendable JWT decoder

.rubocop.yml

@@ -50,7 +50,7 @@ Style/SignalException:
   Enabled: false
 
 Metrics/AbcSize:
-  Max: 20
+  Max: 21
 
 Metrics/ClassLength:
   Max: 101

lib/jwt.rb

@@ -1,6 +1,7 @@
 # frozen_string_literal: true
 
 require 'jwt/base64'
+require 'jwt/extension'
 require 'jwt/json'
 require 'jwt/decode'
 require 'jwt/default_options'
@@ -15,6 +16,10 @@
 module JWT
   include JWT::DefaultOptions
 
+  def self.included(cls)
+    cls.extend(JWT::Extension::ClassMethods)
+  end
+
   module_function
 
   def encode(payload, key, algorithm = 'HS256', header_fields = {})

lib/jwt/decode.rb

@@ -15,14 +15,12 @@ def initialize(jwt, key, verify, options, &keyfinder)
       @options = options
       @segments = jwt.split('.')
       @verify = verify
-      @signature = ''
       @keyfinder = keyfinder
     end
 
     def decode_segments
       validate_segment_count!
       if @verify
-        decode_crypto
         verify_signature
         verify_claims
       end
@@ -40,7 +38,7 @@ def verify_signature
       @key = find_key(&@keyfinder) if @keyfinder
       @key = ::JWT::JWK::KeyFinder.new(jwks: @options[:jwks]).key_for(header['kid']) if @options[:jwks]
 
-      Signature.verify(header['alg'], @key, signing_input, @signature)
+      Signature.verify(header['alg'], @key, signing_input, signature)
     end
 
     def options_includes_algo_in_header?
@@ -85,24 +83,33 @@ def segment_length
       @segments.count
     end
 
-    def decode_crypto
-      @signature = JWT::Base64.url_decode(@segments[2] || '')
+    def signature
+      @signature ||= JWT::Base64.url_decode(@segments[2] || '')
     end
 
     def header
-      @header ||= parse_and_decode @segments[0]
+      @header ||= decode_and_parse_header(@segments[0])
     end
 
     def payload
-      @payload ||= parse_and_decode @segments[1]
+      @payload ||= decode_and_parse_payload(@segments[1])
     end
 
     def signing_input
       @segments.first(2).join('.')
     end
 
-    def parse_and_decode(segment)
-      JWT::JSON.parse(JWT::Base64.url_decode(segment))
+    def decode_and_parse_header(raw_header)
+      json_parse(JWT::Base64.url_decode(raw_header))
+    end
+
+    def decode_and_parse_payload(raw_payload)
+      raw_payload = @options[:decode_payload_proc].call(header, raw_payload, signature) if @options[:decode_payload_proc]
+      json_parse(JWT::Base64.url_decode(raw_payload))
+    end
+
+    def json_parse(decoded_segment)
+      JWT::JSON.parse(decoded_segment)
     rescue ::JSON::ParserError
       raise JWT::DecodeError, 'Invalid segment encoding'
     end

lib/jwt/extension.rb

@@ -0,0 +1,26 @@
+module JWT
+  module Extension
+    module ClassMethods
+      def decode_payload(&block)
+        @decode_payload_block = block
+      end
+
+      def decode(payload, options = {})
+        segments = ::JWT::Decode.new(payload,
+          options.delete(:key),
+          true,
+          create_decode_options(options)).decode_segments
+        {
+          header: segments.last,
+          payload: segments.first
+        }
+      end
+
+      private
+
+      def create_decode_options(given_options)
+        ::JWT::DefaultOptions::DEFAULT_OPTIONS.merge(decode_payload_proc: @decode_payload_block).merge(given_options)
+      end
+    end
+  end
+end

spec/extension_spec.rb

@@ -0,0 +1,39 @@
+# frozen_string_literal: true
+
+require 'securerandom'
+
+RSpec.describe JWT::Extension do
+  subject(:extension) do
+    Class.new do
+      include JWT
+    end
+  end
+
+  let(:secret) { SecureRandom.hex }
+  let(:payload) { { 'pay' => 'load'} }
+  let(:encoded_payload) { ::JWT.encode(payload, secret, 'HS256') }
+
+  describe '.decode' do
+    it { is_expected.to respond_to(:decode) }
+
+    context 'when nothing special is defined' do
+      it 'verifies a token and returns the data' do
+        expect(extension.decode(encoded_payload, key: secret)).to eq(header: { 'alg' => 'HS256' }, payload: payload)
+      end
+    end
+
+    context 'when a decode_payload block is given' do
+      before do
+        extension.decode_payload do |_header, raw_payload, _signature|
+          payload_content = JWT::JSON.parse(JWT::Base64.url_decode(raw_payload))
+          payload_content['pay'].reverse!
+          JWT::Base64.url_encode(JWT::JSON.generate(payload_content))
+        end
+      end
+
+      it 'lets before decode process the raw payload before verifying' do
+        expect(extension.decode(encoded_payload, key: secret)).to eq(header: { 'alg' => 'HS256' }, payload: {'pay' => 'daol'})
+      end
+    end
+  end
+end