[keycloak-dev] Support for more flexible SAML handling for Attribute, AttributeValue and AttributeStatement elements

[Thomas Darimont]

Hello,

In my current project I needed to support extensive customizations of SAML
"Attribute"-elements in Keycloak SAML responses.
Unfortunately keycloaks
"org.keycloak.protocol.saml.mappers.SAMLAttributeStatementMapper" was not
powerful enough to support
all required customizations.

Some features that were missing / problems I faced were:
- Support for custom AttributeValue types, c.f. "xsd:anyURI" to represent
urn:oid values
- Support for complex nested elements (see example below)
- SAMLAttributeStatementMapper's effectivenes is restricted by
org.keycloak.saml.processing.core.saml.v2.writers.BaseWriter.writeAttributeTypeWithoutRootTag
- Lacking support for custom core protocol implementations, like "saml"
- provider-id "saml" should not be used to identify the "protocol"

Here is an example "AttributeStatement" fragment that I needed to support.
...
        <AttributeStatement>
            <Attribute Name="Role"
NameFormat="urn:oasis:names:tc:SAML:2.0:attrname-format:basic">
                <AttributeValue xmlns:xsi="
http://www.w3.org/2001/XMLSchema-instance"
xsi:type="xsd:string">dummy</AttributeValue>
            </Attribute>

            <Attribute FriendlyName="XSPA Organization ID"
                Name="urn:oasis:names:tc:xspa:1.0:subject:organization-id"

NameFormat="urn:oasis:names:tc:SAML:2.0:attrname-format:uri">
                <AttributeValue xmlns:xsi="
http://www.w3.org/2001/XMLSchema-instance"
xsi:type="xsd:anyURI">urn:oid:1.2.3.4.5.6.7.8.9.10.11.12</AttributeValue>
            </Attribute>

            <Attribute FriendlyName="Acme Role"
Name="urn:oasis:names:tc:xacml:2.0:subject:role"

NameFormat="urn:oasis:names:tc:SAML:2.0:attrname-format:uri">
                <AttributeValue xmlns:xsi="
http://www.w3.org/2001/XMLSchema-instance" xsi:type="xsd:anyType">
                <Role code="PRA" codeSystem="1.2.3.4.5.6.7.8.9.10.11.12"
                          codeSystemName="IHEXDShealthcareFacilityTypeCode"
displayName="Doctor's office" xmlns="urn:hl7-org:v3"/>
                </AttributeValue>
            </Attribute>
        </AttributeStatement>
...
a full example for a complete SAML response can be found here for
reference:
https://gist.github.com/thomasdarimont/02a7861562d0684861b213a1d16b7047

I eventually managed to customize the SAML responses as needed, but my
approach is probably not 100% future proof...
If anyone here knows a better way to do this, I'm all ears :)

Nevertheless, it would be great if Keycloaks support for custom SAML
attributes would be more flexible in the future.

Anyways, here is what I had to do in order to get the required SAML
customizations working:

- Create a "org.keycloak.protocol.LoginProtocolFactory" extension and
reference "CustomSamlProtocolFactory"
- Create class CustomSamlProtocolFactory extends SamlProtocolFactory
- Create class CustomSamlProtocol extends SamlProtocol and return a new
CustomSamlProtocol instance in CustomSamlProtocolFactory#create(..) (see
below)
- override
org.keycloak.protocol.saml.SamlProtocol#buildAuthenticatedResponse
  -> protected Response
buildAuthenticatedResponse(AuthenticatedClientSessionModel clientSession,
String redirectUri, Document samlDocument, JaxrsSAML2BindingBuilder
bindingBuilder)
  -> samlDocument contains the SAML document populated by Keycloak BEFORE
it is signed and encrypted
  -> clientSession gives you access to the current user etc.
  -> add custom SAML processing as needed
  -> call return super.buildAuthenticatedResponse(clientSession,
redirectUri, samlDocument, bindingBuilder); at the end, this potentially
signs and encrypts the SAML document
- Make keycloak aware of that extension... (tricky, read on)

Note that the "CustomSamlProtocolFactory" uses the same provider ID "saml"
as the "SamlProtocolFactory" does -> this leads to a name clash...
unfortunantely Keycloak currently doesn't support to have custom
implementations for core protocols.

Here comes the hack:
In order to let Keycloak pickup my "CustomSamlProtocolFactory"
implementation before it's own "SamlProtocolFactory", I needed to copy the
provider jar to
`$KEYCLOAK_HOME/modules/system/layers/keycloak/org/keycloak/keycloak-services/main`
and change the keycloak-services module.xml to:

    <resources>
        <resource-root
path="simple-custom-saml-protocol-1.0.0.0-SNAPSHOT.jar"/> <!-- added this
-->
        <resource-root path="keycloak-services-4.8.3.Final.jar"/>
    </resources>

With that change Keycloak picked up my custom "CustomSamlProtocolFactory"
first and I could adjust the SAML document as required.

Btw. you might be wondering why I didn't just create a new
"CustomSamlProtocolFactory" with a dedicated provider-id like "saml-custom"
and use that for clients. We'll that was my first try, but unfortunately
this is currently not possible, since the Admin-Console
UI for the SAML configuration seems to be hard-coded against the
provider-id "saml"... so I'm stuck with the current "solution".
It would be great if the "LoginProtocolFactory" SPI would support custom
protocol implementations, while reusing protocol configuration options and
the admin-console ui.

Cheers,
Thomas

Ps.:

FYI this is what a SAMLAttributeStatementMapper allows you to do currently
(Keycloak 4.8.3.Final):

