From dc359e56d4801cf714a3845113dd357a38f6e9b3 Mon Sep 17 00:00:00 2001
From: Luca Leonardo Scorcia <l.scorcia@gmail.com>
Date: Thu, 4 Mar 2021 08:47:06 -0500
Subject: [PATCH] KEYCLOAK-17329 Allow emitting custom elements in SAML
 metadata extensions

---
 .../dom/saml/v2/metadata/ExtensionsType.java  | 12 +++++++++
 .../saml/v2/writers/SAMLMetadataWriter.java   | 26 +++++++++++++------
 2 files changed, 30 insertions(+), 8 deletions(-)

diff --git a/saml-core-api/src/main/java/org/keycloak/dom/saml/v2/metadata/ExtensionsType.java b/saml-core-api/src/main/java/org/keycloak/dom/saml/v2/metadata/ExtensionsType.java
index 38b08c131c6..8eb9e4ec2cb 100755
--- a/saml-core-api/src/main/java/org/keycloak/dom/saml/v2/metadata/ExtensionsType.java
+++ b/saml-core-api/src/main/java/org/keycloak/dom/saml/v2/metadata/ExtensionsType.java
@@ -90,6 +90,18 @@ public class ExtensionsType {
         return Collections.unmodifiableList(this.any);
     }
 
+    public List<Element> getDomElements() {
+        List<Element> output = new ArrayList<Element>();
+
+        for (Object o : this.any) {
+            if (o instanceof Element) {
+                output.add((Element) o);
+            }
+        }
+
+        return Collections.unmodifiableList(output);
+    }
+
     public EntityAttributes getEntityAttributes() {
         for (Object o : this.any) {
             if (o instanceof EntityAttributes) {
diff --git a/saml-core/src/main/java/org/keycloak/saml/processing/core/saml/v2/writers/SAMLMetadataWriter.java b/saml-core/src/main/java/org/keycloak/saml/processing/core/saml/v2/writers/SAMLMetadataWriter.java
index 072703b7ea5..21203e39d45 100755
--- a/saml-core/src/main/java/org/keycloak/saml/processing/core/saml/v2/writers/SAMLMetadataWriter.java
+++ b/saml-core/src/main/java/org/keycloak/saml/processing/core/saml/v2/writers/SAMLMetadataWriter.java
@@ -90,7 +90,7 @@ public class SAMLMetadataWriter extends BaseWriter {
         }
         ExtensionsType extensions = entities.getExtensions();
         if (extensions != null) {
-            StaxUtil.writeDOMElement(writer, extensions.getElement());
+            write(extensions);
         }
 
         List<Object> entityDescriptors = entities.getEntityDescriptor();
@@ -126,7 +126,7 @@ public class SAMLMetadataWriter extends BaseWriter {
         }
         ExtensionsType extensions = entityDescriptor.getExtensions();
         if (extensions != null) {
-            StaxUtil.writeDOMElement(writer, extensions.getElement());
+            write(extensions);
         }
 
         List<EntityDescriptorType.EDTChoiceType> choiceTypes = entityDescriptor.getChoiceType();
@@ -297,7 +297,7 @@ public class SAMLMetadataWriter extends BaseWriter {
         }
         ExtensionsType extensions = attributeAuthority.getExtensions();
         if (extensions != null) {
-            StaxUtil.writeDOMElement(writer, extensions.getElement());
+            write(extensions);
         }
 
         List<KeyDescriptorType> keyDescriptorList = attributeAuthority.getKeyDescriptor();
@@ -392,7 +392,7 @@ public class SAMLMetadataWriter extends BaseWriter {
 
         ExtensionsType extensions = org.getExtensions();
         if (extensions != null) {
-            StaxUtil.writeDOMElement(writer, extensions.getElement());
+            write(extensions);
         }
 
         // Write the name
@@ -434,13 +434,14 @@ public class SAMLMetadataWriter extends BaseWriter {
     public void write(ContactType contact) throws ProcessingException {
         StaxUtil.writeStartElement(writer, METADATA_PREFIX, JBossSAMLConstants.CONTACT_PERSON.get(), JBossSAMLURIConstants.METADATA_NSURI.get());
 
-        ExtensionsType extensions = contact.getExtensions();
-        if (extensions != null) {
-            StaxUtil.writeDOMElement(writer, extensions.getElement());
-        }
         ContactTypeType attribs = contact.getContactType();
         StaxUtil.writeAttribute(writer, JBossSAMLConstants.CONTACT_TYPE.get(), attribs.value());
 
+        ExtensionsType extensions = contact.getExtensions();
+        if (extensions != null) {
+            write(extensions);
+        }
+
         // Write the name
         String company = contact.getCompany();
         if (company != null) {
@@ -480,6 +481,15 @@ public class SAMLMetadataWriter extends BaseWriter {
         StaxUtil.flush(writer);
     }
 
+    public void write(ExtensionsType extensions) throws ProcessingException {
+        StaxUtil.writeStartElement(writer, METADATA_PREFIX, JBossSAMLConstants.EXTENSIONS__METADATA.get(), JBossSAMLURIConstants.METADATA_NSURI.get());
+
+        for (Element extension : extensions.getDomElements())
+            StaxUtil.writeDOMElement(writer, extension);
+
+        StaxUtil.writeEndElement(writer);
+    }
+
     public void writeKeyDescriptor(KeyDescriptorType keyDescriptor) throws ProcessingException {
         StaxUtil.writeStartElement(writer, METADATA_PREFIX, JBossSAMLConstants.KEY_DESCRIPTOR.get(), JBossSAMLURIConstants.METADATA_NSURI.get());