updated samlengine

1 files changed, 17 insertions, 6 deletions

diff --git a/id/server/stork2-saml-engine/src/main/java/eu/stork/peps/auth/engine/SAMLEngine.java b/id/server/stork2-saml-engine/src/main/java/eu/stork/peps/auth/engine/SAMLEngine.java
index f8f50fc74..48718242b 100644
--- a/id/server/stork2-saml-engine/src/main/java/eu/stork/peps/auth/engine/SAMLEngine.java
+++ b/id/server/stork2-saml-engine/src/main/java/eu/stork/peps/auth/engine/SAMLEngine.java
@@ -50,8 +50,11 @@ import org.slf4j.LoggerFactory;
 import org.w3c.dom.Document;

 import org.w3c.dom.Element;

 

+import eu.stork.peps.auth.engine.core.CustomAttributeQuery;

 import eu.stork.peps.auth.engine.core.SAMLEngineSignI;

 import eu.stork.peps.auth.engine.core.STORKSAMLCore;

+import eu.stork.peps.auth.engine.core.impl.CustomAttributeQueryMarshaller;

+import eu.stork.peps.auth.engine.core.impl.CustomAttributeQueryUnmarshaller;

 import eu.stork.peps.auth.engine.core.impl.SignModuleFactory;

 import eu.stork.peps.configuration.ConfigurationCreator;

 import eu.stork.peps.configuration.ConfigurationReader;

@@ -237,7 +240,7 @@ public class SAMLEngine {
      * @return the SAML core properties

      */

     protected final STORKSAMLCore getSamlCoreProperties() {

-	return samlCore;

+    	return samlCore;

     }

 

     /**

@@ -259,7 +262,11 @@ public class SAMLEngine {
 	    final MarshallerFactory marshallerFactory = Configuration

 		    .getMarshallerFactory();

 

-	    final Marshaller marshaller = marshallerFactory

+	    final Marshaller marshaller;

+	    if (samlToken.getElementQName().toString().endsWith(CustomAttributeQuery.DEFAULT_ELEMENT_LOCAL_NAME))

+	    	marshaller = new CustomAttributeQueryMarshaller();

+	    else

+	    	marshaller = marshallerFactory

 		    .getMarshaller(samlToken);

 

 	    docBuilder = dbf.newDocumentBuilder();

@@ -326,9 +333,9 @@ public class SAMLEngine {
      */

     protected final byte[] signAndMarshall(final SignableSAMLObject samlToken)

 	    throws SAMLEngineException {

-	LOG.debug("Marshall Saml Token.");

-	SignableSAMLObject signElement = sign(samlToken);

-	return marshall(signElement);

+    	LOG.debug("Marshall Saml Token.");

+    	SignableSAMLObject signElement = sign(samlToken);

+    	return marshall(signElement);

     }

 

     /**

@@ -362,7 +369,11 @@ public class SAMLEngine {
             final UnmarshallerFactory unmarshallerFact = Configuration.getUnmarshallerFactory();

 	    // Unmarshall using the SAML Token root element

             if (unmarshallerFact != null && root != null){

-                final Unmarshaller unmarshaller = unmarshallerFact.getUnmarshaller(root);

+            	final Unmarshaller unmarshaller;

+            	if (root.getLocalName().equals(CustomAttributeQuery.DEFAULT_ELEMENT_LOCAL_NAME))

+            		unmarshaller = new CustomAttributeQueryUnmarshaller();

+            	else

+            		unmarshaller = unmarshallerFact.getUnmarshaller(root);

                 try {

 	    return unmarshaller.unmarshall(root);

                 } catch (NullPointerException e){