Commit 747f6df

@@ -2,11 +2,15 @@
 
 xml.EncryptedData xmlns: ::Xml::Kit::Namespaces::XMLENC do
   xml.EncryptionMethod Algorithm: symmetric_algorithm
-  xml.KeyInfo xmlns: ::Xml::Kit::Namespaces::XMLDSIG do
-    xml.EncryptedKey xmlns: ::Xml::Kit::Namespaces::XMLENC do
-      xml.EncryptionMethod Algorithm: asymmetric_algorithm
-      xml.CipherData do
-        xml.CipherValue asymmetric_cipher_value
+  if key_info
+    key_info.call(xml)
+  else
+    xml.KeyInfo xmlns: ::Xml::Kit::Namespaces::XMLDSIG do
+      xml.EncryptedKey xmlns: ::Xml::Kit::Namespaces::XMLENC do
+        xml.EncryptionMethod Algorithm: asymmetric_algorithm
+        xml.CipherData do
+          xml.CipherValue asymmetric_cipher_value
+        end
       end
     end
   end

@@ -7,12 +7,14 @@ module Xml
       attr_reader :asymmetric_cipher_value
       attr_reader :symmetric_algorithm
       attr_reader :symmetric_cipher_value
+      attr_reader :key_info
 
       def initialize(
         raw_xml,
         public_key,
         symmetric_algorithm: ::Xml::Kit::Crypto::SymmetricCipher::DEFAULT_ALGORITHM,
-        asymmetric_algorithm: ::Xml::Kit::Crypto::RsaCipher::ALGORITHM
+        asymmetric_algorithm: ::Xml::Kit::Crypto::RsaCipher::ALGORITHM,
+        key_info: nil
       )
         @symmetric_algorithm = symmetric_algorithm
         symmetric_cipher = symmetric(symmetric_algorithm)
@@ -21,6 +23,7 @@ module Xml
         @asymmetric_algorithm = asymmetric_algorithm
         asymmetric_cipher = asymmetric(asymmetric_algorithm, public_key)
         @asymmetric_cipher_value = Base64.strict_encode64(asymmetric_cipher.encrypt(symmetric_cipher.key))
+        @key_info = key_info
       end
 
       def to_xml(xml: ::Builder::XmlMarkup.new)

@@ -18,17 +18,19 @@ module Xml
       attr_accessor :encryption_certificate
 
       # Returns the generated XML document with an XML Digital Signature and XML Encryption.
-      def to_xml(xml: ::Builder::XmlMarkup.new)
-        signatures.complete(render(self, xml: xml))
+      def to_xml(xml: ::Builder::XmlMarkup.new, pretty: false)
+        result = signatures.complete(render(self, xml: xml))
+        pretty ? Nokogiri::XML(result).to_xml(indent: 2) : result
       end
 
-      def encryption_for(xml:)
+      def encryption_for(xml:, key_info: nil)
         if encrypt?
           temp = ::Builder::XmlMarkup.new
           yield temp
           ::Xml::Kit::Encryption.new(
             signatures.complete(temp.target!),
-            encryption_certificate.public_key
+            encryption_certificate.public_key,
+            key_info: key_info
           ).to_xml(xml: xml)
         else
           yield xml

@@ -2,7 +2,12 @@ xml.instruct!
 xml.Item ID: id, xmlns: 'https://www.example.org/item#' do
   signature_for reference_id: id, xml: xml
   xml.Encrypted xmlns: 'https://www.example.org/item#' do
-    encryption_for xml: xml do |xml|
+    key_info = lambda do |xml|
+      xml.KeyInfo xmlns: ::Xml::Kit::Namespaces::XMLDSIG do
+        xml.RetrievalMethod xmlns: ::Xml::Kit::Namespaces::XMLDSIG, URI: "#EK", Type: "http://www.w3.org/2001/04/xmlenc#EncryptedKey"
+      end
+    end
+    encryption_for(xml: xml, key_info: key_info) do |xml|
       xml.EncryptMe do
         xml.Secret "secret"
       end

@@ -97,11 +97,18 @@ RSpec.describe ::Xml::Kit::Templatable do
     specify { expect(subject.to_xml).to match_xsd('item') }
 
     context "with the key extracted to the header" do
+      let(:xml_hash) { Hash.from_xml(subject.to_xml) }
+
       before do
         subject.template_path = './spec/fixtures/item-extracted-key.builder'
       end
 
       specify { expect(subject.to_xml).to match_xsd('item-extracted-key') }
+      specify { expect(xml_hash["Item"]["Encrypted"]["EncryptedData"]["KeyInfo"]["RetrievalMethod"]).to be_present }
+      specify { expect(xml_hash["Item"]["Encrypted"]["EncryptedData"]["KeyInfo"]["RetrievalMethod"]["xmlns"]).to eql(::Xml::Kit::Namespaces::XMLDSIG) }
+      specify { expect(xml_hash["Item"]["Encrypted"]["EncryptedData"]["KeyInfo"]["RetrievalMethod"]["URI"]).to eql("#EK") }
+      specify { expect(xml_hash["Item"]["Encrypted"]["EncryptedData"]["KeyInfo"]["RetrievalMethod"]["Type"]).to eql("http://www.w3.org/2001/04/xmlenc#EncryptedKey") }
+      specify { expect(xml_hash["Item"]["Encrypted"]["EncryptedData"]["KeyInfo"]["EncryptedKey"]).to be_nil }
     end
   end
 end