jwt-pq 0.3.0 → 0.5.0

data/lib/jwt/pq/hybrid_key.rb

@@ -2,17 +2,52 @@
 
 module JWT
   module PQ
-    # Composite key combining an Ed25519 keypair with an ML-DSA keypair
-    # for hybrid EdDSA + ML-DSA JWT signatures.
+    # A composite key that pairs an Ed25519 keypair with an ML-DSA keypair
+    # for hybrid `EdDSA+ML-DSA-*` JWT signatures.
+    #
+    # Hybrid mode concatenates the two signatures (`ed25519 || ml_dsa`) so
+    # that a verifier only accepts the token if **both** signatures are
+    # valid. The classical half remains secure against today's attackers
+    # while the post-quantum half resists a future cryptographically
+    # relevant quantum computer.
+    #
+    # Requires the `ed25519` gem (or `jwt-eddsa`, which depends on it).
+    # Use {JWT::PQ.hybrid_available?} to probe availability.
+    #
+    # @example Generate a hybrid key and encode a JWT
+    #   key = JWT::PQ::HybridKey.generate(:ml_dsa_65)
+    #   token = JWT.encode({ sub: "u-1" }, key, "EdDSA+ML-DSA-65")
+    #
+    # @example Verification-only hybrid key
+    #   verifier = JWT::PQ::HybridKey.new(
+    #     ed25519: ed25519_verify_key,
+    #     ml_dsa:  JWT::PQ::Key.from_public_key(:ml_dsa_65, pub_bytes)
+    #   )
     class HybridKey
-      attr_reader :ed25519_signing_key, :ed25519_verify_key, :ml_dsa_key
+      # @return [Ed25519::SigningKey, nil] Ed25519 signing key, or nil for
+      #   verification-only.
+      attr_reader :ed25519_signing_key
 
-      # @param ed25519 [Ed25519::SigningKey, Ed25519::VerifyKey] Ed25519 key
-      # @param ml_dsa [JWT::PQ::Key] ML-DSA key
+      # @return [Ed25519::VerifyKey] Ed25519 verification key.
+      attr_reader :ed25519_verify_key
+
+      # @return [JWT::PQ::Key] ML-DSA keypair (public-only or full).
+      attr_reader :ml_dsa_key
+
+      # Build a hybrid key from existing Ed25519 and ML-DSA components.
+      #
+      # Pass an `Ed25519::SigningKey` for a full signing key, or an
+      # `Ed25519::VerifyKey` for verification-only.
+      #
+      # @param ed25519 [Ed25519::SigningKey, Ed25519::VerifyKey] Ed25519 key.
+      # @param ml_dsa [JWT::PQ::Key] ML-DSA keypair.
+      # @raise [MissingDependencyError] if the `ed25519` gem is not installed.
+      # @raise [KeyError] if `ed25519` is not one of the accepted key types.
       def initialize(ed25519:, ml_dsa:)
         require_eddsa_dependency!
 
         @ml_dsa_key = ml_dsa
+        @op_mutex = Mutex.new
 
         case ed25519
         when Ed25519::SigningKey
@@ -26,7 +61,15 @@ module JWT
         end
       end
 
-      # Generate a new hybrid keypair.
+      # Generate a fresh hybrid keypair.
+      #
+      # Creates both an Ed25519 `SigningKey` and an ML-DSA keypair of the
+      # requested parameter set.
+      #
+      # @param ml_dsa_algorithm [Symbol, String] one of `:ml_dsa_44`,
+      #   `:ml_dsa_65`, `:ml_dsa_87`. Defaults to `:ml_dsa_65`.
+      # @return [HybridKey] a full hybrid keypair (signing + verification).
+      # @raise [MissingDependencyError] if the `ed25519` gem is not installed.
       def self.generate(ml_dsa_algorithm = :ml_dsa_65)
         require_eddsa_dependency!
 
