Implement SAML custom attributes support for Identity Provider

docs/changelog/128176.yaml

@@ -0,0 +1,5 @@
+pr: 128176
+summary: Implement SAML custom attributes support for Identity Provider
+area: Authentication
+type: enhancement
+issues: []

x-pack/plugin/identity-provider/qa/idp-rest-tests/src/javaRestTest/java/org/elasticsearch/xpack/idp/IdentityProviderAuthenticationIT.java

@@ -18,23 +18,38 @@
 import org.elasticsearch.common.util.concurrent.ThreadContext;
 import org.elasticsearch.core.Nullable;
 import org.elasticsearch.xcontent.ObjectPath;
+import org.elasticsearch.xcontent.XContentBuilder;
 import org.elasticsearch.xcontent.json.JsonXContent;
 import org.elasticsearch.xpack.core.security.action.saml.SamlPrepareAuthenticationResponse;
 import org.elasticsearch.xpack.idp.saml.sp.SamlServiceProviderIndex;
 import org.junit.Before;
+import org.w3c.dom.Document;
+import org.w3c.dom.Element;
+import org.w3c.dom.NodeList;
+import org.xml.sax.InputSource;
 
 import java.io.IOException;
+import java.io.StringReader;
 import java.nio.charset.StandardCharsets;
+import java.util.ArrayList;
 import java.util.Base64;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
 
+import javax.xml.parsers.DocumentBuilder;
+import javax.xml.parsers.DocumentBuilderFactory;
+import javax.xml.xpath.XPath;
+import javax.xml.xpath.XPathConstants;
+import javax.xml.xpath.XPathFactory;
+
 import static org.hamcrest.Matchers.containsInAnyOrder;
 import static org.hamcrest.Matchers.containsString;
 import static org.hamcrest.Matchers.equalTo;
 import static org.hamcrest.Matchers.hasSize;
 import static org.hamcrest.Matchers.instanceOf;
