Merge pull request #11278 from vitlibar/fix-crash-set-default-role-with-wrong-args

Fix crash when SET DEFAULT ROLE is called with wrong arguments.

src/Access/ExtendedRoleSet.cpp

@@ -68,15 +68,27 @@ void ExtendedRoleSet::init(const ASTExtendedRoleSet & ast, const AccessControlMa
 {
     all = ast.all;
 
-    auto name_to_id = [id_mode{ast.id_mode}, manager](const String & name) -> UUID
+    auto name_to_id = [&ast, manager](const String & name) -> UUID
     {
-        if (id_mode)
+        if (ast.id_mode)
             return parse<UUID>(name);
         assert(manager);
-        auto id = manager->find<User>(name);
-        if (id)
-            return *id;
-        return manager->getID<Role>(name);
+        if (ast.can_contain_users && ast.can_contain_roles)
+        {
+            auto id = manager->find<User>(name);
+            if (id)
+                return *id;
+            return manager->getID<Role>(name);
+        }
+        else if (ast.can_contain_users)
+        {
+            return manager->getID<User>(name);
+        }
+        else
+        {
+            assert(ast.can_contain_roles);
+            return manager->getID<Role>(name);
+        }
     };
 
     if (!ast.names.empty() && !all)

src/Interpreters/InterpreterSetRoleQuery.cpp

@@ -62,7 +62,7 @@ void InterpreterSetRoleQuery::setRole(const ASTSetRoleQuery & query)
 
 void InterpreterSetRoleQuery::setDefaultRole(const ASTSetRoleQuery & query)
 {
-    context.checkAccess(AccessType::CREATE_USER | AccessType::DROP_USER);
+    context.checkAccess(AccessType::ALTER_USER);
 
     auto & access_control = context.getAccessControlManager();
     std::vector<UUID> to_users = ExtendedRoleSet{*query.to_users, access_control, context.getUserID()}.getMatchingIDs(access_control);

src/Parsers/ASTExtendedRoleSet.h

@@ -15,7 +15,10 @@ public:
     bool all = false;
     Strings except_names;
     bool except_current_user = false;
-    bool id_mode = false;  /// If true then `names` and `except_names` keeps UUIDs, not names.
+
+    bool id_mode = false;          /// true if `names` and `except_names` keep UUIDs, not names.
+    bool can_contain_roles = true; /// true if this set can contain names of roles.
+    bool can_contain_users = true; /// true if this set can contain names of users.
 
     bool empty() const { return names.empty() && !current_user && !all; }
     void replaceCurrentUserTagWithName(const String & current_user_name);

src/Parsers/ParserCreateUserQuery.cpp

@@ -201,6 +201,7 @@ namespace
                 return false;
 
             default_roles = typeid_cast<std::shared_ptr<ASTExtendedRoleSet>>(ast);
+            default_roles->can_contain_users = false;
             return true;
         });
     }

src/Parsers/ParserSetRoleQuery.cpp

@@ -18,6 +18,7 @@ namespace
                 return false;
 
             roles = typeid_cast<std::shared_ptr<ASTExtendedRoleSet>>(ast);
+            roles->can_contain_users = false;
             return true;
         });
     }
@@ -34,6 +35,7 @@ namespace
                 return false;
 
             to_users = typeid_cast<std::shared_ptr<ASTExtendedRoleSet>>(ast);
+            to_users->can_contain_roles = false;
             return true;
         });
     }

tests/integration/test_default_role/test.py

@@ -0,0 +1,77 @@
+import pytest
+from helpers.cluster import ClickHouseCluster
+from helpers.test_tools import TSV
+import re
+
+cluster = ClickHouseCluster(__file__)
+instance = cluster.add_instance('instance')
+
+
+@pytest.fixture(scope="module", autouse=True)
+def started_cluster():
+    try:
+        cluster.start()
+        
+        instance.query("CREATE USER john")
+        instance.query("CREATE ROLE rx")
+        instance.query("CREATE ROLE ry")
+
+        yield cluster
+
+    finally:
+        cluster.shutdown()
+
+
+@pytest.fixture(autouse=True)
+def reset_users_and_roles():
+    instance.query("CREATE USER OR REPLACE john")
+    yield
+
+
+def test_set_default_roles():
+    assert instance.query("SHOW CURRENT ROLES", user="john") == ""
+
+    instance.query("GRANT rx, ry TO john")
+    assert instance.query("SHOW CURRENT ROLES", user="john") == TSV( [['rx', 0, 1], ['ry', 0, 1]] )
+
+    instance.query("SET DEFAULT ROLE NONE TO john")
+    assert instance.query("SHOW CURRENT ROLES", user="john") == ""
+
+    instance.query("SET DEFAULT ROLE rx TO john")
+    assert instance.query("SHOW CURRENT ROLES", user="john") == TSV( [['rx', 0, 1]] )
+
+    instance.query("SET DEFAULT ROLE ry TO john")
+    assert instance.query("SHOW CURRENT ROLES", user="john") == TSV( [['ry', 0, 1]] )
+
+    instance.query("SET DEFAULT ROLE ALL TO john")
+    assert instance.query("SHOW CURRENT ROLES", user="john") == TSV( [['rx', 0, 1], ['ry', 0, 1]] )
+
+    instance.query("SET DEFAULT ROLE ALL EXCEPT rx TO john")
+    assert instance.query("SHOW CURRENT ROLES", user="john") == TSV( [['ry', 0, 1]] )
+
+
+def test_alter_user():
+    assert instance.query("SHOW CURRENT ROLES", user="john") == ""
+
+    instance.query("GRANT rx, ry TO john")
+    assert instance.query("SHOW CURRENT ROLES", user="john") == TSV( [['rx', 0, 1], ['ry', 0, 1]] )
+
+    instance.query("ALTER USER john DEFAULT ROLE NONE")
+    assert instance.query("SHOW CURRENT ROLES", user="john") == ""
+
+    instance.query("ALTER USER john DEFAULT ROLE rx")
+    assert instance.query("SHOW CURRENT ROLES", user="john") == TSV( [['rx', 0, 1]] )
+
+    instance.query("ALTER USER john DEFAULT ROLE ALL")
+    assert instance.query("SHOW CURRENT ROLES", user="john") == TSV( [['rx', 0, 1], ['ry', 0, 1]] )
+
+    instance.query("ALTER USER john DEFAULT ROLE ALL EXCEPT rx")
+    assert instance.query("SHOW CURRENT ROLES", user="john") == TSV( [['ry', 0, 1]] )
+
+
+def test_wrong_set_default_role():
+    assert "There is no user `rx`" in instance.query_and_get_error("SET DEFAULT ROLE NONE TO rx")
+    assert "There is no user `ry`" in instance.query_and_get_error("SET DEFAULT ROLE rx TO ry")
+    assert "There is no role `john`" in instance.query_and_get_error("SET DEFAULT ROLE john TO john")
+    assert "There is no role `john`" in instance.query_and_get_error("ALTER USER john DEFAULT ROLE john")
+    assert "There is no role `john`" in instance.query_and_get_error("ALTER USER john DEFAULT ROLE ALL EXCEPT john")