@@ -36,38 +79,88 @@ module JWT
         new(ed25519: ed_key, ml_dsa: ml_key)
       end
 
-      # Whether both keys have private components (can sign).
+      # @return [Boolean] true when both halves have private components and
+      #   the key can be used for signing.
       def private?
         !@ed25519_signing_key.nil? && @ml_dsa_key.private?
       end
 
-      # The ML-DSA algorithm name (e.g., "ML-DSA-65").
+      # @return [String] the ML-DSA algorithm name (e.g. `"ML-DSA-65"`).
       def algorithm
         @ml_dsa_key.algorithm
       end
 
-      # The hybrid algorithm name (e.g., "EdDSA+ML-DSA-65").
+      # @return [String] the hybrid JWT algorithm name
+      #   (e.g. `"EdDSA+ML-DSA-65"`).
       def hybrid_algorithm
         "EdDSA+#{@ml_dsa_key.algorithm}"
       end
 
-      # Zero and discard private key material from both key components.
-      # After calling this, the key can only be used for verification.
+      # Produce a hybrid signature (Ed25519 ‖ ML-DSA) over `data`.
+      #
+      # Thread-safe: both component signatures are taken under the hybrid
+      # key's own mutex, and {#destroy!} contends on the same mutex. That
+      # guarantees a concurrent `destroy!` cannot zero the Ed25519 seed
+      # while libsodium is mid-sign, and cannot produce a half-signed
+      # output (Ed25519 succeeds, ML-DSA fails because the buffer was
+      # just zeroed). Lock order is hybrid mutex → ML-DSA mutex; callers
+      # must not invoke any Key method while holding another lock that
+      # might be taken by `destroy!`.
+      #
+      # @param data [String] message bytes to sign.
+      # @return [String] concatenated signature — 64 bytes of Ed25519
+      #   followed by the ML-DSA signature.
+      # @raise [KeyError] if either half is missing its private component.
+      def sign(data)
+        @op_mutex.synchronize do
+          raise KeyError, "Ed25519 private key not available — cannot sign" unless @ed25519_signing_key
+
+          ed_sig = @ed25519_signing_key.sign(data)
+          ml_sig = @ml_dsa_key.sign(data)
+          ed_sig + ml_sig
+        end
+      end
+
+      # Zero and discard private key material from both halves.
+      #
+      # After calling this, {#private?} becomes false and the key can only
+      # be used for verification. Idempotent — safe on verification-only keys.
+      #
+      # Thread-safe: serialized on the hybrid key's own mutex (which also
+      # guards {#sign}) and internally delegates to
+      # {JWT::PQ::Key#destroy!}, which uses its own mutex. A concurrent
+      # {#sign} will block until `destroy!` completes and then raise
+      # `KeyError`, never observing a half-destroyed state.
+      #
+      # Ed25519 wipe: `Ed25519::SigningKey` stores the private seed in two
+      # separate Strings — `@seed` (32 bytes) returned by `#to_bytes`, and
+      # `@keypair` (64 bytes = `seed || public_key`) returned by `#keypair`.
+      # Both are `attr_reader`-backed and hand out the internal String by
+      # reference, so we must zero both in place — wiping only `@seed`
+      # leaves the first 32 bytes of `@keypair` holding the seed until GC.
+      #
+      # @return [true]
       def destroy!
-        @ml_dsa_key.destroy!
-        if @ed25519_signing_key
-          seed = @ed25519_signing_key.to_bytes
-          seed.replace("\0" * seed.bytesize)
-          @ed25519_signing_key = nil
+        @op_mutex.synchronize do
+          @ml_dsa_key.destroy!
+          if @ed25519_signing_key
+            seed = @ed25519_signing_key.to_bytes
+            seed.replace("\0" * seed.bytesize)
+            keypair = @ed25519_signing_key.keypair
+            keypair.replace("\0" * keypair.bytesize)
+            @ed25519_signing_key = nil
+          end
         end
         true
       end
 
