From ca55edf22cf92daf1ecdafe87fa03c75082f1b0f Mon Sep 17 00:00:00 2001
From: UnicornChance <chance@defenseunicorns.com>
Date: Wed, 28 Aug 2024 08:03:19 -0600
Subject: [PATCH] fix: restrict colon in group

---
 .../plugin/CustomAWSSAMLGroupMapper.java      | 16 +++++++++++++---
 .../plugin/CustomAWSSAMLGroupMapperTest.java  | 19 +++++++++++++++++++
 2 files changed, 32 insertions(+), 3 deletions(-)

diff --git a/src/plugin/src/main/java/com/defenseunicorns/uds/keycloak/plugin/CustomAWSSAMLGroupMapper.java b/src/plugin/src/main/java/com/defenseunicorns/uds/keycloak/plugin/CustomAWSSAMLGroupMapper.java
index 334fac1..c6d823a 100644
--- a/src/plugin/src/main/java/com/defenseunicorns/uds/keycloak/plugin/CustomAWSSAMLGroupMapper.java
+++ b/src/plugin/src/main/java/com/defenseunicorns/uds/keycloak/plugin/CustomAWSSAMLGroupMapper.java
@@ -33,7 +33,8 @@ public class CustomAWSSAMLGroupMapper extends AbstractSAMLProtocolMapper impleme
         property.setName(AttributeStatementHelper.SAML_ATTRIBUTE_NAME);
         property.setLabel("Group attribute name");
         property.setDefaultValue("member");
-        property.setHelpText("Name of the SAML attribute you want to put your groups into.  i.e. 'member', 'memberOf'.");
+        property.setHelpText(
+                "Name of the SAML attribute you want to put your groups into.  i.e. 'member', 'memberOf'.");
         configProperties.add(property);
         property = new ProviderConfigProperty();
         property.setName(AttributeStatementHelper.FRIENDLY_NAME);
@@ -80,7 +81,7 @@ public List<ProviderConfigProperty> getConfigProperties() {
 
     @Override
     public void transformAttributeStatement(AttributeStatementType attributeStatement, ProtocolMapperModel mappingModel,
-                                            KeycloakSession session, UserSessionModel userSession, AuthenticatedClientSessionModel clientSession) {
+            KeycloakSession session, UserSessionModel userSession, AuthenticatedClientSessionModel clientSession) {
         UserModel user = userSession.getUser();
 
         // Retrieve the user's groups
@@ -92,11 +93,20 @@ public void transformAttributeStatement(AttributeStatementType attributeStatemen
                     .map(ModelToRepresentation::buildGroupPath)
                     .collect(Collectors.toList());
 
+            // Check for any group paths containing a colon
+            for (String groupPath : groupPaths) {
+                if (groupPath.contains(":")) {
+                    throw new IllegalArgumentException(
+                            "Group name contains invalid character ':'. Group name: " + groupPath);
+                }
+            }
+
             // Concatenate the group paths into a single string separated by colons
             String groupsString = String.join(":", groupPaths);
 
             // Create a new SAML attribute
-            AttributeType attribute = new AttributeType(mappingModel.getConfig().get(AttributeStatementHelper.SAML_ATTRIBUTE_NAME));
+            AttributeType attribute = new AttributeType(
+                    mappingModel.getConfig().get(AttributeStatementHelper.SAML_ATTRIBUTE_NAME));
             attribute.setNameFormat(JBossSAMLURIConstants.ATTRIBUTE_FORMAT_BASIC.get());
             attribute.addAttributeValue(groupsString);
             attributeStatement.addAttribute(new ASTChoiceType(attribute));
diff --git a/src/plugin/src/test/java/com/defenseunicorns/uds/keycloak/plugin/CustomAWSSAMLGroupMapperTest.java b/src/plugin/src/test/java/com/defenseunicorns/uds/keycloak/plugin/CustomAWSSAMLGroupMapperTest.java
index d1fb244..d038d66 100644
--- a/src/plugin/src/test/java/com/defenseunicorns/uds/keycloak/plugin/CustomAWSSAMLGroupMapperTest.java
+++ b/src/plugin/src/test/java/com/defenseunicorns/uds/keycloak/plugin/CustomAWSSAMLGroupMapperTest.java
@@ -18,7 +18,10 @@
 import java.util.Map;
 import java.util.stream.Stream;
 
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.fail;
 import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.argThat;
 import static org.mockito.Mockito.*;
 
 public class CustomAWSSAMLGroupMapperTest {
@@ -91,6 +94,22 @@ public void testTransformAttributeStatement_singleGroupParent() {
         }));
     }
 
+   @Test
+    public void testTransformAttributeStatement_groupNameWithColonThrowsException() {
+        GroupModel group = mock(GroupModel.class);
+        when(group.getName()).thenReturn("Admin:IT");
+        when(group.getParent()).thenReturn(null);
+
+        when(mockUser.getGroupsStream()).thenReturn(Stream.of(group));
+
+        try {
+            mapper.transformAttributeStatement(mockAttributeStatement, mockMappingModel, mockSession, mockUserSession, mockClientSession);
+            fail("Expected IllegalArgumentException to be thrown");
+        } catch (IllegalArgumentException e) {
+            assertEquals("Group name contains invalid character ':'. Group name: /Admin:IT", e.getMessage());
+        }
+    }
+
     @Test
     public void testTransformAttributeStatement_multipleGroupsNoParent() {
         GroupModel group1 = mock(GroupModel.class);