+import static org.hamcrest.Matchers.is;
+import static org.hamcrest.Matchers.notNullValue;
 
 public class IdentityProviderAuthenticationIT extends IdpRestTestCase {
 
@@ -74,6 +89,81 @@ public void testRegistrationAndIdpInitiatedSso() throws Exception {
         authenticateWithSamlResponse(samlResponse, null);
     }
 
+    public void testCustomAttributesInIdpInitiatedSso() throws Exception {
+        final Map<String, Object> request = Map.ofEntries(
+            Map.entry("name", "Test SP With Custom Attributes"),
+            Map.entry("acs", SP_ACS),
+            Map.entry("privileges", Map.ofEntries(Map.entry("resource", SP_ENTITY_ID), Map.entry("roles", List.of("sso:(\\w+)")))),
+            Map.entry(
+                "attributes",
+                Map.ofEntries(
+                    Map.entry("principal", "https://idp.test.es.elasticsearch.org/attribute/principal"),
+                    Map.entry("name", "https://idp.test.es.elasticsearch.org/attribute/name"),
+                    Map.entry("email", "https://idp.test.es.elasticsearch.org/attribute/email"),
+                    Map.entry("roles", "https://idp.test.es.elasticsearch.org/attribute/roles")
+                )
+            )
+        );
+        final SamlServiceProviderIndex.DocumentVersion docVersion = createServiceProvider(SP_ENTITY_ID, request);
+        checkIndexDoc(docVersion);
+        ensureGreen(SamlServiceProviderIndex.INDEX_NAME);
+
+        // Create custom attributes
+        Map<String, List<String>> attributesMap = Map.of("department", List.of("engineering", "product"), "region", List.of("APJ"));
+
+        // Generate SAML response with custom attributes
+        final String samlResponse = generateSamlResponseWithAttributes(SP_ENTITY_ID, SP_ACS, null, attributesMap);
+
+        // Parse XML directly from samlResponse (it's not base64 encoded at this point)
+        DocumentBuilderFactory factory = DocumentBuilderFactory.newInstance();
+        factory.setNamespaceAware(true); // Required for XPath
+        DocumentBuilder builder = factory.newDocumentBuilder();
+        Document document = builder.parse(new InputSource(new StringReader(samlResponse)));
+
+        // Create XPath evaluator
+        XPathFactory xPathFactory = XPathFactory.newInstance();
+        XPath xpath = xPathFactory.newXPath();
+
+        // Validate SAML Response structure
+        Element responseElement = (Element) xpath.evaluate("//*[local-name()='Response']", document, XPathConstants.NODE);
+        assertThat("SAML Response element should exist", responseElement, notNullValue());
+
+        Element assertionElement = (Element) xpath.evaluate("//*[local-name()='Assertion']", document, XPathConstants.NODE);
+        assertThat("SAML Assertion element should exist", assertionElement, notNullValue());
+
+        // Validate department attribute
+        NodeList departmentAttributes = (NodeList) xpath.evaluate(
+            "//*[local-name()='Attribute' and @Name='department']/*[local-name()='AttributeValue']",
+            document,
+            XPathConstants.NODESET
+        );
+
+        assertThat("Should have two values for department attribute", departmentAttributes.getLength(), is(2));
+
+        // Verify department values
+        List<String> departmentValues = new ArrayList<>();
+        for (int i = 0; i < departmentAttributes.getLength(); i++) {
+            departmentValues.add(departmentAttributes.item(i).getTextContent());
+        }
+        assertThat(
+            "Department attribute should contain 'engineering' and 'product'",
+            departmentValues,
+            containsInAnyOrder("engineering", "product")
+        );
+
+        // Validate region attribute
+        NodeList regionAttributes = (NodeList) xpath.evaluate(
+            "//*[local-name()='Attribute' and @Name='region']/*[local-name()='AttributeValue']",
+            document,
+            XPathConstants.NODESET
+        );
+
+        assertThat("Should have one value for region attribute", regionAttributes.getLength(), is(1));
+        assertThat("Region attribute should contain 'APJ'", regionAttributes.item(0).getTextContent(), equalTo("APJ"));
+
+        authenticateWithSamlResponse(samlResponse, null);
+    }
+
     public void testRegistrationAndSpInitiatedSso() throws Exception {
         final Map<String, Object> request = Map.ofEntries(
             Map.entry("name", "Test SP"),
@@ -125,17 +215,37 @@ private SamlPrepareAuthenticationResponse generateSamlAuthnRequest(String realmN
         }
     }
 
-    private String generateSamlResponse(String entityId, String acs, @Nullable Map<String, Object> authnState) throws Exception {
+    private String generateSamlResponse(String entityId, String acs, @Nullable Map<String, Object> authnState) throws IOException {
+        return generateSamlResponseWithAttributes(entityId, acs, authnState, null);
+    }
+
+    private String generateSamlResponseWithAttributes(
+        String entityId,
+        String acs,
+        @Nullable Map<String, Object> authnState,
+        @Nullable Map<String, List<String>> attributes
+    ) throws IOException {
         final Request request = new Request("POST", "/_idp/saml/init");
-        if (authnState != null && authnState.isEmpty() == false) {
-            request.setJsonEntity(Strings.format("""
-                {"entity_id":"%s", "acs":"%s","authn_state":%s}
-                """, entityId, acs, Strings.toString(JsonXContent.contentBuilder().map(authnState))));
-        } else {
-            request.setJsonEntity(Strings.format("""
-                {"entity_id":"%s", "acs":"%s"}
-                """, entityId, acs));
+
+        XContentBuilder builder = JsonXContent.contentBuilder();
+        builder.startObject();
+        builder.field("entity_id", entityId);
+        builder.field("acs", acs);
+
+        if (authnState != null) {
+            builder.field("authn_state");
+            builder.map(authnState);
+        }
+
+        if (attributes != null) {
+            builder.field("attributes");
+            builder.map(attributes);
         }
+
+        builder.endObject();
+        String jsonEntity = Strings.toString(builder);
+
+        request.setJsonEntity(jsonEntity);
         request.setOptions(
             RequestOptions.DEFAULT.toBuilder()
                 .addHeader("es-secondary-authorization", basicAuthHeaderValue("idp_user", new SecureString("idp-password".toCharArray())))

x-pack/plugin/identity-provider/src/main/java/org/elasticsearch/xpack/idp/action/SamlInitiateSingleSignOnRequest.java

@@ -12,6 +12,7 @@
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
 import org.elasticsearch.xpack.idp.saml.support.SamlAuthenticationState;
+import org.elasticsearch.xpack.idp.saml.support.SamlInitiateSingleSignOnAttributes;
 
 import java.io.IOException;
 
@@ -22,12 +23,14 @@ public class SamlInitiateSingleSignOnRequest extends LegacyActionRequest {
     private String spEntityId;
     private String assertionConsumerService;
     private SamlAuthenticationState samlAuthenticationState;
+    private SamlInitiateSingleSignOnAttributes attributes;
 
     public SamlInitiateSingleSignOnRequest(StreamInput in) throws IOException {
         super(in);
         spEntityId = in.readString();
         assertionConsumerService = in.readString();
         samlAuthenticationState = in.readOptionalWriteable(SamlAuthenticationState::new);
+        attributes = in.readOptionalWriteable(SamlInitiateSingleSignOnAttributes::new);
     }
 
     public SamlInitiateSingleSignOnRequest() {}
@@ -41,6 +44,17 @@ public ActionRequestValidationException validate() {
         if (Strings.isNullOrEmpty(assertionConsumerService)) {
             validationException = addValidationError("acs is missing", validationException);
         }
+
+        // Validate attributes if present
+        if (attributes != null) {
+            ActionRequestValidationException attributesValidationException = attributes.validate();
+            if (attributesValidationException != null) {
+                for (String error : attributesValidationException.validationErrors()) {
+                    validationException = addValidationError(error, validationException);
+                }
+            }
+        }
+
         return validationException;
     }
 
@@ -68,17 +82,36 @@ public void setSamlAuthenticationState(SamlAuthenticationState samlAuthenticatio
         this.samlAuthenticationState = samlAuthenticationState;
     }
 
+    public SamlInitiateSingleSignOnAttributes getAttributes() {
+        return attributes;
+    }
+
+    public void setAttributes(SamlInitiateSingleSignOnAttributes attributes) {
+        this.attributes = attributes;
+    }
+
     @Override
     public void writeTo(StreamOutput out) throws IOException {
         super.writeTo(out);
         out.writeString(spEntityId);
         out.writeString(assertionConsumerService);
         out.writeOptionalWriteable(samlAuthenticationState);
+        out.writeOptionalWriteable(attributes);
     }
 
     @Override
     public String toString() {
-        return getClass().getSimpleName() + "{spEntityId='" + spEntityId + "', acs='" + assertionConsumerService + "'}";
+        return getClass().getSimpleName()
+            + "{"
+            + "spEntityId='"
+            + spEntityId
+            + "', "
+            + "acs='"
+            + assertionConsumerService
+            + "', "
+            + "attributes="
+            + attributes
+            + "'}";
     }
 
 }

x-pack/plugin/identity-provider/src/main/java/org/elasticsearch/xpack/idp/action/SamlInitiateSingleSignOnResponse.java

@@ -10,6 +10,7 @@
 import org.elasticsearch.common.io.stream.StreamOutput;
 import org.elasticsearch.core.Nullable;
 import org.elasticsearch.xcontent.XContentBuilder;
+import org.elasticsearch.xpack.idp.saml.support.SamlInitiateSingleSignOnAttributes;
 
 import java.io.IOException;
 
@@ -20,19 +21,32 @@ public class SamlInitiateSingleSignOnResponse extends ActionResponse {
     private final String entityId;
     private final String samlStatus;
     private final String error;
+    private final SamlInitiateSingleSignOnAttributes attributes;
 
     public SamlInitiateSingleSignOnResponse(
         String entityId,
         String postUrl,
         String samlResponse,
         String samlStatus,
         @Nullable String error
+    ) {
+        this(entityId, postUrl, samlResponse, samlStatus, error, null);
+    }
+
+    public SamlInitiateSingleSignOnResponse(
+        String entityId,
+        String postUrl,
+        String samlResponse,
+        String samlStatus,
+        @Nullable String error,
+        @Nullable SamlInitiateSingleSignOnAttributes attributes
     ) {
         this.entityId = entityId;
         this.postUrl = postUrl;
         this.samlResponse = samlResponse;
         this.samlStatus = samlStatus;
         this.error = error;
+        this.attributes = attributes;
     }
 
     public String getPostUrl() {
@@ -55,13 +69,18 @@ public String getSamlStatus() {
         return samlStatus;
     }
 
+    public SamlInitiateSingleSignOnAttributes getAttributes() {
+        return attributes;
+    }
+
     @Override
     public void writeTo(StreamOutput out) throws IOException {
         out.writeString(entityId);
         out.writeString(postUrl);
         out.writeString(samlResponse);
         out.writeString(samlStatus);
         out.writeOptionalString(error);
+        out.writeOptionalWriteable(attributes);
     }
 
     public void toXContent(XContentBuilder builder) throws IOException {

x-pack/plugin/identity-provider/src/main/java/org/elasticsearch/xpack/idp/action/TransportSamlInitiateSingleSignOnAction.java

@@ -139,7 +139,7 @@ protected void doExecute(
                         identityProvider
                     );
                     try {
-                        final Response response = builder.build(user, authenticationState);
+                        final Response response = builder.build(user, authenticationState, request.getAttributes());
                         listener.onResponse(
                             new SamlInitiateSingleSignOnResponse(
                                 user.getServiceProvider().getEntityId(),