refactor: set enable_experimental_rbac_check default 1 (#15436)

TCeason · web-flow · commit 0651ef8aedec · 2024-05-08T18:04:31.000+08:00
* refactor: set enable_experimental_rbac_check default 1

* forbiden drop user stage in binder

* fix clippy
diff --git a/src/query/service/src/interpreters/access/privilege_access.rs b/src/query/service/src/interpreters/access/privilege_access.rs
@@ -35,6 +35,7 @@ use databend_common_sql::plans::PresignAction;
 use databend_common_sql::plans::RewriteKind;
 use databend_common_sql::Planner;
 use databend_common_users::RoleCacheManager;
+use databend_common_users::UserApiProvider;
 
 use crate::interpreters::access::AccessChecker;
 use crate::sessions::QueryContext;
@@ -618,13 +619,49 @@ impl AccessChecker for PrivilegeAccess {
                 self.validate_db_access(&plan.catalog, &plan.database, UserPrivilegeType::Drop, plan.if_exists).await?;
             }
             Plan::UndropDatabase(_)
-            | Plan::DropUDF(_)
             | Plan::DropIndex(_)
             | Plan::DropTableIndex(_) => {
                 // undroptable/db need convert name to id. But because of drop, can not find the id. Upgrade Object to Database.
                 self.validate_access(&GrantObject::Global, UserPrivilegeType::Drop)
                     .await?;
             }
+            Plan::DropUDF(plan) => {
+                let udf_name = &plan.udf;
+                if !UserApiProvider::instance().exists_udf(&tenant, udf_name).await? && plan.if_exists {
+                    return Ok(());
+                }
+                if enable_experimental_rbac_check {
+                    let udf = HashSet::from([udf_name]);
+                    self.validate_udf_access(udf).await?;
+                } else {
+                    self.validate_access(&GrantObject::Global, UserPrivilegeType::Drop)
+                                        .await?;
+                }
+            }
+            Plan::DropStage(plan) => {
+                match UserApiProvider::instance().get_stage(&tenant, &plan.name).await {
+                    Ok(stage) => {
+                        if enable_experimental_rbac_check {
+                            let privileges = vec![UserPrivilegeType::Read, UserPrivilegeType::Write];
+                            for privilege in privileges {
+                                self.validate_stage_access(&stage, privilege).await?;
+                            }
+                        } else {
+                            self.validate_access(&GrantObject::Global, UserPrivilegeType::Super)
+                                .await?;
+                        }
+                    }
+                    Err(e) => {
+                        match e.code() {
+                            ErrorCode::UNKNOWN_STAGE if plan.if_exists =>
+                                {
+                                    return Ok(());
+                                }
+                            _ => return Err(e.add_message("error on validating stage access")),
+                        }
+                    }
+                }
+            }
             Plan::UseDatabase(plan) => {
                 let session = self.ctx.get_current_session();
                 if self.has_ownership(&session, &GrantObject::Database(catalog_name.clone(), plan.database.clone())).await? {
@@ -1017,7 +1054,6 @@ impl AccessChecker for PrivilegeAccess {
             | Plan::CreateCatalog(_)
             | Plan::DropCatalog(_)
             | Plan::CreateStage(_)
-            | Plan::DropStage(_)
             | Plan::CreateFileFormat(_)
             | Plan::DropFileFormat(_)
             | Plan::ShowFileFormats(_)
diff --git a/src/query/service/src/interpreters/common/grant.rs b/src/query/service/src/interpreters/common/grant.rs
@@ -78,9 +78,9 @@ pub async fn validate_grant_object_exists(
         }
         GrantObject::UDF(udf) => {
             if !UserApiProvider::instance().exists_udf(&tenant, udf).await? {
-                return Err(databend_common_exception::ErrorCode::UnknownStage(format!(
-                    "udf {udf} not exists"
-                )));
+                return Err(databend_common_exception::ErrorCode::UnknownFunction(
+                    format!("udf {udf} not exists"),
+                ));
             }
         }
         GrantObject::Stage(stage) => {
diff --git a/src/query/service/src/interpreters/interpreter_user_stage_drop.rs b/src/query/service/src/interpreters/interpreter_user_stage_drop.rs
@@ -14,7 +14,6 @@
 
 use std::sync::Arc;
 
-use databend_common_exception::ErrorCode;
 use databend_common_exception::Result;
 use databend_common_management::RoleApi;
 use databend_common_meta_app::principal::OwnershipObject;
@@ -62,13 +61,6 @@ impl Interpreter for DropUserStageInterpreter {
         let tenant = self.ctx.get_tenant();
         let user_mgr = UserApiProvider::instance();
 
-        // Check user stage.
-        if plan.name == "~" {
-            return Err(ErrorCode::StagePermissionDenied(
-                "user stage is not allowed to be dropped",
-            ));
-        }
-
         let stage = user_mgr.get_stage(&tenant, &plan.name).await;
         user_mgr
             .drop_stage(&tenant, &plan.name, plan.if_exists)
diff --git a/src/query/settings/src/settings_default.rs b/src/query/settings/src/settings_default.rs
@@ -649,7 +649,7 @@ impl DefaultSettings {
                     range: Some(SettingRange::String(vec!["rounding".into(), "truncating".into()])),
                 }),
                 ("enable_experimental_rbac_check", DefaultSettingValue {
-                    value: UserSettingValue::UInt64(0),
+                    value: UserSettingValue::UInt64(1),
                     desc: "experiment setting disables stage and udf privilege check(disable by default).",
                     mode: SettingMode::Both,
                     range: Some(SettingRange::Numeric(0..=1)),
diff --git a/src/query/sql/src/planner/binder/binder.rs b/src/query/sql/src/planner/binder/binder.rs
@@ -380,10 +380,17 @@ impl<'a> Binder {
             Statement::DropStage {
                 stage_name,
                 if_exists,
-            } => Plan::DropStage(Box::new(DropStagePlan {
+            } => {
+                // Check user stage.
+                if stage_name == "~" {
+                    return Err(ErrorCode::StagePermissionDenied(
+                        "user stage is not allowed to be dropped",
+                    ));
+                }
+                Plan::DropStage(Box::new(DropStagePlan {
                 if_exists: *if_exists,
                 name: stage_name.clone(),
-            })),
+            }))},
             Statement::RemoveStage { location, pattern } => {
                 self.bind_remove_stage(location, pattern).await?
             }
diff --git a/tests/suites/0_stateless/18_rbac/18_0001_udf_priv.sh b/tests/suites/0_stateless/18_rbac/18_0001_udf_priv.sh
@@ -7,8 +7,6 @@ echo "=== test UDF priv"
 export TEST_USER_PASSWORD="password"
 export TEST_USER_CONNECT="bendsql --user=test-user --password=password --host=${QUERY_MYSQL_HANDLER_HOST} --port ${QUERY_HTTP_HANDLER_PORT}"
 
-echo "set global enable_experimental_rbac_check=1" | $BENDSQL_CLIENT_CONNECT
-
 echo "drop user if exists 'test-user'" | $BENDSQL_CLIENT_CONNECT
 echo "DROP FUNCTION IF EXISTS f1;" |  $BENDSQL_CLIENT_CONNECT
 echo "DROP FUNCTION IF EXISTS f2;" |  $BENDSQL_CLIENT_CONNECT
@@ -141,4 +139,3 @@ echo "drop function if exists a;" | $BENDSQL_CLIENT_CONNECT
 echo "drop function if exists b;" | $BENDSQL_CLIENT_CONNECT
 echo "drop table if exists default.t;" | $BENDSQL_CLIENT_CONNECT
 echo "drop table if exists default.t2;" | $BENDSQL_CLIENT_CONNECT
-echo "unset enable_experimental_rbac_check" | $BENDSQL_CLIENT_CONNECT
diff --git a/tests/suites/0_stateless/18_rbac/18_0002_ownership_cover.sh b/tests/suites/0_stateless/18_rbac/18_0002_ownership_cover.sh
@@ -9,8 +9,6 @@ export TEST_USER_CONNECT="bendsql --user=owner --password=password --host=${QUER
 
 export TEST_TRANSFER_USER_CONNECT="bendsql --user=owner1 --password=password --host=${QUERY_MYSQL_HANDLER_HOST} --port ${QUERY_HTTP_HANDLER_PORT}"
 
-echo "set global enable_experimental_rbac_check=1" | $BENDSQL_CLIENT_CONNECT
-
 ## cleanup
 echo "drop database if exists d_0002;" | $BENDSQL_CLIENT_CONNECT
 echo "drop user if exists '${TEST_USER_NAME}'" | $BENDSQL_CLIENT_CONNECT
@@ -72,13 +70,14 @@ echo "select * from d_0002.t" | $TEST_USER_CONNECT
 
 ## cleanup
 echo "drop database d_0002;" | $BENDSQL_CLIENT_CONNECT
-echo "drop stage hello;" | $BENDSQL_CLIENT_CONNECT
-echo "drop function a;" | $BENDSQL_CLIENT_CONNECT
+echo "drop stage hello;" | $TEST_TRANSFER_USER_CONNECT
+echo "drop stage if exists noexistshello;" | $TEST_TRANSFER_USER_CONNECT
+echo "drop function a;" | $TEST_TRANSFER_USER_CONNECT
+echo "drop function if exists noexistsa;" | $TEST_TRANSFER_USER_CONNECT
 echo "drop user '${TEST_USER_NAME}'" | $BENDSQL_CLIENT_CONNECT
 echo "drop user 'owner1'" | $BENDSQL_CLIENT_CONNECT
 echo "drop role 'r_0002_1'" | $BENDSQL_CLIENT_CONNECT
 echo "drop role 'r_0002'" | $BENDSQL_CLIENT_CONNECT
-echo "unset enable_experimental_rbac_check" | $BENDSQL_CLIENT_CONNECT
 
 echo "=== test ownership: show stmt ==="
 echo "drop user if exists a;" | $BENDSQL_CLIENT_CONNECT
diff --git a/tests/suites/0_stateless/18_rbac/18_0007_privilege_access.sh b/tests/suites/0_stateless/18_rbac/18_0007_privilege_access.sh
@@ -7,8 +7,6 @@ export TEST_USER_PASSWORD="password"
 export TEST_USER_CONNECT="bendsql --user=test-user --password=password --host=${QUERY_MYSQL_HANDLER_HOST} --port ${QUERY_HTTP_HANDLER_PORT}"
 export RM_UUID="sed -E ""s/[a-z0-9]{8}-[a-z0-9]{4}-[a-z0-9]{4}-[a-z0-9]{4}-[a-z0-9]{12}/UUID/g"""
 
-echo "set global enable_experimental_rbac_check=1" | $BENDSQL_CLIENT_CONNECT
-
 stmt "drop database if exists db01;"
 stmt "create database db01;"
 stmt "drop database if exists dbnotexists;"
@@ -275,5 +273,3 @@ echo "drop table if exists t1" | $BENDSQL_CLIENT_CONNECT
 echo "drop table if exists t2" | $BENDSQL_CLIENT_CONNECT
 echo "drop stage if exists s3;" | $BENDSQL_CLIENT_CONNECT
 echo "drop database if exists db01" | $BENDSQL_CLIENT_CONNECT
-
-echo "unset enable_experimental_rbac_check" | $BENDSQL_CLIENT_CONNECT
diff --git a/tests/suites/1_stateful/00_stage/00_0012_stage_priv.sh b/tests/suites/1_stateful/00_stage/00_0012_stage_priv.sh
@@ -9,8 +9,6 @@ export TEST_USER_CONNECT="bendsql --user=u1 --password=password --host=${QUERY_M
 export USER_B_CONNECT="bendsql --user=b --password=password --host=${QUERY_MYSQL_HANDLER_HOST} --port ${QUERY_HTTP_HANDLER_PORT}"
 export RM_UUID="sed -E ""s/[a-z0-9]{8}-[a-z0-9]{4}-[a-z0-9]{4}-[a-z0-9]{4}-[a-z0-9]{12}/UUID/g"""
 
-echo "set global enable_experimental_rbac_check=1" | $BENDSQL_CLIENT_CONNECT
-
 echo "drop table if exists test_table;" | $BENDSQL_CLIENT_CONNECT
 echo "drop user if exists u1;" | $BENDSQL_CLIENT_CONNECT
 echo "drop STAGE if exists s2;" | $BENDSQL_CLIENT_CONNECT
@@ -144,5 +142,3 @@ echo "insert into t select \$1 from 'fs:///tmp/00_0020/' (FILE_FORMAT => 'CSV');
 echo "drop table if exists t" | $BENDSQL_CLIENT_CONNECT
 echo "drop user if exists b" | $BENDSQL_CLIENT_CONNECT
 rm -rf /tmp/00_0012
-
-echo "unset enable_experimental_rbac_check" | $BENDSQL_CLIENT_CONNECT

Original file line number	Diff line number	Diff line change
`@@ -78,9 +78,9 @@ pub async fn validate_grant_object_exists(`
`78`	`78`	`}`
`79`	`79`	`GrantObject::UDF(udf) => {`
`80`	`80`	`if !UserApiProvider::instance().exists_udf(&tenant, udf).await? {`
`81`		`- return Err(databend_common_exception::ErrorCode::UnknownStage(format!(`
`82`		`- "udf {udf} not exists"`
`83`		`- )));`
	`81`	`+ return Err(databend_common_exception::ErrorCode::UnknownFunction(`
	`82`	`+ format!("udf {udf} not exists"),`
	`83`	`+ ));`
`84`	`84`	`}`
`85`	`85`	`}`
`86`	`86`	`GrantObject::Stage(stage) => {`