Merge pull request #321 from mauromol/add-message-xml-post-processing

Allow for extension classes to post-process generated XML

Author: pitbulk

core/src/main/java/com/onelogin/saml2/authn/AuthnRequest.java

@@ -11,8 +11,8 @@
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
-import com.onelogin.saml2.settings.Saml2Settings;
 import com.onelogin.saml2.model.Organization;
+import com.onelogin.saml2.settings.Saml2Settings;
 import com.onelogin.saml2.util.Constants;
 import com.onelogin.saml2.util.Util;
 
@@ -101,7 +101,7 @@ public AuthnRequest(Saml2Settings settings, boolean forceAuthn, boolean isPassiv
 		this.nameIdValueReq = nameIdValueReq;
 
 		StrSubstitutor substitutor = generateSubstitutor(settings);
-		authnRequestString = substitutor.replace(getAuthnRequestTemplate());
+		authnRequestString = postProcessXml(substitutor.replace(getAuthnRequestTemplate()), settings);
 		LOGGER.debug("AuthNRequest --> " + authnRequestString);
 	}
 
@@ -121,6 +121,26 @@ public AuthnRequest(Saml2Settings settings, boolean forceAuthn, boolean isPassiv
 		this(settings, forceAuthn, isPassive, setNameIdPolicy, null);
 	}
 
