Skip to content

Commit

Permalink
[#6306] fix(authz): Fix the OOM of JdbcAuthorizationPlugin (#6312)
Browse files Browse the repository at this point in the history
### What changes were proposed in this pull request?

Fix the OOM of JdbcAuthorizationPlugin. In old implement, we will call
roleUpdate in the roleCreate. We will repeated call, they will cause
oom.

### Why are the changes needed?

Fix: #6306 

### Does this PR introduce _any_ user-facing change?

No.

### How was this patch tested?
Added a new UT.
  • Loading branch information
jerqi authored Jan 17, 2025
1 parent 85b919c commit 7945692
Show file tree
Hide file tree
Showing 2 changed files with 28 additions and 15 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -110,8 +110,13 @@ public Boolean onMetadataUpdated(MetadataObjectChange... changes) throws Runtime
@Override
public Boolean onRoleCreated(Role role) throws AuthorizationPluginException {
List<String> sqls = getCreateRoleSQL(role.name());
boolean createdNewly = false;
for (String sql : sqls) {
executeUpdateSQL(sql, "already exists");
createdNewly = executeUpdateSQL(sql, "already exists");
}

if (!createdNewly) {
return true;
}

if (role.securableObjects() != null) {
Expand Down Expand Up @@ -140,7 +145,6 @@ public Boolean onRoleDeleted(Role role) throws AuthorizationPluginException {
@Override
public Boolean onRoleUpdated(Role role, RoleChange... changes)
throws AuthorizationPluginException {
onRoleCreated(role);
for (RoleChange change : changes) {
if (change instanceof RoleChange.AddSecurableObject) {
SecurableObject object = ((RoleChange.AddSecurableObject) change).getSecurableObject();
Expand Down Expand Up @@ -381,14 +385,15 @@ protected AuthorizationPluginException toAuthorizationPluginException(SQLExcepti
"JDBC authorization plugin fail to execute SQL, error code: %d", se.getErrorCode());
}

public void executeUpdateSQL(String sql, String ignoreErrorMsg) {
public boolean executeUpdateSQL(String sql, String ignoreErrorMsg) {
try (final Connection connection = getConnection()) {
try (final Statement statement = connection.createStatement()) {
statement.executeUpdate(sql);
return true;
}
} catch (SQLException se) {
if (ignoreErrorMsg != null && se.getMessage().contains(ignoreErrorMsg)) {
return;
return false;
}
LOG.error("JDBC authorization plugin exception: ", se);
throw toAuthorizationPluginException(se);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -74,9 +74,10 @@ public List<String> getSetOwnerSQL(
return Collections.emptyList();
}

public void executeUpdateSQL(String sql, String ignoreErrorMsg) {
public boolean executeUpdateSQL(String sql, String ignoreErrorMsg) {
Assertions.assertEquals(expectSQLs.get(currentSQLIndex), sql);
currentSQLIndex++;
return true;
}
};

Expand Down Expand Up @@ -148,23 +149,21 @@ public void testPermissionManagement() {

// Test metalake object and different role change
resetSQLIndex();
expectSQLs = Lists.newArrayList("CREATE ROLE tmp", "GRANT SELECT ON TABLE *.* TO ROLE tmp");
expectSQLs = Lists.newArrayList("GRANT SELECT ON TABLE *.* TO ROLE tmp");
SecurableObject metalakeObject =
SecurableObjects.ofMetalake("metalake", Lists.newArrayList(Privileges.SelectTable.allow()));
RoleChange roleChange = RoleChange.addSecurableObject("tmp", metalakeObject);
plugin.onRoleUpdated(role, roleChange);

resetSQLIndex();
expectSQLs = Lists.newArrayList("CREATE ROLE tmp", "REVOKE SELECT ON TABLE *.* FROM ROLE tmp");
expectSQLs = Lists.newArrayList("REVOKE SELECT ON TABLE *.* FROM ROLE tmp");
roleChange = RoleChange.removeSecurableObject("tmp", metalakeObject);
plugin.onRoleUpdated(role, roleChange);

resetSQLIndex();
expectSQLs =
Lists.newArrayList(
"CREATE ROLE tmp",
"REVOKE SELECT ON TABLE *.* FROM ROLE tmp",
"GRANT CREATE ON TABLE *.* TO ROLE tmp");
"REVOKE SELECT ON TABLE *.* FROM ROLE tmp", "GRANT CREATE ON TABLE *.* TO ROLE tmp");
SecurableObject newMetalakeObject =
SecurableObjects.ofMetalake("metalake", Lists.newArrayList(Privileges.CreateTable.allow()));
roleChange = RoleChange.updateSecurableObject("tmp", metalakeObject, newMetalakeObject);
Expand All @@ -175,7 +174,7 @@ public void testPermissionManagement() {
SecurableObject catalogObject =
SecurableObjects.ofCatalog("catalog", Lists.newArrayList(Privileges.SelectTable.allow()));
roleChange = RoleChange.addSecurableObject("tmp", catalogObject);
expectSQLs = Lists.newArrayList("CREATE ROLE tmp", "GRANT SELECT ON TABLE *.* TO ROLE tmp");
expectSQLs = Lists.newArrayList("GRANT SELECT ON TABLE *.* TO ROLE tmp");
plugin.onRoleUpdated(role, roleChange);

// Test schema object
Expand All @@ -184,8 +183,7 @@ public void testPermissionManagement() {
SecurableObjects.ofSchema(
catalogObject, "schema", Lists.newArrayList(Privileges.SelectTable.allow()));
roleChange = RoleChange.addSecurableObject("tmp", schemaObject);
expectSQLs =
Lists.newArrayList("CREATE ROLE tmp", "GRANT SELECT ON TABLE schema.* TO ROLE tmp");
expectSQLs = Lists.newArrayList("GRANT SELECT ON TABLE schema.* TO ROLE tmp");
plugin.onRoleUpdated(role, roleChange);

// Test table object
Expand All @@ -194,8 +192,18 @@ public void testPermissionManagement() {
SecurableObjects.ofTable(
schemaObject, "table", Lists.newArrayList(Privileges.SelectTable.allow()));
roleChange = RoleChange.addSecurableObject("tmp", tableObject);
expectSQLs =
Lists.newArrayList("CREATE ROLE tmp", "GRANT SELECT ON TABLE schema.table TO ROLE tmp");
expectSQLs = Lists.newArrayList("GRANT SELECT ON TABLE schema.table TO ROLE tmp");
plugin.onRoleUpdated(role, roleChange);

// Test the role with objects
resetSQLIndex();
role =
RoleEntity.builder()
.withId(-1L)
.withName("tmp")
.withSecurableObjects(Lists.newArrayList(tableObject))
.withAuditInfo(AuditInfo.EMPTY)
.build();
plugin.onRoleUpdated(role, roleChange);
}

Expand Down

0 comments on commit 7945692

Please sign in to comment.