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/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(),

x-pack/plugin/identity-provider/src/main/java/org/elasticsearch/xpack/idp/saml/authn/SuccessfulAuthenticationResponseMessageBuilder.java

@@ -18,6 +18,7 @@
 import org.elasticsearch.xpack.idp.saml.support.SamlAuthenticationState;
 import org.elasticsearch.xpack.idp.saml.support.SamlFactory;
 import org.elasticsearch.xpack.idp.saml.support.SamlInit;
+import org.elasticsearch.xpack.idp.saml.support.SamlInitiateSingleSignOnAttributes;
 import org.elasticsearch.xpack.idp.saml.support.SamlObjectSigner;
 import org.opensaml.core.xml.schema.XSString;
 import org.opensaml.saml.saml2.core.Assertion;
@@ -44,6 +45,7 @@
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.List;
+import java.util.Map;
 import java.util.Set;
 
 import static org.opensaml.saml.saml2.core.NameIDType.TRANSIENT;
@@ -66,7 +68,30 @@ public SuccessfulAuthenticationResponseMessageBuilder(SamlFactory samlFactory, C
         this.idp = idp;
     }
 
+    /**
+     * Builds and signs a SAML Response Message with a single assertion for the provided user
+     *
+     * @param user        The user who is authenticated (actually a combination of user+sp)
+     * @param authnState  The authentication state as presented in the SAML request (or {@code null})
+     * @return A SAML Response
+     */
     public Response build(UserServiceAuthentication user, @Nullable SamlAuthenticationState authnState) {
+        return build(user, authnState, null);
+    }
+
+    /**
+     * Builds and signs a SAML Response Message with a single assertion for the provided user
+     *
+     * @param user        The user who is authenticated (actually a combination of user+sp)
+     * @param authnState  The authentication state as presented in the SAML request (or {@code null})
+     * @param customAttributes  Optional custom attributes to include in the response (or {@code null})
+     * @return A SAML Response
+     */
+    public Response build(
+        UserServiceAuthentication user,
+        @Nullable SamlAuthenticationState authnState,
+        @Nullable SamlInitiateSingleSignOnAttributes customAttributes
+    ) {
         logger.debug("Building success response for [{}] from [{}]", user, authnState);
         final Instant now = clock.instant();
         final SamlServiceProvider serviceProvider = user.getServiceProvider();
@@ -87,10 +112,13 @@ public Response build(UserServiceAuthentication user, @Nullable SamlAuthenticati
         assertion.setIssueInstant(now);
         assertion.setConditions(buildConditions(now, serviceProvider));
         assertion.setSubject(buildSubject(now, user, authnState));
-        assertion.getAuthnStatements().add(buildAuthnStatement(now, user));
-        final AttributeStatement attributes = buildAttributes(user);
-        if (attributes != null) {
-            assertion.getAttributeStatements().add(attributes);
+
+        final AuthnStatement authnStatement = buildAuthnStatement(now, user);
+        assertion.getAuthnStatements().add(authnStatement);
+
+        final AttributeStatement attributeStatement = buildAttributes(user, customAttributes);
+        if (attributeStatement != null) {
+            assertion.getAttributeStatements().add(attributeStatement);
         }
         response.getAssertions().add(assertion);
         return sign(response);
@@ -179,7 +207,10 @@ private static String resolveAuthnClass(Set<AuthenticationMethod> authentication
         }
     }
 
-    private AttributeStatement buildAttributes(UserServiceAuthentication user) {
+    private AttributeStatement buildAttributes(
+        UserServiceAuthentication user,
+        @Nullable SamlInitiateSingleSignOnAttributes customAttributes
+    ) {
         final SamlServiceProvider serviceProvider = user.getServiceProvider();
         final AttributeStatement statement = samlFactory.object(AttributeStatement.class, AttributeStatement.DEFAULT_ELEMENT_NAME);
         final List<Attribute> attributes = new ArrayList<>();
@@ -199,27 +230,39 @@ private AttributeStatement buildAttributes(UserServiceAuthentication user) {
         if (name != null) {
             attributes.add(name);
         }
+        // Add custom attributes if provided
+        if (customAttributes != null && customAttributes.getAttributes().isEmpty() == false) {
+            for (Map.Entry<String, List<String>> entry : customAttributes.getAttributes().entrySet()) {
+                Attribute attribute = buildAttribute(entry.getKey(), null, entry.getValue());
+                if (attribute != null) {
+                    attributes.add(attribute);
+                }
+            }
+        }
+
         if (attributes.isEmpty()) {
             return null;
         }
         statement.getAttributes().addAll(attributes);
         return statement;
     }
 
-    private Attribute buildAttribute(String formalName, String friendlyName, String value) {
+    private Attribute buildAttribute(String formalName, @Nullable String friendlyName, String value) {
         if (Strings.isNullOrEmpty(value)) {
             return null;
         }
         return buildAttribute(formalName, friendlyName, List.of(value));
     }
 
-    private Attribute buildAttribute(String formalName, String friendlyName, Collection<String> values) {
+    private Attribute buildAttribute(String formalName, @Nullable String friendlyName, Collection<String> values) {
         if (values.isEmpty() || Strings.isNullOrEmpty(formalName)) {
             return null;
         }
         final Attribute attribute = samlFactory.object(Attribute.class, Attribute.DEFAULT_ELEMENT_NAME);
         attribute.setName(formalName);
-        attribute.setFriendlyName(friendlyName);
+        if (Strings.isNullOrEmpty(friendlyName) == false) {
+            attribute.setFriendlyName(friendlyName);
+        }
         attribute.setNameFormat(Attribute.URI_REFERENCE);
         for (String val : values) {
             final XSString string = samlFactory.object(XSString.class, AttributeValue.DEFAULT_ELEMENT_NAME, XSString.TYPE_NAME);