+	/**
+	 * Allows for an extension class to post-process the AuthnRequest XML generated
+	 * for this request, in order to customize the result.
+	 * <p>
+	 * This method is invoked at construction time, after all the other fields of
+	 * this class have already been initialised. Its default implementation simply
+	 * returns the input XML as-is, with no change.
+	 * 
+	 * @param authnRequestXml
+	 *              the XML produced for this AuthnRequest by the standard
+	 *              implementation provided by {@link AuthnRequest}
+	 * @param settings
+	 *              the settings
+	 * @return the post-processed XML for this AuthnRequest, which will then be
+	 *         returned by any call to {@link #getAuthnRequestXml()}
+	 */
+	protected String postProcessXml(final String authnRequestXml, final Saml2Settings settings) {
+		return authnRequestXml;
+	}
+	
 	/**
 	 * @return the base64 encoded unsigned AuthnRequest (deflated or not)
 	 *

core/src/main/java/com/onelogin/saml2/logout/LogoutRequest.java

@@ -139,7 +139,7 @@ public LogoutRequest(Saml2Settings settings, HttpRequest request, String nameId,
 			this.sessionIndex = sessionIndex;
 
 			StrSubstitutor substitutor = generateSubstitutor(settings);
-			logoutRequestString = substitutor.replace(getLogoutRequestTemplate());
+			logoutRequestString = postProcessXml(substitutor.replace(getLogoutRequestTemplate()), settings);
 		} else {
 			logoutRequestString = Util.base64decodedInflated(samlLogoutRequest);
 			Document doc = Util.loadXML(logoutRequestString);
@@ -224,6 +224,28 @@ public LogoutRequest(Saml2Settings settings, HttpRequest request) {
 		this(settings, request, null, null);
 	}
 
+	/**
+	 * Allows for an extension class to post-process the LogoutRequest XML generated
+	 * for this request, in order to customize the result.
+	 * <p>
+	 * This method is invoked at construction time when no existing LogoutRequest
+	 * message is found in the HTTP request (and hence in the logout request sending
+	 * scenario only), after all the other fields of this class have already been
+	 * initialised. Its default implementation simply returns the input XML as-is,
+	 * with no change.
+	 * 
+	 * @param logoutRequestXml
+	 *              the XML produced for this LogoutRequest by the standard
+	 *              implementation provided by {@link LogoutRequest}
+	 * @param settings
+	 *              the settings
+	 * @return the post-processed XML for this LogoutRequest, which will then be
+	 *         returned by any call to {@link #getLogoutRequestXml()}
+	 */
+	protected String postProcessXml(final String logoutRequestXml, final Saml2Settings settings) {
+		return logoutRequestXml;
+	}
+
 	/**
 	 * @return the base64 encoded unsigned Logout Request (deflated or not)
 	 *

core/src/main/java/com/onelogin/saml2/logout/LogoutResponse.java

@@ -367,7 +367,7 @@ public void build(String inResponseTo, String statusCode) {
 		this.inResponseTo = inResponseTo;
 
 		StrSubstitutor substitutor = generateSubstitutor(settings, statusCode);
-		this.logoutResponseString = substitutor.replace(getLogoutResponseTemplate());
+		this.logoutResponseString = postProcessXml(substitutor.replace(getLogoutResponseTemplate()), settings);
 	}
 
     /**
@@ -388,6 +388,26 @@ public void build() {
 		build(null);
 	}	
 
+	/**
+	 * Allows for an extension class to post-process the LogoutResponse XML
+	 * generated for this response, in order to customize the result.
+	 * <p>
+	 * This method is invoked by {@link #build(String, String)} (and all of its
+	 * overloadings) and hence only in the logout response sending scenario. Its
+	 * default implementation simply returns the input XML as-is, with no change.
+	 * 
+	 * @param logoutResponseXml
+	 *              the XML produced for this LogoutResponse by the standard
+	 *              implementation provided by {@link LogoutResponse}
+	 * @param settings
+	 *              the settings
+	 * @return the post-processed XML for this LogoutResponse, which will then be
+	 *         returned by any call to {@link #getLogoutResponseXml()}
+	 */
+	protected String postProcessXml(final String logoutResponseXml, final Saml2Settings settings) {
+		return logoutResponseXml;
+	}
+	
 	/**
 	 * Substitutes LogoutResponse variables within a string by values.
 	 *

core/src/main/java/com/onelogin/saml2/settings/Metadata.java

@@ -77,7 +77,7 @@ public Metadata(Saml2Settings settings, Calendar validUntilTime, Integer cacheDu
 		this.cacheDuration = cacheDuration;
 
 		StrSubstitutor substitutor = generateSubstitutor(settings);
-		String unsignedMetadataString = substitutor.replace(getMetadataTemplate());
+		String unsignedMetadataString = postProcessXml(substitutor.replace(getMetadataTemplate()), settings);
 
 		LOGGER.debug("metadata --> " + unsignedMetadataString);
 		metadataString = unsignedMetadataString;
@@ -109,11 +109,31 @@ public Metadata(Saml2Settings settings) throws CertificateEncodingException {
 		this.cacheDuration = SECONDS_CACHED;
 
 		StrSubstitutor substitutor = generateSubstitutor(settings);
-		String unsignedMetadataString = substitutor.replace(getMetadataTemplate());
+		String unsignedMetadataString = postProcessXml(substitutor.replace(getMetadataTemplate()), settings);
 
 		LOGGER.debug("metadata --> " + unsignedMetadataString);
 		metadataString = unsignedMetadataString;
 	}
+	
+	/**
+	 * Allows for an extension class to post-process the SAML metadata XML generated
+	 * for this metadata instance, in order to customize the result.
+	 * <p>
+	 * This method is invoked at construction time, after all the other fields of
+	 * this class have already been initialised. Its default implementation simply
+	 * returns the input XML as-is, with no change.
+	 * 
+	 * @param metadataXml
+	 *              the XML produced for this metadata instance by the standard
+	 *              implementation provided by {@link Metadata}
+	 * @param settings
+	 *              the settings
+	 * @return the post-processed XML for this metadata instance, which will then be
+	 *         returned by any call to {@link #getMetadataString()}
+	 */
+	protected String postProcessXml(final String metadataXml, final Saml2Settings settings) {
+		return metadataXml;
+	}
 
 	/**
 	 * Substitutes metadata variables within a string by values.

core/src/test/java/com/onelogin/saml2/test/authn/AuthnRequestTest.java

@@ -5,6 +5,7 @@
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertNotEquals;
 import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertSame;
 import static org.junit.Assert.assertThat;
 import static org.junit.Assert.assertTrue;
 
@@ -382,4 +383,25 @@ public void testAuthNDestination() throws Exception {
 		assertThat(authnRequestStr, containsString("<samlp:AuthnRequest"));
 		assertThat(authnRequestStr, not(containsString("Destination=\"http://idp.example.com/simplesaml/saml2/idp/SSOService.php\"")));
 	}
+	
+	/**
+	 * Tests the postProcessXml method of AuthnRequest
+	 *
+	 * @throws Exception
+	 * 
+	 * @see com.onelogin.saml2.authn.AuthnRequest#postProcessXml
+	 */
+	@Test
+	public void testPostProcessXml() throws Exception {
+		Saml2Settings settings = new SettingsBuilder().fromFile("config/config.min.properties").build();
+		AuthnRequest authnRequest = new AuthnRequest(settings) {
+			@Override
+			protected String postProcessXml(String authRequestXml, Saml2Settings sett) {
+				assertEquals(authRequestXml, super.postProcessXml(authRequestXml, sett));
+				assertSame(settings, sett);
+				return "changed";
+			}
+		};
+		assertEquals("changed", authnRequest.getAuthnRequestXml());
+	}
 }

