Decode class rewritten to increase understanding and to soon support custom algorithms

Author: anakinj

lib/jwt.rb

@@ -2,6 +2,7 @@
 
 require 'base64'
 require 'jwt/extension'
+require 'jwt/decode_token'
 require 'jwt/json'
 require 'jwt/decode'
 require 'jwt/default_options'

lib/jwt/decode.rb

@@ -123,44 +123,20 @@ def algorithm
     end
 
     def header
-      @header ||= decode_and_parse_header(@segments[0])
+      @header ||= parse_and_decode @segments[0]
     end
 
     def payload
-      @payload ||= decode_and_parse_payload(@segments[1])
+      @payload ||= parse_and_decode @segments[1]
     end
 
     def signing_input
       @segments.first(2).join('.')
     end
 
-    def decode_and_parse_header(raw_header)
-      json_parse(decode_header(raw_header))
-    end
-
-    def decode_and_parse_payload(raw_payload)
-      decode_payload(raw_payload)
-    end
-
-    def decode_payload(raw_segment)
-      if @options[:decode_payload_proc]
-        @options[:decode_payload_proc].call(raw_segment, header, @signature)
-      else
-        json_parse(Base64.urlsafe_decode64(raw_segment))
-      end
-    rescue ArgumentError
-      raise JWT::DecodeError, 'Invalid segment encoding'
-    end
-
-    def decode_header(raw_segment)
-      Base64.urlsafe_decode64(raw_segment)
-    rescue ArgumentError
-      raise JWT::DecodeError, 'Invalid segment encoding'
-    end
-
-    def json_parse(decoded_segment)
-      JWT::JSON.parse(decoded_segment)
-    rescue ::JSON::ParserError
+    def parse_and_decode(segment)
+      JWT::JSON.parse(Base64.urlsafe_decode64(segment))
+    rescue ::JSON::ParserError, ArgumentError
       raise JWT::DecodeError, 'Invalid segment encoding'
     end
   end

lib/jwt/decode_token.rb

@@ -0,0 +1,141 @@
+# frozen_string_literal: true
+
+require 'jwt/signature'
+require 'jwt/verify'
+require 'jwt/x5c_key_finder'
+
+module JWT
+  class DecodeToken
+    def initialize(token, options = {})
+      raise ArgumentError, 'Provided token is not a String object' unless token.is_a?(String)
+
+      @token   = token
+      @options = options
+    end
+
+    def decoded_segments
+      validate_segment_count!
+
+      if verify?
+        verify_alg_header!
+        verify_signature!
+        verify_claims!
+      end
+
+      [payload, header]
+    end
+
+    private
+
+    attr_reader :token, :options
+
+    def algorithms
+      @algorithms ||= Array(options[:algorithms])
+    end
+
+    def segments
+      @segments ||= token.split('.')
+    end
+
+    def signature
+      @signature ||= Base64.urlsafe_decode64(segments[2] || '')
+    end
+
+    def header
+      @header ||= decode_header(segments[0])
+    end
+
+    def payload
+      @payload ||= decode_payload(segments[1])
+    end
+
+    def signing_input
+      segments.first(2).join('.')
+    end
+
+    def verify?
+      options[:verify] != false
+    end
+
+    def key
+      @key ||=
+        if options[:jwks]
+          ::JWT::JWK::KeyFinder.new(jwks: options[:jwks]).key_for(header['kid'])
+        elsif (x5c_options = options[:x5c])
+          ::JWT::X5cKeyFinder.new(x5c_options[:root_certificates], x5c_options[:crls]).from(header['x5c'])
+        else
+          options[:key]
+        end
+    end
+
+    def verify_alg_header!
+      return unless valid_algorithms.empty?
+
+      raise JWT::IncorrectAlgorithm, 'Expected a different algorithm'
+    end
+
+    def valid_algorithms
+      @valid_algorithms ||= algorithms.select do |algorithm|
+        if algorithm.is_a?(String)
+          header['alg'] == algorithm
+        else
+          algorithm.valid_alg?(header['alg'])
+        end
+      end
+    end
+
+    def verify_signature!
+      return if valid_algorithms.any? { |algorithm| verify_signature_for?(algorithm, key) }
+
+      raise JWT::VerificationError, 'Signature verification failed'
+    end
+
+    def verify_signature_for?(algorithm, key)
+      if algorithm.is_a?(String)
+        raise JWT::DecodeError, 'No verification key available' unless key
+
+        Array(key).any? { |k| Signature.verify(algorithm, k, signing_input, signature) }
+      else
+        algorithm.verify(signing_input, signature, key: key, header: header, payload: payload)
+      end
+    end
+
+    def verify_claims!
+      Verify.verify_claims(payload, options)
+      Verify.verify_required_claims(payload, options)
+    end
+
+    def validate_segment_count!
+      segment_count = segments.size
+
+      return if segment_count == 3
+      return if segment_count == 2 && (!verify? || header['alg'] == 'none')
+
+      raise JWT::DecodeError, 'Not enough or too many segments'
+    end
+
+    def decode_header(raw_header)
+      decode_segment_default(raw_header)
+    end
+
+    def decode_payload(raw_segment)
+      if @options[:decode_payload_proc]
+        @options[:decode_payload_proc].call(raw_segment, header, signature)
+      else
+        decode_segment_default(raw_segment)
+      end
+    end
+
+    def decode_segment_default(raw_segment)
+      json_parse(Base64.urlsafe_decode64(raw_segment))
+    rescue ArgumentError
+      raise JWT::DecodeError, 'Invalid segment encoding'
+    end
+
+    def json_parse(decoded_segment)
+      JWT::JSON.parse(decoded_segment)
+    rescue ::JSON::ParserError
+      raise JWT::DecodeError, 'Invalid segment encoding'
+    end
+  end
+end

