From ae98d8ea286c4f507beeb1a91f3c298966f11fd6 Mon Sep 17 00:00:00 2001
From: Luca Leonardo Scorcia <l.scorcia@gmail.com>
Date: Sun, 30 May 2021 06:43:58 -0400
Subject: [PATCH] KEYCLOAK-18315 SAML Client - Add parameter to request
 specific AttributeConsumingServiceIndex

---
 .../saml/SAML2AuthnRequestBuilder.java        |  9 +-
 .../broker/saml/SAMLIdentityProvider.java     |  3 +
 .../saml/SAMLIdentityProviderConfig.java      | 25 +++++
 ...amlAttributeConsumingServiceIndexTest.java | 98 +++++++++++++++++++
 .../messages/admin-messages_en.properties     |  2 +
 .../realm-identity-provider-saml.html         |  7 ++
 6 files changed, 142 insertions(+), 2 deletions(-)
 create mode 100644 testsuite/integration-arquillian/tests/base/src/test/java/org/keycloak/testsuite/broker/KcSamlAttributeConsumingServiceIndexTest.java

diff --git a/saml-core/src/main/java/org/keycloak/saml/SAML2AuthnRequestBuilder.java b/saml-core/src/main/java/org/keycloak/saml/SAML2AuthnRequestBuilder.java
index cf7da5cc6e..a4f7f53cea 100755
--- a/saml-core/src/main/java/org/keycloak/saml/SAML2AuthnRequestBuilder.java
+++ b/saml-core/src/main/java/org/keycloak/saml/SAML2AuthnRequestBuilder.java
@@ -75,6 +75,11 @@ public class SAML2AuthnRequestBuilder implements SamlProtocolExtensionsAwareBuil
         return this;
     }
 
+    public SAML2AuthnRequestBuilder attributeConsumingServiceIndex(Integer attributeConsumingServiceIndex) {
+        this.authnRequestType.setAttributeConsumingServiceIndex(attributeConsumingServiceIndex);
+        return this;
+    }
+
     public SAML2AuthnRequestBuilder forceAuthn(boolean forceAuthn) {
         this.authnRequestType.setForceAuthn(forceAuthn);
         return this;
@@ -85,8 +90,8 @@ public class SAML2AuthnRequestBuilder implements SamlProtocolExtensionsAwareBuil
         return this;
     }
 