core/src/test/java/com/onelogin/saml2/test/logout/LogoutRequestTest.java

@@ -5,6 +5,7 @@
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertThat;
 import static org.junit.Assert.assertNull;
+import static org.junit.Assert.assertSame;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNotEquals;
 import static org.junit.Assert.assertNotNull;
@@ -1008,4 +1009,25 @@ public void testGetError() throws Exception {
 	private static HttpRequest newHttpRequest(String requestURL, String samlRequestEncoded) {
 		return new HttpRequest(requestURL, (String)null).addParameter("SAMLRequest", samlRequestEncoded);
 	}
+
+	/**
+	 * Tests the postProcessXml method of LogoutRequest
+	 *
+	 * @throws Exception
+	 * 
+	 * @see com.onelogin.saml2.logout.LogoutRequest#postProcessXml
+	 */
+	@Test
+	public void testPostProcessXml() throws Exception {
+		Saml2Settings settings = new SettingsBuilder().fromFile("config/config.min.properties").build();
+		LogoutRequest logoutRequest = new LogoutRequest(settings) {
+			@Override
+			protected String postProcessXml(String authRequestXml, Saml2Settings sett) {
+				assertEquals(authRequestXml, super.postProcessXml(authRequestXml, sett));
+				assertSame(settings, sett);
+				return "changed";
+			}
+		};
+		assertEquals("changed", logoutRequest.getLogoutRequestXml());
+	}
 }

core/src/test/java/com/onelogin/saml2/test/logout/LogoutResponseTest.java

@@ -5,6 +5,7 @@
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertThat;
 import static org.junit.Assert.assertNull;
+import static org.junit.Assert.assertSame;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNotEquals;
 import static org.junit.Assert.assertNotNull;
@@ -717,4 +718,26 @@ public void testGetError() throws URISyntaxException, IOException, XMLEntityExce
 	private static HttpRequest newHttpRequest(String requestURL, String samlResponseEncoded) {
 		return new HttpRequest(requestURL, (String)null).addParameter("SAMLResponse", samlResponseEncoded);
 	}
+
+	/**
+	 * Tests the postProcessXml method of LogoutResponse
+	 *
+	 * @throws Exception
+	 * 
+	 * @see com.onelogin.saml2.logout.LogoutResponse#postProcessXml
+	 */
+	@Test
+	public void testPostProcessXml() throws Exception {
+		Saml2Settings settings = new SettingsBuilder().fromFile("config/config.min.properties").build();
+		LogoutResponse logoutResponse = new LogoutResponse(settings, null) {
+			@Override
+			protected String postProcessXml(String authRequestXml, Saml2Settings sett) {
+				assertEquals(authRequestXml, super.postProcessXml(authRequestXml, sett));
+				assertSame(settings, sett);
+				return "changed";
+			}
+		};
+		logoutResponse.build();
+		assertEquals("changed", logoutResponse.getLogoutResponseXml());
+	}
 }

core/src/test/java/com/onelogin/saml2/test/settings/MetadataTest.java

@@ -7,7 +7,7 @@
 import static org.junit.Assert.assertThat;
 import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.assertNull;
-
+import static org.junit.Assert.assertSame;
 
 import java.io.IOException;
 import java.security.GeneralSecurityException;
@@ -520,4 +520,25 @@ public void shouldIgnoreCacheDuration() throws CertificateEncodingException, Err
 		assertNull("should not set cache duration attribute", cacheDurationNode);
 
 	}
+	
+	/**
+	 * Tests the postProcessXml method of Metadata
+	 *
+	 * @throws Exception
+	 * 
+	 * @see com.onelogin.saml2.settings.Metadata#postProcessXml
+	 */
+	@Test
+	public void testPostProcessXml() throws Exception {
+		Saml2Settings settings = new SettingsBuilder().fromFile("config/config.min.properties").build();
+		Metadata metadata = new Metadata(settings) {
+			@Override
+			protected String postProcessXml(String authRequestXml, Saml2Settings sett) {
+				assertEquals(authRequestXml, super.postProcessXml(authRequestXml, sett));
+				assertSame(settings, sett);
+				return "changed";
+			}
+		};
+		assertEquals("changed", metadata.getMetadataString());
+	}
 }