lib/jwt/encode.rb

@@ -14,7 +14,13 @@ def initialize(options)
       @options = options
       @payload = options[:payload]
       @key = options[:key]
-      _, @algorithm = Algos.find(options[:algorithm])
+
+      if (@algorithm_implementation = options[:algorithm_implementation]).nil?
+        _, @algorithm = Algos.find(options[:algorithm])
+      else
+        @algorithm = @algorithm_implementation.alg
+      end
+
       @headers = options[:headers].each_with_object({}) { |(key, value), headers| headers[key.to_s] = value }
     end
 
@@ -58,7 +64,13 @@ def encode_payload
     def encode_signature
       return '' if @algorithm == ALG_NONE
 
-      Base64.urlsafe_encode64(JWT::Signature.sign(@algorithm, encoded_header_and_payload, @key), padding: false)
+      Base64.urlsafe_encode64(signature, padding: false)
+    end
+
+    def signature
+      return @algorithm_implementation.sign(encoded_header_and_payload, key: @key) if @algorithm_implementation
+
+      JWT::Signature.sign(@algorithm, encoded_header_and_payload, @key)
     end
 
     def encode(data)

lib/jwt/extension/decode.rb

@@ -18,26 +18,24 @@ def jwk_resolver(&block)
         @jwk_resolver
       end
 
-      def decode!(payload, options = {})
-        ::JWT::Decode.new(payload,
-                          decode_signing_key_from_options(options),
-                          true,
-                          create_decode_options(options)).decode_segments
+      def decode!(token, options = {})
+        Internals.decode!(token, options, self)
       end
 
-      def decode_signing_key_from_options(options)
-        options[:signing_key] || self.signing_key
-      end
-
-      def create_decode_options(given_options)
-        ::JWT::DefaultOptions::DEFAULT_OPTIONS.merge(decode_payload_proc: self.decode_payload,
-                                                     algorithms: self.decoding_algorithms,
-                                                     jwks: self.jwk_resolver)
-          .merge(given_options)
-      end
+      module Internals
+        class << self
+          def decode!(token, options, context)
+            ::JWT::DecodeToken.new(token, build_decode_options(options, context)).decoded_segments
+          end
 
-      def decoding_algorithms
-        (Array(self.algorithm) + Array(self.algorithms)).uniq
+          def build_decode_options(options, context)
+            ::JWT::DefaultOptions::DEFAULT_OPTIONS.merge(key: options[:signing_key] || context.verification_key || context.signing_key,
+                                                         decode_payload_proc: context.decode_payload,
+                                                         algorithms: (Array(context.algorithm) + Array(context.algorithms)).uniq,
+                                                         jwks: context.jwk_resolver)
+              .merge(options)
+          end
+        end
       end
     end
   end

lib/jwt/extension/encode.rb

@@ -14,20 +14,33 @@ def encode_payload(&block)
       end
 
       def encode!(payload, options = {})
-        ::JWT::Encode.new(
-          payload: payload,
-          key: signing_key_from_options(options),
-          algorithm: self.algorithm,
-          encode_payload_proc: self.encode_payload,
-          headers: Array(options[:headers])
-        ).segments
+        Internals.encode!(payload, options, self)
       end
 
-      def signing_key_from_options(options)
-        key = options[:signing_key] || self.signing_key
-        raise ::JWT::SigningKeyMissing, 'No key given for signing' if key.nil?
+      module Internals
+        class << self
+          def encode!(payload, options, context)
+            ::JWT::Encode.new(build_options(payload, options, context)).segments
+          end
 
-        key
+          def build_options(payload, options, context)
+            opts = {
+              payload: payload,
+              key: options[:key] || context.signing_key,
+              encode_payload_proc: context.encode_payload,
+              headers: Array(options[:headers])
+            }
+
+            if (algo = context.algorithm).is_a?(String)
+              opts[:algorithm] = algo
+              raise ::JWT::SigningKeyMissing, 'No key given for signing' if opts[:key].nil?
+            else
+              opts[:algorithm_implementation] = algo
+            end
+
+            opts
+          end
+        end
       end
     end
   end

lib/jwt/extension/keys.rb

@@ -7,6 +7,11 @@ def signing_key(value = nil)
         @signing_key = value unless value.nil?
         @signing_key
       end
+
+      def verification_key(value = nil)
+        @verification_key = value unless value.nil?
+        @verification_key
+      end
     end
   end
 end

spec/extension/example_custom_algorithm_spec.rb

@@ -0,0 +1,49 @@
+# frozen_string_literal: true
+
+RSpec.describe 'Custom Signing algorithm' do
+  let(:payload) { { 'pay' => 'load'} }
+
+  let(:signing_algo) do
+    Class.new do
+      class << self
+        def alg
+          'CustomStatic'
+        end
+
+        def valid_alg?(_alg)
+          true
+        end
+
+        def sign(_to_sign, _options)
+          'static'
+        end
+
+        def verify(_to_verify, signature, _options)
+          signature == 'static'
+        end
+      end
+    end
+  end
+
+  subject(:extension) do
+    algo = signing_algo
+
+    Class.new do
+      include JWT
+      algorithm algo
+    end
+  end
+
+  context 'when encoding' do
+    it 'adds the custom signature to the end' do
+      expect(::Base64.decode64(subject.encode!(payload).split('.')[2])).to eq('static')
+    end
+  end
+
+  context 'when decoding signed token' do
+    let(:presigned_token) { subject.encode!(payload) }
+    it 'verifies and decodes the payload' do
+      expect(subject.decode!(presigned_token)).to eq([{'pay' => 'load'}, {'alg' => 'CustomStatic'}])
+    end
+  end
+end