-    public SAML2AuthnRequestBuilder nameIdPolicy(SAML2NameIDPolicyBuilder nameIDPolicy) {
-        this.authnRequestType.setNameIDPolicy(nameIDPolicy.build());
+    public SAML2AuthnRequestBuilder nameIdPolicy(SAML2NameIDPolicyBuilder nameIDPolicyBuilder) {
+        this.authnRequestType.setNameIDPolicy(nameIDPolicyBuilder.build());
         return this;
     }
 
diff --git a/services/src/main/java/org/keycloak/broker/saml/SAMLIdentityProvider.java b/services/src/main/java/org/keycloak/broker/saml/SAMLIdentityProvider.java
index 7b324ec0ce..5ce7c2eed3 100755
--- a/services/src/main/java/org/keycloak/broker/saml/SAMLIdentityProvider.java
+++ b/services/src/main/java/org/keycloak/broker/saml/SAMLIdentityProvider.java
@@ -130,6 +130,8 @@ public class SAMLIdentityProvider extends AbstractIdentityProvider<SAMLIdentityP
             for (String authnContextDeclRef : getAuthnContextDeclRefUris())
                 requestedAuthnContext.addAuthnContextDeclRef(authnContextDeclRef);
 
+            Integer attributeConsumingServiceIndex = getConfig().getAttributeConsumingServiceIndex();
+
             String loginHint = getConfig().isLoginHint() ? request.getAuthenticationSession().getClientNote(OIDCLoginProtocol.LOGIN_HINT_PARAM) : null;
             Boolean allowCreate = null;
             if (getConfig().getConfig().get(SAMLIdentityProviderConfig.ALLOW_CREATE) == null || getConfig().isAllowCreate())
@@ -143,6 +145,7 @@ public class SAMLIdentityProvider extends AbstractIdentityProvider<SAMLIdentityP
                     .nameIdPolicy(SAML2NameIDPolicyBuilder
                         .format(nameIDPolicyFormat)
                         .setAllowCreate(allowCreate))
+                    .attributeConsumingServiceIndex(attributeConsumingServiceIndex)
                     .requestedAuthnContext(requestedAuthnContext)
                     .subject(loginHint);
 
diff --git a/services/src/main/java/org/keycloak/broker/saml/SAMLIdentityProviderConfig.java b/services/src/main/java/org/keycloak/broker/saml/SAMLIdentityProviderConfig.java
index 372106db18..de34865b9f 100755
--- a/services/src/main/java/org/keycloak/broker/saml/SAMLIdentityProviderConfig.java
+++ b/services/src/main/java/org/keycloak/broker/saml/SAMLIdentityProviderConfig.java
@@ -60,6 +60,7 @@ public class SAMLIdentityProviderConfig extends IdentityProviderModel {
     public static final String AUTHN_CONTEXT_DECL_REFS = "authnContextDeclRefs";
     public static final String SIGN_SP_METADATA = "signSpMetadata";
     public static final String ALLOW_CREATE = "allowCreate";
+    public static final String ATTRIBUTE_CONSUMING_SERVICE_INDEX = "attributeConsumingServiceIndex";
 
     public SAMLIdentityProviderConfig() {
     }
@@ -345,6 +346,30 @@ public class SAMLIdentityProviderConfig extends IdentityProviderModel {
         getConfig().put(ALLOW_CREATE, String.valueOf(allowCreate));
     }
 
+    public Integer getAttributeConsumingServiceIndex() {
+        Integer result = null;
+        String strAttributeConsumingServiceIndex = getConfig().get(ATTRIBUTE_CONSUMING_SERVICE_INDEX);
+        if (strAttributeConsumingServiceIndex != null && !strAttributeConsumingServiceIndex.isEmpty()) {
+            try {
+                result = Integer.parseInt(strAttributeConsumingServiceIndex);
+                if (result < 0) {
+                    result = null;
+                }
+            } catch (NumberFormatException e) {
+                // ignore it and use null
+            }
+        }
+        return result;
+    }
+
+    public void setAttributeConsumingServiceIndex(Integer attributeConsumingServiceIndex) {
+        if (attributeConsumingServiceIndex == null || attributeConsumingServiceIndex < 0) {
+            getConfig().remove(ATTRIBUTE_CONSUMING_SERVICE_INDEX);
+        } else {
+            getConfig().put(ATTRIBUTE_CONSUMING_SERVICE_INDEX, String.valueOf(attributeConsumingServiceIndex));
+        }
+    }
+
     @Override
     public void validate(RealmModel realm) {
         SslRequired sslRequired = realm.getSslRequired();
diff --git a/testsuite/integration-arquillian/tests/base/src/test/java/org/keycloak/testsuite/broker/KcSamlAttributeConsumingServiceIndexTest.java b/testsuite/integration-arquillian/tests/base/src/test/java/org/keycloak/testsuite/broker/KcSamlAttributeConsumingServiceIndexTest.java
new file mode 100644
index 0000000000..7092690b2c
--- /dev/null
+++ b/testsuite/integration-arquillian/tests/base/src/test/java/org/keycloak/testsuite/broker/KcSamlAttributeConsumingServiceIndexTest.java
@@ -0,0 +1,98 @@
+package org.keycloak.testsuite.broker;
+
+import org.keycloak.broker.saml.SAMLIdentityProviderConfig;
+import org.keycloak.dom.saml.v2.protocol.AuthnRequestType;
+import org.keycloak.saml.common.util.DocumentUtil;
+import org.keycloak.saml.processing.api.saml.v2.request.SAML2Request;
+import org.keycloak.testsuite.saml.AbstractSamlTest;
+import org.keycloak.testsuite.updaters.IdentityProviderAttributeUpdater;
+import org.keycloak.testsuite.util.SamlClient;
+import org.keycloak.testsuite.util.SamlClient.Binding;
+import org.keycloak.testsuite.util.SamlClientBuilder;
+import java.io.Closeable;
+
+import org.junit.Assert;
+import org.junit.Test;
+import org.w3c.dom.Document;
+import org.w3c.dom.Node;
+import static org.keycloak.testsuite.broker.BrokerTestTools.getConsumerRoot;
+
+/**
+ * Final class as it's not intended to be overriden.
+ */
+public final class KcSamlAttributeConsumingServiceIndexTest extends AbstractBrokerTest {
+
+    @Override
+    protected BrokerConfiguration getBrokerConfiguration() {
+        return KcSamlBrokerConfiguration.INSTANCE;
+    }
+
+    @Test
+    public void testAttributeConsumingServiceIndexNotSet() throws Exception {
+        // No Attribute Consuming Service Index set -> No attribute added to AuthnRequest
+        try (Closeable idpUpdater = new IdentityProviderAttributeUpdater(identityProviderResource)
+            .update())
+        {
+            // Build the login request document
+            AuthnRequestType loginRep = SamlClient.createLoginRequestDocument(AbstractSamlTest.SAML_CLIENT_ID_SALES_POST + ".dot/ted", getConsumerRoot() + "/sales-post/saml", null);
+            Document doc = SAML2Request.convert(loginRep);
+            new SamlClientBuilder()
+                .authnRequest(getConsumerSamlEndpoint(bc.consumerRealmName()), doc, Binding.POST)
+                .build()   // Request to consumer IdP
+                .login().idp(bc.getIDPAlias()).build()
+                .processSamlResponse(Binding.POST)    // AuthnRequest to producer IdP
+                  .targetAttributeSamlRequest()
+                  .transformDocument((document) -> {
+                    try
+                    {
+                        log.infof("Document: %s", DocumentUtil.asString(document));
+
+                        // Find the AuthnRequest AttributeConsumingServiceIndex attribute
+                        Node attrNode = document.getDocumentElement().getAttributes().getNamedItem("AttributeConsumingServiceIndex");
+                        Assert.assertEquals("Unexpected AttributeConsumingServiceIndex attribute value", null, attrNode);
+                    }
+                    catch (Exception ex)
+                    {
+                        throw new RuntimeException(ex);
+                    }
+                  })
+                  .build()
+                .execute();
+        }
+    }
+
+    @Test
+    public void testAttributeConsumingServiceIndexSet() throws Exception {
+        // Attribute Consuming Service Index set -> Attribute added to AuthnRequest
+        try (Closeable idpUpdater = new IdentityProviderAttributeUpdater(identityProviderResource)
+            .setAttribute(SAMLIdentityProviderConfig.ATTRIBUTE_CONSUMING_SERVICE_INDEX, "15")
+            .update())
+        {
+            // Build the login request document
+            AuthnRequestType loginRep = SamlClient.createLoginRequestDocument(AbstractSamlTest.SAML_CLIENT_ID_SALES_POST + ".dot/ted", getConsumerRoot() + "/sales-post/saml", null);
+            Document doc = SAML2Request.convert(loginRep);
+            new SamlClientBuilder()
+                .authnRequest(getConsumerSamlEndpoint(bc.consumerRealmName()), doc, Binding.POST)
+                .build()   // Request to consumer IdP
+                .login().idp(bc.getIDPAlias()).build()
+                .processSamlResponse(Binding.POST)    // AuthnRequest to producer IdP
+                  .targetAttributeSamlRequest()
+                  .transformDocument((document) -> {
+                    try
+                    {
+                        log.infof("Document: %s", DocumentUtil.asString(document));
+
+                        // Find the AuthnRequest AttributeConsumingServiceIndex attribute
+                        String attrValue = document.getDocumentElement().getAttributes().getNamedItem("AttributeConsumingServiceIndex").getNodeValue();
+                        Assert.assertEquals("Unexpected AttributeConsumingServiceIndex attribute value", "15", attrValue);
+                    }
+                    catch (Exception ex)
+                    {
+                        throw new RuntimeException(ex);
+                    }
+                  })
+                  .build()
+                .execute();
+        }
+    }
+}
\ No newline at end of file
diff --git a/themes/src/main/resources/theme/base/admin/messages/admin-messages_en.properties b/themes/src/main/resources/theme/base/admin/messages/admin-messages_en.properties
index bb391adf2e..ccc94bf721 100644
--- a/themes/src/main/resources/theme/base/admin/messages/admin-messages_en.properties
+++ b/themes/src/main/resources/theme/base/admin/messages/admin-messages_en.properties
@@ -718,6 +718,8 @@ identity-provider.saml.entity-id=Service Provider Entity ID
 identity-provider.saml.entity-id.tooltip=The Entity ID that will be used to uniquely identify this SAML Service Provider
 identity-provider.saml.protocol-endpoints.saml=SAML 2.0 Service Provider Metadata
 identity-provider.saml.protocol-endpoints.saml.tooltip=Shows the configuration of the Service Provider endpoint
+identity-provider.saml.attribute-consuming-service-index=Attribute Consuming Service Index
+identity-provider.saml.attribute-consuming-service-index.tooltip=Index of the Attribute Consuming Service profile to request during authentication
 saml-config=SAML Config
 identity-provider.saml-config.tooltip=SAML SP and external IDP configuration.
 single-signon-service-url=Single Sign-On Service URL
diff --git a/themes/src/main/resources/theme/base/admin/resources/partials/realm-identity-provider-saml.html b/themes/src/main/resources/theme/base/admin/resources/partials/realm-identity-provider-saml.html
index fb1a38d9df..933f5bebfe 100755
--- a/themes/src/main/resources/theme/base/admin/resources/partials/realm-identity-provider-saml.html
+++ b/themes/src/main/resources/theme/base/admin/resources/partials/realm-identity-provider-saml.html
@@ -302,6 +302,13 @@
                 </div>
                 <kc-tooltip>{{:: 'identity-provider.allowed-clock-skew.tooltip' | translate}}</kc-tooltip>
             </div>
+            <div class="form-group">
+                <label class="col-md-2 control-label" for="attributeConsumingServiceIndex">{{:: 'identity-provider.saml.attribute-consuming-service-index' | translate}}</label>
+                <div class="col-md-6">
+                    <input class="form-control" string-to-number type="number" min="0" max="65535" step="1" ng-model="identityProvider.config.attributeConsumingServiceIndex" id="attributeConsumingServiceIndex"/>
+                </div>
+                <kc-tooltip>{{:: 'identity-provider.saml.attribute-consuming-service-index.tooltip' | translate}}</kc-tooltip>
+            </div>
         </fieldset>
         <fieldset>
             <legend collapsed><span class="text">{{:: 'identity-provider.saml.requested-authncontext' | translate}}</span> <kc-tooltip>{{:: 'identity-provider.saml.requested-authncontext.tooltip' | translate}}</kc-tooltip></legend>