Add support for dynamic key detection during decryption

README.md

@@ -73,6 +73,27 @@ plaintext = JWE.decrypt(encrypted, key)
 puts plaintext #"The quick brown fox jumps over the lazy dog."
 ```
 
+This example sets an extra header and then uses it to detect the key during decryption.
+
+```ruby
+require 'jwe'
+
+keys = {
+  'id-1' => OpenSSL::PKey::RSA.generate(2048),
+  'id-2' => OpenSSL::PKey::RSA.generate(2048)
+}
+payload = "The quick brown fox jumps over the lazy dog."
+
+encrypted = JWE.encrypt(payload, keys['id-2'], headers: {kid: 'id-2'})
+puts encrypted
+
+plaintext = JWE.decrypt(encrypted, nil) do |headers|
+  kid = headers['kid']
+  keys[kid]
+end
+puts plaintext #"The quick brown fox jumps over the lazy dog."
+```
+
 ## Available Algorithms
 
 The RFC 7518 JSON Web Algorithms (JWA) spec defines the algorithms for [encryption](https://tools.ietf.org/html/rfc7518#section-5.1)

lib/jwe.rb

@@ -19,14 +19,15 @@ class InvalidData < RuntimeError; end
   VALID_ENC = ['A128CBC-HS256', 'A192CBC-HS384', 'A256CBC-HS512', 'A128GCM', 'A192GCM', 'A256GCM'].freeze
   VALID_ZIP = ['DEF'].freeze
 
-  def self.encrypt(payload, key, alg: 'RSA-OAEP', enc: 'A128GCM', zip: nil)
+  def self.encrypt(payload, key, alg: 'RSA-OAEP', enc: 'A128GCM', zip: nil, headers: {})
     raise ArgumentError.new("\"#{alg}\" is not a valid alg method") unless VALID_ALG.include?(alg)
     raise ArgumentError.new("\"#{enc}\" is not a valid enc method") unless VALID_ENC.include?(enc)
     raise ArgumentError.new("\"#{zip}\" is not a valid zip method") unless zip.nil? || zip == '' || VALID_ZIP.include?(zip)
     raise ArgumentError.new('The key must not be nil or blank') if key.nil? || (key.is_a?(String) && key.strip == '')
 
     header = { alg: alg, enc: enc }
     header[:zip] = zip if zip && zip != ''
+    header.merge!(headers) if headers.is_a?(Hash)
 
     cipher = Enc.for(enc).new
     cipher.cek = key if alg == 'dir'
@@ -39,14 +40,16 @@ def self.encrypt(payload, key, alg: 'RSA-OAEP', enc: 'A128GCM', zip: nil)
     Serialization::Compact.encode(header.to_json, encrypted_cek, cipher.iv, ciphertext, cipher.tag)
   end
 
-  def self.decrypt(payload, key)
+  def self.decrypt(payload, key, &keyfinder)
     header, enc_key, iv, ciphertext, tag = Serialization::Compact.decode(payload)
     header = JSON.parse(header)
     base64header = payload.split('.').first
 
     raise ArgumentError.new("\"#{header['alg']}\" is not a valid alg method") unless VALID_ALG.include?(header['alg'])
     raise ArgumentError.new("\"#{header['enc']}\" is not a valid enc method") unless VALID_ENC.include?(header['enc'])
     raise ArgumentError.new("\"#{header['zip']}\" is not a valid zip method") unless header['zip'].nil? || VALID_ZIP.include?(header['zip'])
+
+    key = yield(header) if keyfinder
     raise ArgumentError.new('The key must not be nil or blank') if key.nil? || (key.is_a?(String) && key.strip == '')
 
     cek = Alg.for(header['alg']).new(key).decrypt(enc_key)

spec/jwe_spec.rb

@@ -28,6 +28,27 @@
     end
   end
 
+  describe 'when using extra headers' do
+    it 'roundtrips' do
+      encrypted = JWE.encrypt(plaintext, rsa_key, headers: {kid: 'some-kid-1'})
+      result = JWE.decrypt(encrypted, rsa_key)
+      header, _ = JWE::Serialization::Compact.decode(encrypted)
+      header = JSON.parse(header)
+
+      expect(header['kid']).to eq 'some-kid-1'
+      expect(result).to eq plaintext
+    end
+
+    it 'allows dynamic key detection' do
+      encrypted = JWE.encrypt(plaintext, rsa_key)
+      result = JWE.decrypt(encrypted, nil) do |header|
+        rsa_key
+      end
+
+      expect(result).to eq plaintext
+    end
+  end
+
   it 'raises when passed a bad alg' do
     expect { JWE.encrypt(plaintext, rsa_key, alg: 'TEST') }.to raise_error(ArgumentError)
   end