KEYCLOAK-9102 Add tests for Saml RelayState

saml-core/src/main/java/org/keycloak/saml/processing/api/saml/v2/request/SAML2Request.java

@@ -75,6 +75,21 @@ public class SAML2Request {
         this.nameIDFormat = nameIDFormat;
     }
 
+    /**
+     * Create authentication request with protocolBinding defaulting to POST
+     *
+     * @param id
+     * @param assertionConsumerURL
+     * @param destination
+     * @param issuerValue
+     * @return
+     * @throws ConfigurationException
+     */
+    public AuthnRequestType createAuthnRequestType(String id, String assertionConsumerURL, String destination,
+                                                   String issuerValue) throws ConfigurationException {
+        return createAuthnRequestType(id, assertionConsumerURL, destination, issuerValue, JBossSAMLURIConstants.SAML_HTTP_POST_BINDING.getUri());
+    }
+
     /**
      * Create an authentication request
      *
@@ -82,18 +97,19 @@ public class SAML2Request {
      * @param assertionConsumerURL
      * @param destination
      * @param issuerValue
+     * @param protocolBindingUri
      *
      * @return
      *
      * @throws ConfigurationException
      */
     public AuthnRequestType createAuthnRequestType(String id, String assertionConsumerURL, String destination,
-                                                   String issuerValue) throws ConfigurationException {
+                                                   String issuerValue, URI protocolBinding) throws ConfigurationException {
         XMLGregorianCalendar issueInstant = XMLTimeUtil.getIssueInstant();
 
         AuthnRequestType authnRequest = new AuthnRequestType(id, issueInstant);
         authnRequest.setAssertionConsumerServiceURL(URI.create(assertionConsumerURL));
-        authnRequest.setProtocolBinding(JBossSAMLURIConstants.SAML_HTTP_POST_BINDING.getUri());
+        authnRequest.setProtocolBinding(protocolBinding);
         if (destination != null) {
             authnRequest.setDestination(URI.create(destination));
         }

testsuite/integration-arquillian/tests/base/src/main/java/org/keycloak/testsuite/util/SamlClient.java

@@ -54,11 +54,13 @@ import java.security.PublicKey;
 import java.util.Arrays;
 import java.util.LinkedList;
 import java.util.List;
+import java.util.Optional;
 import java.util.UUID;
 
 import org.jboss.logging.Logger;
 import static org.hamcrest.Matchers.*;
 import static org.junit.Assert.assertThat;
+import static org.keycloak.saml.common.constants.GeneralConstants.RELAY_STATE;
 import static org.keycloak.testsuite.util.Matchers.statusCodeIsHC;
 
 /**
@@ -126,6 +128,14 @@ public class SamlClient {
                 return null;
             }
 
+            @Override
+            public String extractRelayState(CloseableHttpResponse response) throws IOException {
+                assertThat(response, statusCodeIsHC(Response.Status.OK));
+                String responsePage = EntityUtils.toString(response.getEntity(), "UTF-8");
+                response.close();
+                return extractSamlRelayStateFromForm(responsePage);
+            }
+
             @Override
             public HttpPost createSamlSignedRequest(URI samlEndpoint, String relayState, Document samlRequest, String realmPrivateKey, String realmPublicKey) {
                 return createSamlPostMessage(samlEndpoint, relayState, samlRequest, GeneralConstants.SAML_REQUEST_KEY, realmPrivateKey, realmPublicKey);
@@ -160,7 +170,7 @@ public class SamlClient {
                 }
 
                 if (relayState != null) {
-                    parameters.add(new BasicNameValuePair(GeneralConstants.RELAY_STATE, relayState));
+                    parameters.add(new BasicNameValuePair(RELAY_STATE, relayState));
                 }
 
                 UrlEncodedFormEntity formEntity;
@@ -245,6 +255,14 @@ public class SamlClient {
                 }
             }
 
+            @Override
+            public String extractRelayState(CloseableHttpResponse response) throws IOException {
+                assertThat(response, statusCodeIsHC(Response.Status.FOUND));
+                String location = response.getFirstHeader("Location").getValue();
+                response.close();
+                return extractRelayStateFromRedirect(location);
+            }
+
             @Override
             public HttpUriRequest createSamlSignedRequest(URI samlEndpoint, String relayState, Document samlRequest, String realmPrivateKey, String realmPublicKey) {
                 throw new UnsupportedOperationException("Not implemented yet.");
@@ -262,6 +280,8 @@ public class SamlClient {
         public abstract HttpUriRequest createSamlUnsignedResponse(URI samlEndpoint, String relayState, Document samlRequest);
 
         public abstract HttpUriRequest createSamlSignedResponse(URI samlEndpoint, String relayState, Document samlRequest, String realmPrivateKey, String realmPublicKey);
+
+        public abstract String extractRelayState(CloseableHttpResponse response) throws IOException;
     }
 
     private static final Logger LOG = Logger.getLogger(SamlClient.class);
@@ -288,6 +308,35 @@ public class SamlClient {
         return SAMLRequestParser.parseResponsePostBinding(respElement.val());
     }
 
+    /**
+     * Extracts and parses value of RelayState input field of a form present in the given page.
+     *
+     * @param responsePage HTML code of the page
+     * @return
+     */
+    public static String extractSamlRelayStateFromForm(String responsePage) {
+        assertThat(responsePage, containsString("form name=\"saml-post-binding\""));
+        org.jsoup.nodes.Document theResponsePage = Jsoup.parse(responsePage);
+        Elements samlRelayStates = theResponsePage.select("input[name=RelayState]");
+
+        if (samlRelayStates.isEmpty()) return null;
+
+        return samlRelayStates.first().val();
+    }
+
+    /**
+     * Extracts and parses value of RelayState query parameter from the given URI.
+     *
+     * @param responseUri
+     * @return
+     */
+    public static String extractRelayStateFromRedirect(String responseUri) {
+        List<NameValuePair> params = URLEncodedUtils.parse(URI.create(responseUri), "UTF-8");
+
+        return params.stream().filter(nameValuePair -> nameValuePair.getName().equals(RELAY_STATE))
+                .findFirst().map(NameValuePair::getValue).orElse(null);
+    }
+
     /**
      * Extracts and parses value of SAMLResponse query parameter from the given URI.
      *

testsuite/integration-arquillian/tests/base/src/main/java/org/keycloak/testsuite/util/SamlClientBuilder.java

@@ -183,6 +183,21 @@ public class SamlClientBuilder {
           .executeAndTransform(responseBinding::extractResponse);
     }
 
+    /** Returns RelayState from Saml response. Note that the redirects are disabled for this to work. */
+    public String getSamlRelayState(Binding responseBinding) {
+        return doNotFollowRedirects()
+                .executeAndTransform(responseBinding::extractRelayState);
+    }
+
+    /** Provide possibility to consume RelayState from saml response. Note that the redirects are disabled for this to work. */
+    public SamlClientBuilder assertSamlRelayState(Binding responseBinding, Consumer<String> relayStateConsumer) {
+        if (responseBinding.equals(Binding.REDIRECT)) doNotFollowRedirects();
+        return addStep((client, currentURI, currentResponse, context) -> {
+            relayStateConsumer.accept(responseBinding.extractRelayState(currentResponse));
+            return null;
+        });
+    }
+
     /** Returns SAML request or response as replied from server. Note that the redirects are disabled for this to work. */
     public ModifySamlResponseStepBuilder processSamlResponse(Binding responseBinding) {
         return

testsuite/integration-arquillian/tests/base/src/main/java/org/keycloak/testsuite/util/saml/CreateAuthnRequestStepBuilder.java

@@ -26,6 +26,8 @@ import org.keycloak.saml.processing.api.saml.v2.request.SAML2Request;
 import org.keycloak.testsuite.util.SamlClient.Binding;
 import java.net.URI;
 import java.util.UUID;
+import java.util.function.Supplier;
+
 import org.apache.http.client.methods.CloseableHttpResponse;
 import org.apache.http.client.methods.HttpUriRequest;
 import org.apache.http.client.protocol.HttpClientContext;
@@ -44,7 +46,7 @@ public class CreateAuthnRequestStepBuilder extends SamlDocumentStepBuilder<Authn
 
     private final Document forceLoginRequestDocument;
 
-    private String relayState;
+    private Supplier<String> relayState;
 
     public CreateAuthnRequestStepBuilder(URI authServerSamlUrl, String issuer, String assertionConsumerURL, Binding requestBinding, SamlClientBuilder clientBuilder) {
         super(clientBuilder);
@@ -67,16 +69,14 @@ public class CreateAuthnRequestStepBuilder extends SamlDocumentStepBuilder<Authn
         this.assertionConsumerURL = null;
     }
 
-    public String assertionConsumerURL() {
+    public CreateAuthnRequestStepBuilder relayState(Supplier<String> relayState) {
-        return assertionConsumerURL;
-    }
-
-    public String relayState() {
-        return relayState;
-    }
-
-    public void relayState(String relayState) {
         this.relayState = relayState;
+        return this;
+    }
+
+    public CreateAuthnRequestStepBuilder relayState(String relayState) {
+        this.relayState = () -> relayState;
+        return this;
     }
 
     public CreateAuthnRequestStepBuilder signWith(String signingPrivateKeyPem, String signingPublicKeyPem) {
@@ -97,6 +97,7 @@ public class CreateAuthnRequestStepBuilder extends SamlDocumentStepBuilder<Authn
         }
 
         Document samlDoc = DocumentUtil.getDocument(transformed);
+        String relayState = this.relayState == null ? null : this.relayState.get();
         return this.signingPrivateKeyPem == null
           ? requestBinding.createSamlUnsignedRequest(authServerSamlUrl, relayState, samlDoc)
           : requestBinding.createSamlSignedRequest(authServerSamlUrl, relayState, samlDoc, signingPrivateKeyPem, signingPublicKeyPem);
@@ -109,7 +110,8 @@ public class CreateAuthnRequestStepBuilder extends SamlDocumentStepBuilder<Authn
 
         try {
             SAML2Request samlReq = new SAML2Request();
-            AuthnRequestType loginReq = samlReq.createAuthnRequestType(UUID.randomUUID().toString(), assertionConsumerURL, this.authServerSamlUrl.toString(), issuer);
+            AuthnRequestType loginReq = samlReq.createAuthnRequestType(UUID.randomUUID().toString(),
+                    assertionConsumerURL, this.authServerSamlUrl.toString(), issuer, requestBinding.getBindingUri());
 
             return SAML2Request.convert(loginReq);
         } catch (ConfigurationException | ParsingException | ProcessingException ex) {

testsuite/integration-arquillian/tests/base/src/main/java/org/keycloak/testsuite/util/saml/IdPInitiatedLoginBuilder.java

@@ -16,15 +16,19 @@
  */
 package org.keycloak.testsuite.util.saml;
 
+import org.keycloak.saml.common.constants.GeneralConstants;
 import org.keycloak.testsuite.util.SamlClient.Step;
 import org.keycloak.testsuite.util.SamlClientBuilder;
 import java.net.URI;
+import java.util.function.Supplier;
+
 import org.apache.http.client.methods.CloseableHttpResponse;
 import org.apache.http.client.methods.HttpGet;
 import org.apache.http.client.methods.HttpUriRequest;
 import org.apache.http.client.protocol.HttpClientContext;
 import org.apache.http.impl.client.CloseableHttpClient;
 
+
 /**
  *
  * @author hmlnarik
@@ -34,6 +38,7 @@ public class IdPInitiatedLoginBuilder implements Step {
     private final SamlClientBuilder clientBuilder;
     private final URI authServerSamlUrl;
     private final String clientId;
+    private Supplier<String> relayState;
 
     public IdPInitiatedLoginBuilder(URI authServerSamlUrl, String clientId, SamlClientBuilder clientBuilder) {
         this.clientBuilder = clientBuilder;
@@ -43,7 +48,22 @@ public class IdPInitiatedLoginBuilder implements Step {
 
     @Override
     public HttpUriRequest perform(CloseableHttpClient client, URI currentURI, CloseableHttpResponse currentResponse, HttpClientContext context) throws Exception {
-        return new HttpGet(authServerSamlUrl.toString() + "/clients/" + this.clientId);
+        return new HttpGet(authServerSamlUrl.toString() + "/clients/" + this.clientId + getRelayStateQueryParamString());
+    }
+
+    private String getRelayStateQueryParamString(){
+        if (relayState == null) return "";
+        return "?" + GeneralConstants.RELAY_STATE + "=" + relayState.get();
+    }
+
+    public IdPInitiatedLoginBuilder relayState(String relayState) {
+        this.relayState = () -> relayState;
+        return this;
+    }
+
+    public IdPInitiatedLoginBuilder relayState(Supplier<String> relayState) {
+        this.relayState = relayState;
+        return this;
     }
 
     public SamlClientBuilder build() {

testsuite/integration-arquillian/tests/base/src/test/java/org/keycloak/testsuite/saml/SamlRelayStateTest.java

@@ -0,0 +1,170 @@
+package org.keycloak.testsuite.saml;
+
+import org.junit.Ignore;
+import org.junit.Test;
+import org.keycloak.protocol.saml.SamlProtocol;
+import org.keycloak.testsuite.admin.concurrency.AbstractConcurrencyTest;
+import org.keycloak.testsuite.updaters.ClientAttributeUpdater;
+import org.keycloak.testsuite.util.SamlClient;
+import org.keycloak.testsuite.util.SamlClientBuilder;
+
+import java.io.Closeable;
+import java.util.List;
+import java.util.UUID;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+
+/**
+ * @author mhajas
+ */
+public class SamlRelayStateTest extends AbstractSamlTest {
+
+    private static final String RELAY_STATE = "/importantRelayState";
+
+    @Test
+    public void testRelayStateDoesNotRetainBetweenTwoRequestsPost() throws Exception {
+        new SamlClientBuilder()
+                .authnRequest(getAuthServerSamlEndpoint(REALM_NAME), SAML_CLIENT_ID_SALES_POST, SAML_ASSERTION_CONSUMER_URL_SALES_POST, SamlClient.Binding.POST)
+                    .relayState(RELAY_STATE)
+                    .build()
+
+                .login().user(bburkeUser).build()
+                .assertSamlRelayState(SamlClient.Binding.POST,
+                        relayState -> assertThat(relayState).isEqualTo(RELAY_STATE))
+
+                .authnRequest(getAuthServerSamlEndpoint(REALM_NAME), SAML_CLIENT_ID_SALES_POST, SAML_ASSERTION_CONSUMER_URL_SALES_POST, SamlClient.Binding.POST)
+                    .build()
+                .followOneRedirect()
+                .assertSamlRelayState(SamlClient.Binding.POST,
+                        relayState -> assertThat(relayState).isNull())
+                .execute();
+    }
+
+    @Test
+    public void testRelayStateDoesNotRetainBetweenTwoRequestsRedirect() throws Exception {
+        String url = adminClient.realm(REALM_NAME).clients().findByClientId(SAML_CLIENT_ID_SALES_POST).get(0)
+                .getAttributes().get(SamlProtocol.SAML_ASSERTION_CONSUMER_URL_POST_ATTRIBUTE);
+        try (Closeable c = ClientAttributeUpdater.forClient(adminClient, REALM_NAME, SAML_CLIENT_ID_SALES_POST)
+                .setAttribute(SamlProtocol.SAML_ASSERTION_CONSUMER_URL_POST_ATTRIBUTE, null)
+                .setAttribute(SamlProtocol.SAML_ASSERTION_CONSUMER_URL_REDIRECT_ATTRIBUTE, url)
+                .update()) {
+            new SamlClientBuilder()
+                    .authnRequest(getAuthServerSamlEndpoint(REALM_NAME), SAML_CLIENT_ID_SALES_POST, SAML_ASSERTION_CONSUMER_URL_SALES_POST, SamlClient.Binding.REDIRECT)
+                        .relayState(RELAY_STATE)
+                        .build()
+
+                    .login().user(bburkeUser).build()
+                    .assertSamlRelayState(SamlClient.Binding.REDIRECT,
+                            relayState -> assertThat(relayState).isEqualTo(RELAY_STATE))
+
+                    .authnRequest(getAuthServerSamlEndpoint(REALM_NAME), SAML_CLIENT_ID_SALES_POST, SAML_ASSERTION_CONSUMER_URL_SALES_POST, SamlClient.Binding.REDIRECT)
+                        .build()
+                    .assertSamlRelayState(SamlClient.Binding.REDIRECT,
+                            relayState -> assertThat(relayState).isNull())
+                    .execute();
+        }
+    }
+
+    @Test
+    public void testRelayStateDoesNotRetainBetweenTwoRequestsIdpInitiatedPost() throws Exception {
+        new SamlClientBuilder()
+                .idpInitiatedLogin(getAuthServerSamlEndpoint(REALM_NAME), "sales-post")
+                    .relayState(RELAY_STATE)
+                    .build()
+                .login().user(bburkeUser).build()
+
+                .assertSamlRelayState(SamlClient.Binding.POST,
+                        relayState -> assertThat(relayState).isEqualTo(RELAY_STATE))
+
+                .idpInitiatedLogin(getAuthServerSamlEndpoint(REALM_NAME), "sales-post")
+                    .build()
+
+                .assertSamlRelayState(SamlClient.Binding.POST,
+                        relayState -> assertThat(relayState).isNull())
+                .execute();
+    }
+
+    @Test
+    public void testRelayStateDoesNotRetainBetweenTwoRequestsIdpInitiatedRedirect() throws Exception {
+        String url = adminClient.realm(REALM_NAME).clients().findByClientId(SAML_CLIENT_ID_SALES_POST).get(0)
+                .getAttributes().get(SamlProtocol.SAML_ASSERTION_CONSUMER_URL_POST_ATTRIBUTE);
+        try (Closeable c = ClientAttributeUpdater.forClient(adminClient, REALM_NAME, SAML_CLIENT_ID_SALES_POST)
+                .setAttribute(SamlProtocol.SAML_ASSERTION_CONSUMER_URL_POST_ATTRIBUTE, null)
+                .setAttribute(SamlProtocol.SAML_ASSERTION_CONSUMER_URL_REDIRECT_ATTRIBUTE, url)
+                .update()) {
+            new SamlClientBuilder()
+                    .idpInitiatedLogin(getAuthServerSamlEndpoint(REALM_NAME), "sales-post")
+                        .relayState(RELAY_STATE)
+                        .build()
+                    .login().user(bburkeUser).build()
+
+                    .assertSamlRelayState(SamlClient.Binding.REDIRECT,
+                            relayState -> assertThat(relayState).isEqualTo(RELAY_STATE))
+
+                    .idpInitiatedLogin(getAuthServerSamlEndpoint(REALM_NAME), "sales-post")
+                        .build()
+
+                    .assertSamlRelayState(SamlClient.Binding.REDIRECT,
+                            relayState -> assertThat(relayState).isNull())
+                    .execute();
+        }
+    }
+
+    @Test
+    public void testRelayStateForSameAuthSession() throws Exception {
+        new SamlClientBuilder()
+                .authnRequest(getAuthServerSamlEndpoint(REALM_NAME), SAML_CLIENT_ID_SALES_POST, SAML_ASSERTION_CONSUMER_URL_SALES_POST, SamlClient.Binding.POST)
+                    .relayState(RELAY_STATE)
+                    .build()
+
+                .authnRequest(getAuthServerSamlEndpoint(REALM_NAME), SAML_CLIENT_ID_SALES_POST, SAML_ASSERTION_CONSUMER_URL_SALES_POST, SamlClient.Binding.POST)
+                    .build()
+
+                .login().user(bburkeUser).build()
+                .assertSamlRelayState(SamlClient.Binding.POST,
+                        relayState -> assertThat(relayState).isNull())
+                .execute();
+    }
+
+    @Test
+    public void testRelayStateForSameAuthSessionIDPInitiated() throws Exception {
+        new SamlClientBuilder()
+                .idpInitiatedLogin(getAuthServerSamlEndpoint(REALM_NAME), "sales-post")
+                    .relayState(RELAY_STATE)
+                    .build()
+
+                .idpInitiatedLogin(getAuthServerSamlEndpoint(REALM_NAME), "sales-post").build()
+                .login().user(bburkeUser).build()
+                .assertSamlRelayState(SamlClient.Binding.POST,
+                        relayState -> assertThat(relayState).isNull())
+                .execute();
+    }
+
+    @Test
+    @Ignore("KEYCLOAK-5179")
+    public void relayStateConcurrencyTest() throws Exception {
+        ThreadLocal<UUID> tl = new ThreadLocal<>();
+
+        List<SamlClient.Step> steps = new SamlClientBuilder()
+                .addStep(() -> tl.set(UUID.randomUUID()))
+                .authnRequest(getAuthServerSamlEndpoint(REALM_NAME), SAML_CLIENT_ID_SALES_POST, SAML_ASSERTION_CONSUMER_URL_SALES_POST, SamlClient.Binding.POST)
+                    .relayState(() -> tl.get().toString())
+                    .build()
+
+                .login().user(bburkeUser).build()
+                .assertSamlRelayState(SamlClient.Binding.POST, relayState -> {
+                    assertThat(relayState).isNotNull();
+                    assertThat(relayState).isEqualTo(tl.get().toString());
+                })
+                .getSteps();
+
+        SamlClient client = new SamlClient();
+        client.execute(steps);
+        steps.remove(2); // removing login as it should not be necessary anymore
+
+        AbstractConcurrencyTest.run(2, 10, this, (threadIndex, keycloak, realm) -> {
+            client.execute(steps);
+        });
+    }
+}