// SimpleSamlMapper
...
    @Override
    public void transformAttributeStatement(AttributeStatementType
attributeStatement, ProtocolMapperModel mappingModel, KeycloakSession
session, UserSessionModel userSession, AuthenticatedClientSessionModel
clientSession) {

        // transform attributeStatement here
        LOGGER.infof("transformAttributeStatement");

        AttributeType bubu = new AttributeType("bubu");
        bubu.setFriendlyName("FriendlyBubu");

bubu.setNameFormat("urn:oasis:names:tc:SAML:2.0:attrname-format:basic");
        bubu.setName("Bubu");

        bubu.addAttributeValue("Object allowed but only Strings or
NameIDType supported here...");
        // see: bottom of
org.keycloak.saml.processing.core.saml.v2.writers.BaseWriter.writeAttributeTypeWithoutRootTag

        // would be great to have support for AttributeValue type
customizations, even if value is String.
        // would be great to have support for complex XML element as values

        attributeStatement.addAttribute(new
AttributeStatementType.ASTChoiceType(bubu));
    }
...

// CustomSamlProtocolFactory
import org.keycloak.models.KeycloakSession;
import org.keycloak.protocol.LoginProtocol;
import org.keycloak.protocol.saml.SamlProtocolFactory;

public class CustomSamlProtocolFactory extends SamlProtocolFactory {

    public LoginProtocol create(KeycloakSession session) {
        return (new CustomSamlProtocol()).setSession(session);
    }
}


// CustomSamlProtocol
import org.keycloak.models.AuthenticatedClientSessionModel;
import org.keycloak.protocol.saml.JaxrsSAML2BindingBuilder;
import org.keycloak.protocol.saml.SamlProtocol;
import org.keycloak.saml.common.constants.JBossSAMLURIConstants;
import org.keycloak.saml.common.exceptions.ConfigurationException;
import org.keycloak.saml.common.exceptions.ProcessingException;
import org.w3c.dom.Document;
import org.w3c.dom.Element;

import javax.ws.rs.core.Response;
import java.io.IOException;

public class CustomSamlProtocol extends SamlProtocol {

    protected Response
buildAuthenticatedResponse(AuthenticatedClientSessionModel clientSession,
String redirectUri, Document samlDocument, JaxrsSAML2BindingBuilder
bindingBuilder) throws ConfigurationException, ProcessingException,
IOException {

        // TODO use Keycloak provider information from this.session
        Element assertionElement = (Element)
samlDocument.getElementsByTagNameNS(JBossSAMLURIConstants.ASSERTION_NSURI.get(),
"Assertion").item(0);

        Element attributeStatementElement = (Element)
assertionElement.getElementsByTagNameNS(JBossSAMLURIConstants.ASSERTION_NSURI.get(),
"AttributeStatement").item(0);
        if (attributeStatementElement == null) {
            attributeStatementElement =
samlDocument.createElementNS(JBossSAMLURIConstants.ASSERTION_NSURI.get(),
"AttributeStatement");
            assertionElement.appendChild(attributeStatementElement);
        }
        // TODO pull information from user attributes

attributeStatementElement.appendChild(newSamlAttributeElement(samlDocument,
null, "Role", JBossSAMLURIConstants.ATTRIBUTE_FORMAT_BASIC.get(), "dummy",
"xsd:string"));

        // see: http://oid-info.com/get/

attributeStatementElement.appendChild(newSamlAttributeElement(samlDocument,
"XSPA Organization ID",
"urn:oasis:names:tc:xspa:1.0:subject:organization-id",
JBossSAMLURIConstants.ATTRIBUTE_FORMAT_URI.get(),
"urn:oid:1.2.3.4.5.6.7.8.9.10.11.12", "xsd:anyURI"));

        Element roleElement =
samlDocument.createElementNS("urn:hl7-org:v3", "Role");
        roleElement.setAttribute("code", "PRA");
        roleElement.setAttribute("codeSystem",
"1.2.3.4.5.6.7.8.9.10.11.12");
        roleElement.setAttribute("codeSystemName",
"IHEXDShealthcareFacilityTypeCode");
        roleElement.setAttribute("displayName", "Doctor's office");

attributeStatementElement.appendChild(newSamlAttributeElement(samlDocument,
"Acme Role", "urn:oasis:names:tc:xacml:2.0:subject:role",
JBossSAMLURIConstants.ATTRIBUTE_FORMAT_URI.get(), roleElement,
"xsd:anyType"));

        return super.buildAuthenticatedResponse(clientSession, redirectUri,
samlDocument, bindingBuilder);
    }

    private Element newSamlAttributeElement(Document samlDocument, String
friendlyName, String name, String nameFormat, Object value, String type) {

        Element targetSamlAttributeElement =
samlDocument.createElementNS(JBossSAMLURIConstants.ASSERTION_NSURI.get(),
"Attribute");

        if (friendlyName != null) {
            targetSamlAttributeElement.setAttribute("FriendlyName",
friendlyName);
        }
        targetSamlAttributeElement.setAttribute("Name", name);
        if (nameFormat != null) {
            targetSamlAttributeElement.setAttribute("NameFormat",
nameFormat);
        }

        Element samlAttributeValue =
samlDocument.createElementNS(JBossSAMLURIConstants.ASSERTION_NSURI.get(),
"AttributeValue");
        samlAttributeValue.setAttribute("xmlns:xsi", "
http://www.w3.org/2001/XMLSchema-instance");
        samlAttributeValue.setAttribute("xsi:type", type);
        targetSamlAttributeElement.appendChild(samlAttributeValue);

        if (value instanceof String) {
            samlAttributeValue.setTextContent((String) value);
        } else if (value instanceof Element) {
            samlAttributeValue.appendChild((Element) value);
        } else if (value != null) {
            samlAttributeValue.setTextContent(value.toString());
        } else {
            samlAttributeValue.setTextContent(String.valueOf(value));
        }

        return targetSamlAttributeElement;
    }
}