Allow replacement string to substitute against original user in impersonation

core/trino-main/src/test/java/io/trino/security/TestFileBasedSystemAccessControl.java

@@ -19,6 +19,7 @@
 import io.trino.plugin.base.security.DefaultSystemAccessControl;
 import io.trino.plugin.base.security.FileBasedSystemAccessControl;
 import io.trino.spi.QueryId;
+import io.trino.spi.TrinoException;
 import io.trino.spi.connector.CatalogSchemaName;
 import io.trino.spi.connector.SchemaTableName;
 import io.trino.spi.security.AccessDeniedException;
@@ -102,6 +103,26 @@ public void testCanImpersonateUserOperations()
         assertThatThrownBy(() -> accessControlManager.checkCanImpersonateUser(Identity.ofUser("invalid-other"), "test"))
                 .isInstanceOf(AccessDeniedException.class);
 
+        accessControlManager.checkCanImpersonateUser(Identity.ofUser("svc_tenant"), "svc_tenant_prod");
+        assertThatThrownBy(() -> accessControlManager.checkCanImpersonateUser(Identity.ofUser("svc_tenant"), "svc_tenant_other"))
+                .isInstanceOf(AccessDeniedException.class);
+        assertThatThrownBy(() -> accessControlManager.checkCanImpersonateUser(Identity.ofUser("svc_tenant"), "svc_other_prod"))
+                .isInstanceOf(AccessDeniedException.class);
+
+        accessControlManager.checkCanImpersonateUser(Identity.ofUser("external_corp_dept"), "internal-dept-corp-sandbox");
+        assertThatThrownBy(() -> accessControlManager.checkCanImpersonateUser(Identity.ofUser("external_corp_dept"), "internal-corp-dept-sandbox"))
+                .isInstanceOf(AccessDeniedException.class);
+        assertThatThrownBy(() -> accessControlManager.checkCanImpersonateUser(Identity.ofUser("external_corp_dept"), "invalid"))
+                .isInstanceOf(AccessDeniedException.class);
+
+        accessControlManager.checkCanImpersonateUser(Identity.ofUser("missing_replacement_group"), "anything");
+        assertThatThrownBy(() -> accessControlManager.checkCanImpersonateUser(Identity.ofUser("incorrect_number_of_replacements_groups_group"), "$2_group_prod"))
+                .isInstanceOf(TrinoException.class)
+                .hasMessageContaining("new_user in impersonation rule refers to a capturing group that does not exist in original_user");
+        assertThatThrownBy(() -> accessControlManager.checkCanImpersonateUser(Identity.ofUser("incorrect_number_of_replacements_groups_group"), "group_group_prod"))
+                .isInstanceOf(TrinoException.class)
+                .hasMessageContaining("new_user in impersonation rule refers to a capturing group that does not exist in original_user");
+
         AccessControlManager accessControlManagerWithPrincipal = newAccessControlManager(transactionManager, "catalog_principal.json");
         accessControlManagerWithPrincipal.checkCanImpersonateUser(Identity.ofUser("anything"), "anythingElse");
     }

core/trino-main/src/test/resources/catalog_impersonation.json

@@ -21,6 +21,26 @@
         {
             "original_user": ".*",
             "new_user": "test"
+        },
+        {
+            "original_user": "svc_(.*)",
+            "new_user": "svc_$1_prod",
+            "allow": true
+        },
+        {
+            "original_user": "external_(.*)_(.*)",
+            "new_user": "internal-$2-$1-sandbox",
+            "allow": true
+        },
+        {
+            "original_user": "missing_replacement_(.*)",
+            "new_user": ".*",
+            "allow": true
+        },
+        {
+            "original_user": "incorrect_number_of_replacements_groups_(.*)",
+            "new_user": "$2_$1_prod",
+            "allow": true
         }
     ]
 }

docs/src/main/sphinx/security/file-system-access-control.rst

@@ -544,7 +544,9 @@ Each impersonation rule is composed of the following fields:
 * ``original_role`` (optional): regex to match against role names of the
   requesting impersonation. Defaults to ``.*``.
 * ``new_user`` (required): regex to match against the user that will be
-  impersonated.
+  impersonated. May contain references to subsequences captured during the match
+  against *original_user*, and each reference will be replaced by the result of
+  evaluating the corresponding group respectively.
 * ``allow`` (optional): boolean indicating if the authentication should be
   allowed. Defaults to ``true``.
 
@@ -553,7 +555,9 @@ The impersonation rules are a bit different than the other rules: The attribute
 Doing so it was possible to make the attribute ``allow`` optional.
 
 The following example allows the ``admin`` role, to impersonate any user, except
-for ``bob``. It also allows any user to impersonate the ``test`` user:
+for ``bob``. It also allows any user to impersonate the ``test`` user. It also
+allows a user in the form ``team_backend`` to impersonate the
+``team_backend_sandbox`` user, but not arbitrary users:
 
 .. literalinclude:: user-impersonation.json
     :language: json

docs/src/main/sphinx/security/user-impersonation.json

@@ -12,6 +12,11 @@
         {
             "original_user": ".*",
             "new_user": "test"
+        },
+        {
+            "original_user": "team_(.*)",
+            "new_user": "team_$1_sandbox",
+            "allow": true
         }
     ]
 }

lib/trino-plugin-toolkit/src/main/java/io/trino/plugin/base/security/ImpersonationRule.java

@@ -16,12 +16,15 @@
 import com.fasterxml.jackson.annotation.JsonAlias;
 import com.fasterxml.jackson.annotation.JsonCreator;
 import com.fasterxml.jackson.annotation.JsonProperty;
+import io.trino.spi.TrinoException;
 
 import java.util.Optional;
 import java.util.Set;
+import java.util.regex.Matcher;
 import java.util.regex.Pattern;
 
 import static com.google.common.base.MoreObjects.firstNonNull;
+import static io.trino.spi.StandardErrorCode.CONFIGURATION_INVALID;
 import static java.lang.Boolean.TRUE;
 import static java.util.Objects.requireNonNull;
 
@@ -47,9 +50,28 @@ public ImpersonationRule(
 
     public Optional<Boolean> match(String originalUser, Set<String> originalRoles, String newUser)
     {
+        Pattern replacedNewUserPattern = newUserPattern;
+        if (originalUserPattern.isPresent()) {
+            Matcher matcher = originalUserPattern.get().matcher(originalUser);
+            if (matcher.matches()) {
+                StringBuilder stringBuilder = new StringBuilder();
+                try {
+                    matcher.appendReplacement(stringBuilder, newUserPattern.pattern());
+                }
+                catch (IndexOutOfBoundsException e) {
+                    throw new TrinoException(
+                            CONFIGURATION_INVALID,
+                            "new_user in impersonation rule refers to a capturing group that does not exist in original_user",
+                            e);
+                }
+
+                replacedNewUserPattern = Pattern.compile(stringBuilder.toString());
+            }
+        }
+
         if (originalUserPattern.map(regex -> regex.matcher(originalUser).matches()).orElse(true) &&
                 originalRolePattern.map(regex -> originalRoles.stream().anyMatch(role -> regex.matcher(role).matches())).orElse(true) &&
-                newUserPattern.matcher(newUser).matches()) {
+                replacedNewUserPattern.matcher(newUser).matches()) {
             return Optional.of(allow);
         }
         return Optional.empty();