Make extending algorithms easier

Co-authored-by: Matteo Pierro <[email protected]>

CHANGELOG.md

@@ -11,6 +11,7 @@
 **Fixes and enhancements:**
 
 - Refactor claim validators into their own classes [#605](https://github.com/jwt/ruby-jwt/pull/605) ([@anakinj](https://github.com/anakinj), [@MatteoPierro](https://github.com/MatteoPierro))
+- Allow extending available algorithms [#607](https://github.com/jwt/ruby-jwt/pull/607) ([@anakinj](https://github.com/anakinj))
 - Your contribution here
 
 ## [v2.8.2](https://github.com/jwt/ruby-jwt/tree/v2.8.2) (2024-06-18)

README.md

@@ -223,18 +223,22 @@ puts decoded_token
 
 ### **Custom algorithms**
 
-An object implementing custom signing or verification behaviour can be passed in the `algorithm` option when encoding and decoding. The given object needs to implement the method `valid_alg?` and `verify` and/or `alg` and `sign`, depending if object is used for encoding or decoding.
+When encoding or decoding a token, you can pass in a custom object through the `algorithm` option to handle signing or verification. This custom object must include or extend the `JWT::JWA::SigningAlgorithm` module and implement certain methods:
+
+- For decoding/verifying: The object must implement the methods `alg` and `verify`.
+- For encoding/signing: The object must implement the methods `alg` and `sign`.
+
+For customization options check the details from `JWT::JWA::SigningAlgorithm`.
+
 
 ```ruby
 module CustomHS512Algorithm
+  extend JWT::JWA::SigningAlgorithm
+
   def self.alg
     'HS512'
   end
 
-  def self.valid_alg?(alg_to_validate)
-    alg_to_validate == alg
-  end
-
   def self.sign(data:, signing_key:)
     OpenSSL::HMAC.digest(OpenSSL::Digest.new('sha512'), data, signing_key)
   end

lib/jwt/decode.rb

@@ -90,7 +90,7 @@ def allowed_algorithms
     end
 
     def resolve_allowed_algorithms
-      algs = given_algorithms.map { |alg| JWA.create(alg) }
+      algs = given_algorithms.map { |alg| JWA.resolve(alg) }
 
       sort_by_alg_header(algs)
     end

lib/jwt/encode.rb

@@ -6,14 +6,11 @@
 module JWT
   # Encoding logic for JWT
   class Encode
-    ALG_KEY = 'alg'
-
     def initialize(options)
       @payload          = options[:payload]
       @key              = options[:key]
-      @algorithm        = JWA.create(options[:algorithm])
+      @algorithm        = JWA.resolve(options[:algorithm])
       @headers          = options[:headers].transform_keys(&:to_s)
-      @headers[ALG_KEY] = @algorithm.alg
     end
 
     def segments
@@ -40,7 +37,7 @@ def encoded_header_and_payload
     end
 
     def encode_header
-      encode_data(@headers)
+      encode_data(@headers.merge(@algorithm.header(signing_key: @key)))
     end
 
     def encode_payload

lib/jwt/jwa.rb

@@ -8,54 +8,35 @@
   raise if defined?(RbNaCl)
 end
 
-require_relative 'jwa/hmac'
-require_relative 'jwa/eddsa'
+require_relative 'jwa/signing_algorithm'
+
 require_relative 'jwa/ecdsa'
-require_relative 'jwa/rsa'
-require_relative 'jwa/ps'
+require_relative 'jwa/eddsa'
+require_relative 'jwa/hmac'
 require_relative 'jwa/none'
+require_relative 'jwa/ps'
+require_relative 'jwa/rsa'
 require_relative 'jwa/unsupported'
 require_relative 'jwa/wrapper'
 
+if JWT.rbnacl_6_or_greater?
+  require_relative 'jwa/hmac_rbnacl'
+elsif JWT.rbnacl?
+  require_relative 'jwa/hmac_rbnacl_fixed'
+end
+
 module JWT
   module JWA
-    ALGOS = [Hmac, Ecdsa, Rsa, Eddsa, Ps, None, Unsupported].tap do |l|
-      if ::JWT.rbnacl_6_or_greater?
-        require_relative 'jwa/hmac_rbnacl'
-        l << Algos::HmacRbNaCl
-      elsif ::JWT.rbnacl?
-        require_relative 'jwa/hmac_rbnacl_fixed'
-        l << Algos::HmacRbNaClFixed
-      end
-    end.freeze
-
     class << self
-      def find(algorithm)
-        indexed[algorithm&.downcase]
-      end
-
-      def create(algorithm)
-        return algorithm if JWA.implementation?(algorithm)
-
-        Wrapper.new(*find(algorithm))
-      end
-
-      def implementation?(algorithm)
-        (algorithm.respond_to?(:valid_alg?) && algorithm.respond_to?(:verify)) ||
-          (algorithm.respond_to?(:alg) && algorithm.respond_to?(:sign))
-      end
+      def resolve(algorithm)
+        return find(algorithm) if algorithm.is_a?(String) || algorithm.is_a?(Symbol)
 
-      private
-
-      def indexed
-        @indexed ||= begin
-          fallback = [nil, Unsupported]
-          ALGOS.each_with_object(Hash.new(fallback)) do |cls, hash|
-            cls.const_get(:SUPPORTED).each do |alg|
-              hash[alg.downcase] = [alg, cls]
-            end
-          end
+        unless algorithm.is_a?(SigningAlgorithm)
+          Deprecations.warning('Custom algorithms are required to include JWT::JWA::SigningAlgorithm')
+          return Wrapper.new(algorithm)
         end
+
+        algorithm
       end
     end
   end

lib/jwt/jwa/ecdsa.rb

@@ -2,8 +2,35 @@
 
 module JWT
   module JWA
-    module Ecdsa
-      module_function
+    class Ecdsa
+      include JWT::JWA::SigningAlgorithm
+
+      def initialize(alg, digest)
+        @alg = alg
+        @digest = OpenSSL::Digest.new(digest)
+      end
+
+      def sign(data:, signing_key:)
+        curve_definition = curve_by_name(signing_key.group.curve_name)
+        key_algorithm = curve_definition[:algorithm]
+        if alg != key_algorithm
+          raise IncorrectAlgorithm, "payload algorithm is #{alg} but #{key_algorithm} signing key was provided"
+        end
+
+        asn1_to_raw(signing_key.dsa_sign_asn1(digest.digest(data)), signing_key)
+      end
+
+      def verify(data:, signature:, verification_key:)
+        curve_definition = curve_by_name(verification_key.group.curve_name)
+        key_algorithm = curve_definition[:algorithm]
+        if alg != key_algorithm
+          raise IncorrectAlgorithm, "payload algorithm is #{alg} but #{key_algorithm} verification key was provided"
+        end
+
+        verification_key.dsa_verify_asn1(digest.digest(data), raw_to_asn1(signature, verification_key))
+      rescue OpenSSL::PKey::PKeyError
+        raise JWT::VerificationError, 'Signature verification raised'
+      end
 
       NAMED_CURVES = {
         'prime256v1' => {
@@ -28,36 +55,22 @@ module Ecdsa
         }
       }.freeze
 
-      SUPPORTED = NAMED_CURVES.map { |_, c| c[:algorithm] }.uniq.freeze
+      NAMED_CURVES.each_value do |v|
+        register_algorithm(new(v[:algorithm], v[:digest]))
+      end
 
-      def sign(algorithm, msg, key)
-        curve_definition = curve_by_name(key.group.curve_name)
-        key_algorithm = curve_definition[:algorithm]
-        if algorithm != key_algorithm
-          raise IncorrectAlgorithm, "payload algorithm is #{algorithm} but #{key_algorithm} signing key was provided"
+      def self.curve_by_name(name)
+        NAMED_CURVES.fetch(name) do
+          raise UnsupportedEcdsaCurve, "The ECDSA curve '#{name}' is not supported"
         end
-
-        digest = OpenSSL::Digest.new(curve_definition[:digest])
-        asn1_to_raw(key.dsa_sign_asn1(digest.digest(msg)), key)
       end
 
-      def verify(algorithm, public_key, signing_input, signature)
-        curve_definition = curve_by_name(public_key.group.curve_name)
-        key_algorithm = curve_definition[:algorithm]
-        if algorithm != key_algorithm
-          raise IncorrectAlgorithm, "payload algorithm is #{algorithm} but #{key_algorithm} verification key was provided"
-        end
+      private
 
-        digest = OpenSSL::Digest.new(curve_definition[:digest])
-        public_key.dsa_verify_asn1(digest.digest(signing_input), raw_to_asn1(signature, public_key))
-      rescue OpenSSL::PKey::PKeyError
-        raise JWT::VerificationError, 'Signature verification raised'
-      end
+      attr_reader :digest
 
       def curve_by_name(name)
-        NAMED_CURVES.fetch(name) do
-          raise UnsupportedEcdsaCurve, "The ECDSA curve '#{name}' is not supported"
-        end
+        self.class.curve_by_name(name)
       end
 
       def raw_to_asn1(signature, private_key)

lib/jwt/jwa/eddsa.rb

@@ -2,41 +2,33 @@
 
 module JWT
   module JWA
-    module Eddsa
-      SUPPORTED = %w[ED25519 EdDSA].freeze
-      SUPPORTED_DOWNCASED = SUPPORTED.map(&:downcase).freeze
+    class Eddsa
+      include JWT::JWA::SigningAlgorithm
 
-      class << self
-        def sign(algorithm, msg, key)
-          unless key.is_a?(RbNaCl::Signatures::Ed25519::SigningKey)
-            raise EncodeError, "Key given is a #{key.class} but has to be an RbNaCl::Signatures::Ed25519::SigningKey"
-          end
-
-          validate_algorithm!(algorithm)
+      def initialize(alg)
+        @alg = alg
+      end
 
-          key.sign(msg)
+      def sign(data:, signing_key:)
+        unless signing_key.is_a?(RbNaCl::Signatures::Ed25519::SigningKey)
+          raise_encode_error!("Key given is a #{signing_key.class} but has to be an RbNaCl::Signatures::Ed25519::SigningKey")
         end
 
-        def verify(algorithm, public_key, signing_input, signature)
-          unless public_key.is_a?(RbNaCl::Signatures::Ed25519::VerifyKey)
-            raise DecodeError, "key given is a #{public_key.class} but has to be a RbNaCl::Signatures::Ed25519::VerifyKey"
-          end
-
-          validate_algorithm!(algorithm)
+        signing_key.sign(data)
+      end
 
-          public_key.verify(signature, signing_input)
-        rescue RbNaCl::CryptoError
-          false
+      def verify(data:, signature:, verification_key:)
+        unless verification_key.is_a?(RbNaCl::Signatures::Ed25519::VerifyKey)
+          raise_decode_error!("key given is a #{verification_key.class} but has to be a RbNaCl::Signatures::Ed25519::VerifyKey")
         end
 
-        private
-
-        def validate_algorithm!(algorithm)
-          return if SUPPORTED_DOWNCASED.include?(algorithm.downcase)
-
-          raise IncorrectAlgorithm, "Algorithm #{algorithm} not supported. Supported algoritms are #{SUPPORTED.join(', ')}"
-        end
+        verification_key.verify(signature, data)
+      rescue RbNaCl::CryptoError
+        false
       end
+
+      register_algorithm(new('ED25519'))
+      register_algorithm(new('EdDSA'))
     end
   end
 end

lib/jwt/jwa/hmac.rb

@@ -2,35 +2,40 @@
 
 module JWT
   module JWA
-    module Hmac
-      module_function
+    class Hmac
+      include JWT::JWA::SigningAlgorithm
 
-      MAPPING = {
-        'HS256' => OpenSSL::Digest::SHA256,
-        'HS384' => OpenSSL::Digest::SHA384,
-        'HS512' => OpenSSL::Digest::SHA512
-      }.freeze
-
-      SUPPORTED = MAPPING.keys
+      def initialize(alg, digest)
+        @alg = alg
+        @digest = digest
+      end
 
-      def sign(algorithm, msg, key)
-        key ||= ''
+      def sign(data:, signing_key:)
+        signing_key ||= ''
 
-        raise JWT::DecodeError, 'HMAC key expected to be a String' unless key.is_a?(String)
+        raise_verify_error!('HMAC key expected to be a String') unless signing_key.is_a?(String)
 
-        OpenSSL::HMAC.digest(MAPPING[algorithm].new, key, msg)
+        OpenSSL::HMAC.digest(digest.new, signing_key, data)
       rescue OpenSSL::HMACError => e
-        if key == '' && e.message == 'EVP_PKEY_new_mac_key: malloc failure'
-          raise JWT::DecodeError, 'OpenSSL 3.0 does not support nil or empty hmac_secret'
+        if signing_key == '' && e.message == 'EVP_PKEY_new_mac_key: malloc failure'
+          raise_verify_error!('OpenSSL 3.0 does not support nil or empty hmac_secret')
         end
 
         raise e
       end
 
-      def verify(algorithm, key, signing_input, signature)
-        SecurityUtils.secure_compare(signature, sign(algorithm, signing_input, key))
+      def verify(data:, signature:, verification_key:)
+        SecurityUtils.secure_compare(signature, sign(data: data, signing_key: verification_key))
       end
 
+      register_algorithm(new('HS256', OpenSSL::Digest::SHA256))
+      register_algorithm(new('HS384', OpenSSL::Digest::SHA384))
+      register_algorithm(new('HS512', OpenSSL::Digest::SHA512))
+
+      private
+
+      attr_reader :digest
+
       # Copy of https://github.com/rails/rails/blob/v7.0.3.1/activesupport/lib/active_support/security_utils.rb
       # rubocop:disable Naming/MethodParameterName, Style/StringLiterals, Style/NumericPredicate
       module SecurityUtils