Refine nested JWT API and validation

Author: ydah

README.md

@@ -339,19 +339,16 @@ inner_jwt = JWT.encode(inner_payload, inner_key, 'HS256')
 
 # Wrap it in an outer JWT with a different key/algorithm
 outer_key = OpenSSL::PKey::RSA.generate(2048)
-nested_jwt = JWT::NestedToken.sign(
-  inner_jwt,
-  algorithm: 'RS256',
-  key: outer_key
-)
+nested = JWT::NestedToken.new(inner_jwt)
+nested.sign!(algorithm: 'RS256', key: outer_key)
+nested_jwt = nested.jwt
 ```
 
 ### Decoding a Nested JWT
 
 ```ruby
 # Decode and verify all nesting levels
-tokens = JWT::NestedToken.decode(
-  nested_jwt,
+tokens = JWT::NestedToken.new(nested_jwt).verify!(
   keys: [
     { algorithm: 'RS256', key: outer_key.public_key },
     { algorithm: 'HS256', key: inner_key }
@@ -362,20 +359,6 @@ inner_payload = tokens.last.payload
 # => { 'user_id' => 123, 'role' => 'admin' }
 ```
 
-### Using JWT::Token.wrap
-
-You can also use `JWT::Token.wrap` to create nested tokens:
-
-```ruby
-inner = JWT::Token.new(payload: { sub: 'user' })
-inner.sign!(algorithm: 'HS256', key: 'inner_secret')
-
-outer = JWT::Token.wrap(inner)
-outer.sign!(algorithm: 'RS256', key: rsa_private_key)
-
-nested_jwt = outer.jwt
-```
-
 ### Checking for Nested JWTs
 
 ```ruby

lib/jwt/encoded_token.rb

@@ -13,7 +13,7 @@ module JWT
   #   encoded_token = JWT::EncodedToken.new(token.jwt)
   #   encoded_token.verify_signature!(algorithm: 'HS256', key: 'secret')
   #   encoded_token.payload # => {'pay' => 'load'}
-  class EncodedToken
+  class EncodedToken # rubocop:disable Metrics/ClassLength
     DEFAULT_CLAIMS = [:exp].freeze
 
     private_constant(:DEFAULT_CLAIMS)
@@ -205,7 +205,7 @@ def nested?
     def inner_token
       return nil unless nested?
 
-      EncodedToken.new(unverified_payload)
+      EncodedToken.new(decode_nested_payload)
     end
 
     # Unwraps all nesting levels and returns an array of tokens.
@@ -218,11 +218,13 @@ def inner_token
     #   all_tokens = token.unwrap_all
     #   all_tokens.first # => outermost token
     #   all_tokens.last  # => innermost token
-    def unwrap_all
+    def unwrap_all(max_depth:)
       tokens = [self]
       current = self
 
       while current.nested?
+        raise JWT::DecodeError, "Nested JWT exceeds maximum depth of #{max_depth}" if tokens.length >= max_depth
+
         current = current.inner_token
         tokens << current
       end
@@ -249,6 +251,14 @@ def decode_payload
       parse_and_decode(encoded_payload)
     end
 
+    def decode_nested_payload
+      raise JWT::DecodeError, 'Encoded payload is empty' if encoded_payload == ''
+
+      return encoded_payload if unencoded_payload?
+
+      ::JWT::Base64.url_decode(encoded_payload || '')
+    end
+
     def unencoded_payload?
       header['b64'] == false
     end

lib/jwt/nested_token.rb

@@ -1,111 +1,119 @@
 # frozen_string_literal: true
 
 module JWT
-  # Provides functionality for creating and decoding Nested JWTs
-  # as defined in RFC 7519 Section 5.2, Section 7.1 Step 5, and Appendix A.2.
+  # Represents a Nested JWT as defined in RFC 7519 Section 5.2, Section 7.1 Step 5, and Appendix A.2.
   #
-  # A Nested JWT is a JWT that is used as the payload of another JWT,
-  # allowing for multiple layers of signing or encryption.
+  # A Nested JWT wraps an existing JWT string as the payload of another signed JWT.
   #
   # @example Creating a Nested JWT
   #   inner_jwt = JWT.encode({ user_id: 123 }, 'inner_secret', 'HS256')
-  #   nested_jwt = JWT::NestedToken.sign(
-  #     inner_jwt,
-  #     algorithm: 'RS256',
-  #     key: rsa_private_key
-  #   )
+  #   nested = JWT::NestedToken.new(inner_jwt)
+  #   nested.sign!(algorithm: 'RS256', key: rsa_private_key)
+  #   nested.jwt
   #
-  # @example Decoding a Nested JWT
-  #   tokens = JWT::NestedToken.decode(
-  #     nested_jwt,
+  # @example Verifying a Nested JWT
+  #   nested = JWT::NestedToken.new(nested_jwt)
+  #   tokens = nested.verify!(
   #     keys: [
   #       { algorithm: 'RS256', key: rsa_public_key },
   #       { algorithm: 'HS256', key: 'inner_secret' }
   #     ]
   #   )
-  #   inner_payload = tokens.last.payload
+  #   tokens.last.payload
   #
   # @see https://datatracker.ietf.org/doc/html/rfc7519#section-5.2 RFC 7519 Section 5.2
   class NestedToken
-    # The content type header value for nested JWTs as per RFC 7519
     CTY_JWT = 'JWT'
+    MAX_DEPTH = 10
 
-    class << self
-      # Wraps an inner JWT with an outer JWS, creating a Nested JWT.
-      # Automatically sets the `cty` (content type) header to "JWT" as required by RFC 7519.
-      #
-      # @param inner_jwt [String] the inner JWT string to wrap
-      # @param algorithm [String] the signing algorithm for the outer JWS (e.g., 'RS256', 'HS256')
-      # @param key [Object] the signing key for the outer JWS
-      # @param header [Hash] additional header fields to include (cty is automatically set)
-      # @return [String] the Nested JWT string
-      #
-      # @raise [JWT::EncodeError] if signing fails
-      #
-      # @example Basic usage with HS256
-      #   inner_jwt = JWT.encode({ sub: 'user' }, 'secret', 'HS256')
-      #   nested = JWT::NestedToken.sign(inner_jwt, algorithm: 'HS256', key: 'outer_secret')
-      #
-      # @example With RSA and custom headers
-      #   nested = JWT::NestedToken.sign(
-      #     inner_jwt,
-      #     algorithm: 'RS256',
-      #     key: rsa_private_key,
-      #     header: { kid: 'my-key-id' }
-      #   )
-      def sign(inner_jwt, algorithm:, key:, header: {})
-        outer_header = header.merge('cty' => CTY_JWT)
-        token = Token.new(payload: inner_jwt, header: outer_header)
-        token.sign!(algorithm: algorithm, key: key)
-        token.jwt
-      end
+    # @return [String] the current JWT string represented by this instance.
+    attr_reader :jwt
+
+    # @return [Array<JWT::EncodedToken>, nil] verified tokens ordered from outermost to innermost.
+    attr_reader :tokens
+
+    # @param jwt [String] the JWT string to wrap or verify.
+    # @raise [ArgumentError] if the provided JWT is not a String.
+    def initialize(jwt)
+      raise ArgumentError, 'Provided JWT must be a String' unless jwt.is_a?(String)
+
+      @jwt = jwt
+    end
+
+    # Wraps the current JWT string in an outer JWS and replaces {#jwt} with the nested JWT.
+    # The payload is base64url-encoded directly from the JWT string (without JSON string encoding).
+    #
+    # @param algorithm [String, Object] the algorithm to use for signing.
+    # @param key [String, JWT::JWK::KeyBase] the key to use for signing.
+    # @param header [Hash] additional header fields to include in the outer token.
+    # @return [nil]
+    def sign!(algorithm:, key:, header: {})
+      signer = JWA.create_signer(algorithm: algorithm, key: key)
+      outer_header = (header || {})
+                     .transform_keys(&:to_s)
+                     .merge('cty' => CTY_JWT)
+
+      outer_header.merge!(signer.jwa.header) { |_header_key, old, _new| old }
+
+      encoded_header = ::JWT::Base64.url_encode(JWT::JSON.generate(outer_header))
+      encoded_payload = ::JWT::Base64.url_encode(jwt)
+      signing_input = [encoded_header, encoded_payload].join('.')
+      signature = signer.sign(data: signing_input)
+
+      @jwt = [encoded_header, encoded_payload, ::JWT::Base64.url_encode(signature)].join('.')
+      @tokens = nil
+      nil
+    end
 
-      # Decodes and verifies a Nested JWT, unwrapping all nesting levels.
-      # Each level's signature is verified using the corresponding key configuration.
-      #
-      # @param token [String] the Nested JWT string to decode
-      # @param keys [Array<Hash>] an array of key configurations for each nesting level,
-      #   ordered from outermost to innermost. Each hash should contain:
-      #   - `:algorithm` [String] the expected algorithm
-      #   - `:key` [Object] the verification key
-      # @return [Array<JWT::EncodedToken>] array of tokens from outermost to innermost
-      #
-      # @raise [JWT::DecodeError] if decoding fails at any level
-      # @raise [JWT::VerificationError] if signature verification fails at any level
-      #
-      # @example Decoding a two-level nested JWT
-      #   tokens = JWT::NestedToken.decode(
-      #     nested_jwt,
-      #     keys: [
-      #       { algorithm: 'RS256', key: rsa_public_key },
-      #       { algorithm: 'HS256', key: 'inner_secret' }
-      #     ]
-      #   )
-      #   inner_token = tokens.last
-      #   inner_token.payload # => { 'user_id' => 123 }
-      def decode(token, keys:)
-        tokens = []
-        current_token = token
-
-        keys.each_with_index do |key_config, index|
-          encoded_token = EncodedToken.new(current_token)
-          encoded_token.verify_signature!(
-            algorithm: key_config[:algorithm],
-            key: key_config[:key]
-          )
-
-          tokens << encoded_token
-
-          if encoded_token.nested?
-            current_token = encoded_token.unverified_payload
-          elsif index < keys.length - 1
-            raise JWT::DecodeError, 'Token is not nested but more keys were provided'
-          end
-        end
-
-        tokens.each(&:verify_claims!)
-        tokens
+    # Verifies signatures of all nested levels and the claims of the innermost token.
+    #
+    # @param keys [Array<Hash>] key configuration per nesting level (outermost to innermost).
+    # @param claims [Array<Symbol>, Hash, nil] claim verification options for the innermost token.
+    # @return [Array<JWT::EncodedToken>] verified tokens from outermost to innermost.
+    def verify!(keys:, claims: nil)
+      verify_signatures!(keys: keys)
+      verify_claims!(claims: claims)
+      tokens
+    end
+
+    # Verifies signatures of all nested levels.
+    #
+    # @param keys [Array<Hash>] key configuration per nesting level (outermost to innermost).
+    # @return [Array<JWT::EncodedToken>] verified tokens from outermost to innermost.
+    def verify_signatures!(keys:)
+      @tokens = EncodedToken.new(jwt).unwrap_all(max_depth: MAX_DEPTH)
+      validate_key_count!(keys)
+
+      tokens.each_with_index do |token, index|
+        key_config = keys[index]
+        token.verify_signature!(
+          algorithm: key_config[:algorithm],
+          key: key_config[:key],
+          key_finder: key_config[:key_finder]
+        )
       end
+
+      tokens
+    end
+
+    # Verifies claims of the innermost token after signatures have been verified.
+    #
+    # @param claims [Array<Symbol>, Hash, nil] claim verification options for the innermost token.
+    # @return [Array<JWT::EncodedToken>] verified tokens from outermost to innermost.
+    def verify_claims!(claims: nil)
+      raise JWT::DecodeError, 'Verify nested token signatures before verifying claims' unless tokens
+
+      innermost_token = tokens.last
+      claims.is_a?(Array) ? innermost_token.verify_claims!(*claims) : innermost_token.verify_claims!(claims)
+      tokens
+    end
+
+    private
+
+    def validate_key_count!(keys)
+      return if keys.length == tokens.length
+
+      raise JWT::DecodeError, "Expected #{tokens.length} key configurations, got #{keys.length}"
     end
   end
 end

lib/jwt/token.rb

@@ -127,31 +127,5 @@ def valid_claims?(*options)
     #
     # @return [String] the JWT token as a string.
     alias to_s jwt
-
-    class << self
-      # Wraps another JWT token, creating a Nested JWT.
-      # Sets the `cty` (content type) header to "JWT" as required by RFC 7519 Section 5.2.
-      #
-      # @param inner_token [JWT::Token, String] the token to wrap. Can be a JWT::Token instance
-      #   or a JWT string.
-      # @param header [Hash] additional header fields for the outer token
-      # @return [JWT::Token] a new token with the inner token as its payload and cty header set
-      #
-      # @example Wrapping a token
-      #   inner = JWT::Token.new(payload: { sub: 'user' })
-      #   inner.sign!(algorithm: 'HS256', key: 'secret')
-      #   outer = JWT::Token.wrap(inner)
-      #   outer.sign!(algorithm: 'RS256', key: rsa_private)
-      #
-      # @example Wrapping a JWT string
-      #   jwt_string = JWT.encode({ sub: 'user' }, 'secret', 'HS256')
-      #   outer = JWT::Token.wrap(jwt_string)
-      #
-      # @see https://datatracker.ietf.org/doc/html/rfc7519#section-5.2 RFC 7519 Section 5.2
-      def wrap(inner_token, header: {})
-        jwt_string = inner_token.is_a?(Token) ? inner_token.jwt : inner_token
-        new(payload: jwt_string, header: header.merge('cty' => 'JWT'))
-      end
-    end
   end
 end