+      # @return [String] short diagnostic string — never contains key material.
       def inspect
         "#<#{self.class} algorithm=#{hybrid_algorithm} private=#{private?}>"
       end
       alias to_s inspect
 
+      # @api private
       def self.require_eddsa_dependency!
         require "ed25519"
       rescue LoadError

data/lib/jwt/pq/jwk.rb

@@ -1,23 +1,42 @@
 # frozen_string_literal: true
 
 require "base64"
+require "json"
 require "openssl"
 
 module JWT
   module PQ
-    # JWK (JSON Web Key) support for ML-DSA keys.
+    # JWK (JSON Web Key) import/export for ML-DSA keys.
     #
-    # Follows the draft-ietf-cose-dilithium conventions:
-    #   kty: "AKP" (Algorithm Key Pair)
-    #   alg: "ML-DSA-44", "ML-DSA-65", or "ML-DSA-87"
-    #   pub: base64url-encoded public key
-    #   priv: base64url-encoded private key (optional)
+    # Follows the `draft-ietf-cose-dilithium` conventions for the `AKP`
+    # ("Algorithm Key Pair") key type:
+    #
+    # - `kty`: `"AKP"`
+    # - `alg`: `"ML-DSA-44"`, `"ML-DSA-65"`, or `"ML-DSA-87"`
+    # - `pub`: base64url-encoded public key (no padding)
+    # - `priv`: base64url-encoded private key (optional, no padding)
+    # - `kid`: RFC 7638 thumbprint over the required members
+    #
+    # @example Export and re-import
+    #   jwk = JWT::PQ::JWK.new(key).export
+    #   restored = JWT::PQ::JWK.import(jwk)
+    #
+    # @see https://datatracker.ietf.org/doc/draft-ietf-cose-dilithium/ draft-ietf-cose-dilithium
+    # @see https://www.rfc-editor.org/rfc/rfc7638 RFC 7638 (JWK Thumbprint)
     class JWK
+      # Algorithm names accepted in the `alg` field.
       ALGORITHMS = MlDsa::ALGORITHMS.keys.freeze
+
+      # Value of the `kty` field for all ML-DSA JWKs.
       KTY = "AKP"
 
+      # @return [JWT::PQ::Key] the wrapped key.
       attr_reader :key
 
+      # Wrap a {JWT::PQ::Key} for JWK operations.
+      #
+      # @param key [JWT::PQ::Key] the key to export or thumbprint.
+      # @raise [KeyError] if `key` is not a {JWT::PQ::Key}.
       def initialize(key)
         raise KeyError, "Expected a JWT::PQ::Key, got #{key.class}" unless key.is_a?(JWT::PQ::Key)
 
@@ -25,6 +44,14 @@ module JWT
       end
 
       # Export the key as a JWK hash.
