Fixed ZMSUtils to correctly determine PrincipalType (#2556)

Signed-off-by: Shimaoka Shuya <sshimaok@lycorp.co.jp> Co-authored-by: Shimaoka Shuya <sshimaok@lycorp.co.jp>
AthenZ · Mar 22, 2024 · e9eac1e · e9eac1e
1 parent 2cb4cfb
commit e9eac1e
Show file tree

Hide file tree

Showing 5 changed files with 70 additions and 4 deletions.
diff --git a/libs/java/auth_core/src/main/java/com/yahoo/athenz/auth/util/StringUtils.java b/libs/java/auth_core/src/main/java/com/yahoo/athenz/auth/util/StringUtils.java
@@ -100,4 +100,14 @@ public static boolean requestUriMatch(String uri, Set<String> uriSet,
         }
         return false;
     }
+
+    public static int countMatches(final CharSequence str, final char ch) {
+        int count = 0;
+        for (int i = 0; i < str.length(); i++) {
+            if (str.charAt(i) == ch) {
+                count++;
+            }
+        }
+        return count;
+    }
 }
diff --git a/libs/java/auth_core/src/test/java/com/yahoo/athenz/auth/util/StringUtilsTest.java b/libs/java/auth_core/src/test/java/com/yahoo/athenz/auth/util/StringUtilsTest.java
@@ -98,4 +98,11 @@ public void testRequestUriMatch() {
         assertFalse(StringUtils.requestUriMatch("/zts/v1/domain/athenz/service/zms", uriSet, uriList));
         assertTrue(StringUtils.requestUriMatch("/zts/v1/domain/athenz/service/zms/publickey/zms1", uriSet, uriList));
     }
+
+    @Test
+    public void testCountMatches() {
+        assertEquals(StringUtils.countMatches("user", '.'), 0);
+        assertEquals(StringUtils.countMatches("user.joe", '.'), 1);
+        assertEquals(StringUtils.countMatches("home.joe.service", '.'), 2);
+    }
 }