+      #
+      # By default, only the public material is included. Pass
+      # `include_private: true` to emit the `priv` field as well (and only
+      # when the wrapped key actually has a private component).
+      #
+      # @param include_private [Boolean] include the `priv` field. Default: false.
+      # @return [Hash{Symbol=>String}] a JWK with `:kty`, `:alg`, `:pub`,
+      #   `:kid`, and optionally `:priv`.
       def export(include_private: false)
         jwk = {
           kty: KTY,
@@ -39,16 +66,31 @@ module JWT
       end
 
       # Import a Key from a JWK hash.
+      #
+      # Accepts string or symbol keys. Validates `kty`, `alg`, and the
+      # presence/base64url-ness of `pub` (and `priv` if present).
+      #
+      # @param jwk_hash [Hash] a JWK object.
+      # @return [JWT::PQ::Key] a key reconstructed from the JWK — with a
+      #   private component iff the JWK carried a `priv` field.
+      # @raise [KeyError] on missing/wrong `kty`, missing/unsupported `alg`,
+      #   missing `pub`, wrong field types, or invalid base64url in
+      #   `pub`/`priv`.
       def self.import(jwk_hash)
+        raise KeyError, "Expected a Hash for JWK, got #{jwk_hash.class}" unless jwk_hash.is_a?(Hash)
+
         jwk = normalize_keys(jwk_hash)
 
         validate_kty!(jwk)
         alg = validate_alg!(jwk)
         raise KeyError, "Missing 'pub' in JWK" unless jwk.key?("pub")
+        raise KeyError, "'pub' must be a String, got #{jwk["pub"].class}" unless jwk["pub"].is_a?(String)
 
         pub_bytes = decode_field(jwk, "pub")
 
         if jwk.key?("priv")
+          raise KeyError, "'priv' must be a String, got #{jwk["priv"].class}" unless jwk["priv"].is_a?(String)
+
           priv_bytes = decode_field(jwk, "priv")
           Key.new(algorithm: alg, public_key: pub_bytes, private_key: priv_bytes)
         else
@@ -56,14 +98,38 @@ module JWT
         end
       end
 
-      # JWK Thumbprint (RFC 7638) for key identification.
-      # Uses the required members: alg, kty, pub.
+      # Compute the JWK Thumbprint (RFC 7638) used as `kid`.
+      #
+      # Delegates to {JWT::PQ::Key#jwk_thumbprint}, which memoizes the
+      # result on the key — repeated calls on the same key avoid
+      # recomputing the canonical JSON + SHA-256 digest.
+      #
+      # @return [String] base64url-encoded SHA-256 thumbprint.
       def thumbprint
-        canonical = "{\"alg\":\"#{@key.algorithm}\",\"kty\":\"#{KTY}\",\"pub\":\"#{base64url_encode(@key.public_key)}\"}"
+        @key.jwk_thumbprint
+      end
+
+      # Compute an RFC 7638 thumbprint from algorithm + public key bytes
+      # without allocating a {JWK} or {JWT::PQ::Key} wrapper.
+      #
+      # @api private
+      # @param algorithm [String] canonical algorithm name.
+      # @param public_key [String] raw public key bytes.
+      # @return [String] base64url-encoded SHA-256 thumbprint.
+      def self.compute_thumbprint(algorithm, public_key)
+        # RFC 7638 §3.2: canonical JSON over the required members in
+        # lexicographic order (alg, kty, pub), no whitespace. Using
+        # `JSON.generate` over an ordered Hash instead of string
+        # interpolation so a future algorithm or key-byte change that
+        # introduces a character needing JSON escape does not silently
+        # produce a divergent thumbprint.
+        pub_b64 = ::Base64.urlsafe_encode64(public_key, padding: false)
+        canonical = JSON.generate({ alg: algorithm, kty: KTY, pub: pub_b64 })
         digest = OpenSSL::Digest::SHA256.digest(canonical)
-        base64url_encode(digest)
+        ::Base64.urlsafe_encode64(digest, padding: false)
       end
 
+      # @api private
       def self.validate_kty!(jwk)
         kty = jwk["kty"]
         raise KeyError, "Missing 'kty' in JWK" unless kty
@@ -71,6 +137,7 @@ module JWT
       end
       private_class_method :validate_kty!
 
+      # @api private
       def self.validate_alg!(jwk)
         alg = jwk["alg"]
         raise KeyError, "Missing 'alg' in JWK" unless alg
@@ -80,11 +147,13 @@ module JWT
       end
       private_class_method :validate_alg!
 
+      # @api private
       def self.normalize_keys(hash)
         hash.transform_keys(&:to_s)
       end
       private_class_method :normalize_keys
 
+      # @api private
       def self.decode_field(jwk, field)
         ::Base64.urlsafe_decode64(jwk[field])
       rescue ArgumentError => e

data/lib/jwt/pq/jwk_set.rb

@@ -0,0 +1,196 @@
+# frozen_string_literal: true
+
+require "json"
+
+module JWT
+  module PQ
+    # A set of JWKs (RFC 7517 §5) for publication and `kid`-based lookup.
+    #
+    # Typical producer flow — publish verification keys on
+    # `/.well-known/jwks.json`:
+    #
+    # @example Publish a JWKS
+    #   jwks = JWT::PQ::JWKSet.new([key_current, key_next])
+    #   File.write("jwks.json", jwks.to_json)
+    #
+    # Typical consumer flow — pick the right key for an incoming JWT by the
+    # `kid` header:
+    #
+    # @example Resolve a verification key by kid
+    #   jwks = JWT::PQ::JWKSet.import(JSON.parse(fetch_jwks))
+    #   _payload, header = JWT.decode(token, nil, false) # unverified peek
+    #   key = jwks[header["kid"]] or raise "unknown kid"
+    #   payload, = JWT.decode(token, key, true, algorithms: [header["alg"]])
+    #
+    # The set indexes members by their RFC 7638 JWK Thumbprint, which is the
+    # `kid` {JWK#export} emits. If you import a JWKS that uses custom (non-
+    # thumbprint) `kid` values, lookup by those custom values is not
+    # supported — rely on thumbprints when generating the set.
+    #
+    # @see https://www.rfc-editor.org/rfc/rfc7517#section-5 RFC 7517 §5
+    class JWKSet
+      include Enumerable
+
+      # Build a set from zero or more {JWT::PQ::Key}s.
+      #
+      # @param keys [Array<JWT::PQ::Key>, JWT::PQ::Key, nil] initial members.
+      # @raise [KeyError] if any element is not a {JWT::PQ::Key}.
+      def initialize(keys = [])
+        @keys = []
+        @kid_index = {}
+        Array(keys).each { |k| add(k) }
+      end
+
+      # Add a key to the set.
+      #
+      # Idempotent: if a key with the same RFC 7638 thumbprint is already
+      # in the set, the call is a no-op (Set semantics). The thumbprint
+      # is computed before any mutation, so a failure to derive the
+      # `kid` leaves the set unchanged.
+      #
+      # @param key [JWT::PQ::Key] the key to add.
+      # @return [JWKSet] self, for chaining.
+      # @raise [KeyError] if `key` is not a {JWT::PQ::Key}.
+      def add(key)
+        raise KeyError, "Expected a JWT::PQ::Key, got #{key.class}" unless key.is_a?(JWT::PQ::Key)
+
+        kid = key.jwk_thumbprint
+        return self if @kid_index.key?(kid)
+
+        @keys << key
+        @kid_index[kid] = key
+        self
+      end
+
+      # Iterate over the keys in insertion order.
+      #
+      # @yieldparam key [JWT::PQ::Key]
+      # @return [Enumerator] when called without a block.
+      def each(&)
+        @keys.each(&)
+      end
+
+      # @return [Integer] number of keys in the set.
+      def size
+        @keys.size
+      end
+      alias length size
+
+      # @return [Boolean] true if the set is empty.
+      def empty?
+        @keys.empty?
+      end
+
+      # Look up a key by its RFC 7638 thumbprint (the `kid` from {JWK#export}).
+      #
+      # @param kid [String] the thumbprint to match.
+      # @return [JWT::PQ::Key, nil] the matching key, or nil if not in the set.
+      def find(kid)
+        @kid_index[kid]
+      end
+      alias [] find
+
+      # @return [Array<JWT::PQ::Key>] a frozen snapshot of the keys in the set.
+      def keys
+        @keys.dup.freeze
+      end
+
+      # Export the set as a JWKS hash.
+      #
+      # @param include_private [Boolean] include the `priv` field on each
+      #   member that has a private component. Default: false.
+      # @return [Hash{Symbol=>Array<Hash>}] a hash with a single `:keys`
+      #   member, suitable for serialization with {#to_json}.
+      def export(include_private: false)
+        { keys: @keys.map { |k| JWK.new(k).export(include_private: include_private) } }
+      end
+
+      # Serialize the set as a JWKS JSON document.
+      #
+      # Always emits public-only keys — the `priv` field is never
+      # written out. This keeps the method safe for arbitrary nesting
+      # inside other JSON (e.g. `{ jwks: set }.to_json`), where Ruby's
+      # stdlib JSON passes a generator state as a positional argument.
+      # To publish private material (unusual), call
+      # `JSON.generate(set.export(include_private: true))` explicitly.
+      #
+      # @return [String] a JSON document ready for `/.well-known/jwks.json`.
+      def to_json(*)
+        export.to_json(*)
+      end
+
+      # @return [String] short diagnostic string — never contains key material.
+      def inspect
+        "#<#{self.class} size=#{size}>"
+      end
+      alias to_s inspect
+
+      # Import a JWKS from a Hash or JSON string.
+      #
+      # Each member is reconstructed via {JWT::PQ::JWK.import}; malformed
+      # members raise {KeyError}.
+      #
+      # ML-DSA public keys are ~1.3–2.6 KB each, so a JWKS with N keys is
+      # at least N × ~2 KB. When ingesting untrusted JWKS payloads (e.g.
+      # a remote `/.well-known/jwks.json`), bound the HTTP body size
+      # before calling `import` — this method does not cap the number of
+      # members.
+      #
+      # @param source [Hash, String] a JWKS hash or JSON string with a
+      #   `"keys"` array.
+      # @return [JWKSet] a new set with all parsed members.
+      # @raise [KeyError] if `source` is not a Hash/String, if the `keys`
+      #   field is missing or not an Array, or if any member fails to import.
+      def self.import(source)
+        hash = coerce_to_hash(source)
+        raise KeyError, "Expected Hash for JWKS body, got #{hash.class}" unless hash.is_a?(Hash)
+
+        hash = hash.transform_keys(&:to_s)
+        raise KeyError, "Missing 'keys' in JWKS" unless hash.key?("keys")
+        raise KeyError, "'keys' must be an Array" unless hash["keys"].is_a?(Array)
+
+        members = hash["keys"].map { |jwk| JWT::PQ::JWK.import(jwk) }
+        new(members)
+      end
+
+      # @api private
+      def self.coerce_to_hash(source)
+        case source
+        when String
+          begin
+            ::JSON.parse(source)
+          rescue ::JSON::ParserError => e
+            raise KeyError, "Invalid JSON for JWKS: #{e.message}"
+          end
+        when Hash then source
+        else raise KeyError, "Expected Hash or JSON String for JWKS, got #{source.class}"
+        end
+      end
+      private_class_method :coerce_to_hash
+
+      # Fetch a JWKS from a URL, honouring the process-global cache.
+      #
+      # Convenience wrapper around {Loader#fetch} using
+      # {Loader.default} — the cache is shared across all callers, so
+      # repeated hits on the same URL within `cache_ttl` seconds return
+      # the in-memory set without touching the network.
+      #
+      # See {Loader} for the full option reference (cache TTL, timeouts,
+      # body-size cap, HTTPS enforcement, ETag-based revalidation).
+      #
+      # @example Verify a token using a remote JWKS
+      #   jwks = JWT::PQ::JWKSet.fetch("https://issuer.example/.well-known/jwks.json")
+      #   _payload, header = JWT.decode(token, nil, false)
+      #   key = jwks[header["kid"]] or raise "unknown kid"
+      #   payload, = JWT.decode(token, key, true, algorithms: [header["alg"]])
+      #
+      # @param url [String] absolute JWKS URL.
+      # @return [JWKSet] the parsed set of verification keys.
+      # @raise [JWKSFetchError] on fetch failure (see {Loader#fetch}).
+      # @raise [KeyError] if the fetched body is not a valid JWKS.
+      def self.fetch(url, **)
+        Loader.default.fetch(url, **)
+      end
+    end
+  end
+end