diff --git a/servers/zms/src/main/java/com/yahoo/athenz/zms/utils/ZMSUtils.java b/servers/zms/src/main/java/com/yahoo/athenz/zms/utils/ZMSUtils.java
@@ -259,7 +259,7 @@ public static Principal.Type principalType(final String memberName, final String
 
         if (ZMSUtils.isUserDomainPrincipal(memberName, userDomainPrefix, addlUserCheckDomainPrefixList)) {
             return Principal.Type.USER;
-        } else if (memberName.startsWith(headlessUserDomainPrefix)) {
+        } else if (ZMSUtils.isHeadlessUserDomainPrincipal(memberName, headlessUserDomainPrefix)) {
             return Principal.Type.USER_HEADLESS;
         } else if (memberName.contains(AuthorityConsts.GROUP_SEP)) {
             return Principal.Type.GROUP;
@@ -271,20 +271,24 @@ public static Principal.Type principalType(final String memberName, final String
     public static boolean isUserDomainPrincipal(final String memberName, final String userDomainPrefix,
             final List<String> addlUserCheckDomainPrefixList) {
 
-        if (memberName.startsWith(userDomainPrefix)) {
+        if (memberName.startsWith(userDomainPrefix) && StringUtils.countMatches(memberName, '.') == 1) {
             return true;
         }
 
         if (addlUserCheckDomainPrefixList != null) {
             for (String prefix : addlUserCheckDomainPrefixList) {
-                if (memberName.startsWith(prefix)) {
+                if (memberName.startsWith(prefix) && StringUtils.countMatches(memberName, '.') == 1) {
                     return true;
                 }
             }
         }
 
         return false;
     }
+
+    public static boolean isHeadlessUserDomainPrincipal(final String memberName, final String headlessUserDomainPrefix) {
+        return memberName.startsWith(headlessUserDomainPrefix) && StringUtils.countMatches(memberName, '.') == 1;
+    }
 
     public static String extractObjectName(String domainName, String fullName, String objType) {
 

diff --git a/servers/zms/src/test/java/com/yahoo/athenz/zms/ZMSDeleteUserTest.java b/servers/zms/src/test/java/com/yahoo/athenz/zms/ZMSDeleteUserTest.java
@@ -29,7 +29,6 @@
 
 import static com.yahoo.athenz.common.messaging.DomainChangeMessage.ObjectType.*;
 import static org.testng.Assert.*;
-import static org.testng.Assert.assertTrue;
 
 public class ZMSDeleteUserTest {
 
@@ -82,6 +81,12 @@ public void testDeleteUser() {
                 "Test SubDomain21", "testOrg", zmsTestInitializer.getAdminUser());
         zmsImpl.postSubDomain(ctx, "user.jack", auditRef, subDom2);
 
+        ServiceIdentity service2 = new ServiceIdentity().setName(ResourceUtils.serviceResourceName("user.jack.sub1", "api"));
+        zmsImpl.putServiceIdentity(ctx, "user.jack.sub1", "api", auditRef, false, service2);
+
+        ServiceIdentity service3 = new ServiceIdentity().setName(ResourceUtils.serviceResourceName("user.jack.sub1", "service"));
+        zmsImpl.putServiceIdentity(ctx, "user.jack.sub1", "service", auditRef, false, service3);
+
         Role role1 = zmsTestInitializer.createRoleObject(domainName, "role1", null,
                 "user.joe", "user.jack.sub1.service");
         zmsImpl.putRole(ctx, domainName, "role1", auditRef, false, role1);

diff --git a/servers/zms/src/test/java/com/yahoo/athenz/zms/utils/ZMSUtilsTest.java b/servers/zms/src/test/java/com/yahoo/athenz/zms/utils/ZMSUtilsTest.java
@@ -133,6 +133,46 @@ public void testGetAudtLogMsgBuilder() {
         assertNotNull(msgBuilder);
         assertTrue(msgBuilder.who().contains("who-roles=[role1, role2]"), msgBuilder.who());
     }
+
+    @Test
+    public void testPrincipalType() {
+        // Set different strings between user and home domain
+        String userDomain = "user";
+        String userDomain2 = "user2";
+        String homeDomain = "home";
+        String headlessDomain = "headless";
+        String topLevelDomain = "athenz";
+        String groupSep = ":group";
+        List<String> addlUserCheckDomainPrefixList = Arrays.asList(userDomain2);
+
+        // GROUP
+        assertEquals(ZMSUtils.principalType(homeDomain + ".joe" + groupSep + ".test-group", userDomain, addlUserCheckDomainPrefixList, headlessDomain), Principal.Type.GROUP);
+        assertEquals(ZMSUtils.principalType(topLevelDomain + groupSep + ".test-group", userDomain, addlUserCheckDomainPrefixList, headlessDomain), Principal.Type.GROUP);
+        // USER
+        assertEquals(ZMSUtils.principalType(userDomain + ".joe", userDomain, addlUserCheckDomainPrefixList, headlessDomain), Principal.Type.USER);
+        assertEquals(ZMSUtils.principalType(userDomain2 + ".joe", userDomain, addlUserCheckDomainPrefixList, headlessDomain), Principal.Type.USER);
+        // USER_HEADLESS
+        assertEquals(ZMSUtils.principalType(headlessDomain + ".joe", userDomain, addlUserCheckDomainPrefixList, headlessDomain), Principal.Type.USER_HEADLESS);
+        // SERVICE
+        assertEquals(ZMSUtils.principalType(topLevelDomain + ".test-service", userDomain, addlUserCheckDomainPrefixList, headlessDomain), Principal.Type.SERVICE);
+        assertEquals(ZMSUtils.principalType(homeDomain + ".joe" + ".test-service", userDomain, addlUserCheckDomainPrefixList, headlessDomain), Principal.Type.SERVICE);
+
+        // Set same strings between user and home domain.
+        userDomain = "personal";
+        homeDomain = userDomain;
+
+        // GROUP
+        assertEquals(ZMSUtils.principalType(homeDomain + ".joe" + groupSep + ".test-group", userDomain, addlUserCheckDomainPrefixList, headlessDomain), Principal.Type.GROUP);
+        assertEquals(ZMSUtils.principalType(topLevelDomain + groupSep + ".test-group", userDomain, addlUserCheckDomainPrefixList, headlessDomain), Principal.Type.GROUP);
+        // USER
+        assertEquals(ZMSUtils.principalType(userDomain + ".joe", userDomain, addlUserCheckDomainPrefixList, headlessDomain), Principal.Type.USER);
+        assertEquals(ZMSUtils.principalType(userDomain2 + ".joe", userDomain, addlUserCheckDomainPrefixList, headlessDomain), Principal.Type.USER);
+        // USER_HEADLESS
+        assertEquals(ZMSUtils.principalType(headlessDomain + ".joe", userDomain, addlUserCheckDomainPrefixList, headlessDomain), Principal.Type.USER_HEADLESS);
+        // SERVICE
+        assertEquals(ZMSUtils.principalType(topLevelDomain + ".test-service", userDomain, addlUserCheckDomainPrefixList, headlessDomain), Principal.Type.SERVICE);
+        assertEquals(ZMSUtils.principalType(homeDomain + ".joe" + ".test-service", userDomain, addlUserCheckDomainPrefixList, headlessDomain), Principal.Type.SERVICE);
+    }
 
     @Test
     public void testExtractRoleName() {