feat(query): add config network_policy_whitelist (databendlabs#17065)

* feat(query): add config network_policy_whitelist * z * z * z * z * z * z * z
TCeason · Dec 18, 2024 · b38c46f · b38c46f
1 parent d4fbfa2
commit b38c46f
Show file tree

Hide file tree

Showing 10 changed files with 39 additions and 5 deletions.
diff --git a/.github/actions/setup_bendsql/action.yml b/.github/actions/setup_bendsql/action.yml
@@ -8,15 +8,16 @@ runs:
       shell: bash
       run: |
         if command -v bendsql >/dev/null 2>&1; then
+          bendsql --version
           exit 0
         fi
         case $RUNNER_PROVIDER in
           aws)
-            aws s3 cp s3://databend-ci/packages/bendsql_$(dpkg --print-architecture).deb /tmp/bendsql.deb
+            aws s3 cp s3://databend-ci/packages/bendsql_$(dpkg --print-architecture).deb /tmp/bendsql.deb --no-progress
             sudo dpkg -i /tmp/bendsql.deb
             ;;
           gcp)
-            gsutil cp gs://databend-ci/packages/bendsql_$(dpkg --print-architecture).deb /tmp/bendsql.deb
+            gsutil cp gs://databend-ci/packages/bendsql_$(dpkg --print-architecture).deb /tmp/bendsql.deb --no-progress
             sudo dpkg -i /tmp/bendsql.deb
             ;;
           *)
@@ -30,6 +31,7 @@ runs:
       shell: bash
       run: |
         if command -v bendsql >/dev/null 2>&1; then
+          bendsql --version
           exit 0
         fi
         brew install databendcloud/homebrew-tap/bendsql

diff --git a/Cargo.lock b/Cargo.lock
diff --git a/scripts/ci/deploy/config/databend-query-node-1.toml b/scripts/ci/deploy/config/databend-query-node-1.toml
@@ -44,6 +44,8 @@ udf_server_allow_insecure = true
 
 cloud_control_grpc_server_address = "http://0.0.0.0:50051"
 
+# network_policy_whitelist = ['127.0.0.0/8']
+
 [[query.users]]
 name = "root"
 auth_type = "no_password"

diff --git a/scripts/ci/deploy/config/databend-query-node-native.toml b/scripts/ci/deploy/config/databend-query-node-native.toml
@@ -36,6 +36,8 @@ enable_udf_server = true
 udf_server_allow_list = ['http://0.0.0.0:8815']
 udf_server_allow_insecure = true
 
+# network_policy_whitelist = ['127.0.0.0/8']
+
 [[query.users]]
 name = "root"
 auth_type = "no_password"

diff --git a/src/common/exception/Cargo.toml b/src/common/exception/Cargo.toml
@@ -18,6 +18,7 @@ arrow-flight = { workspace = true }
 arrow-schema = { workspace = true }
 backtrace = { workspace = true, features = ["std", "serialize-serde"] }
 bincode = { workspace = true }
+cidr = { workspace = true }
 geozero = { workspace = true }
 gimli = { workspace = true }
 http = { workspace = true }

diff --git a/src/common/exception/src/exception_into.rs b/src/common/exception/src/exception_into.rs
@@ -52,6 +52,12 @@ impl From<std::net::AddrParseError> for ErrorCode {
     }
 }
 
+impl From<cidr::errors::NetworkParseError> for ErrorCode {
+    fn from(error: cidr::errors::NetworkParseError) -> Self {
+        ErrorCode::BadAddressFormat(format!("Bad network format, cause: {}", error))
+    }
+}
+
 impl From<std::str::Utf8Error> for ErrorCode {
     fn from(error: std::str::Utf8Error) -> Self {
         ErrorCode::Internal(format!("Invalid Utf8, cause: {}", error))

diff --git a/src/query/config/src/config.rs b/src/query/config/src/config.rs
@@ -1687,6 +1687,10 @@ pub struct QueryConfig {
     #[clap(long, value_name = "VALUE", default_value = "50")]
     pub max_cached_queries_profiles: usize,
 
+    /// A list of network that not to be checked by network policy.
+    #[clap(long, value_name = "VALUE")]
+    pub network_policy_whitelist: Vec<String>,
+
     #[clap(skip)]
     pub settings: HashMap<String, SettingValue>,
 }
@@ -1776,6 +1780,7 @@ impl TryInto<InnerQueryConfig> for QueryConfig {
             cloud_control_grpc_server_address: self.cloud_control_grpc_server_address,
             cloud_control_grpc_timeout: self.cloud_control_grpc_timeout,
             max_cached_queries_profiles: self.max_cached_queries_profiles,
+            network_policy_whitelist: self.network_policy_whitelist,
             settings: self
                 .settings
                 .into_iter()
@@ -1878,6 +1883,7 @@ impl From<InnerQueryConfig> for QueryConfig {
             cloud_control_grpc_server_address: inner.cloud_control_grpc_server_address,
             cloud_control_grpc_timeout: inner.cloud_control_grpc_timeout,
             max_cached_queries_profiles: inner.max_cached_queries_profiles,
+            network_policy_whitelist: inner.network_policy_whitelist,
             settings: HashMap::new(),
         }
     }

diff --git a/src/query/config/src/inner.rs b/src/query/config/src/inner.rs
@@ -246,6 +246,9 @@ pub struct QueryConfig {
     pub cloud_control_grpc_server_address: Option<String>,
     pub cloud_control_grpc_timeout: u64,
     pub max_cached_queries_profiles: usize,
+
+    pub network_policy_whitelist: Vec<String>,
+
     pub settings: HashMap<String, UserSettingValue>,
 }
 
@@ -322,6 +325,7 @@ impl Default for QueryConfig {
             cloud_control_grpc_timeout: 0,
             data_retention_time_in_days_max: 90,
             max_cached_queries_profiles: 50,
+            network_policy_whitelist: Vec::new(),
             settings: HashMap::new(),
         }
     }

diff --git a/src/query/service/tests/it/storages/testdata/configs_table_basic.txt b/src/query/service/tests/it/storages/testdata/configs_table_basic.txt
@@ -116,6 +116,7 @@ DB.Table: 'system'.'configs', Table: configs-table_id:1, ver:0, Engine: SystemCo
 | 'query'   | 'mysql_handler_tcp_keepalive_timeout_secs'      | '120'                                                                                                                                                                                             | ''       |
 | 'query'   | 'mysql_tls_server_cert'                         | ''                                                                                                                                                                                                | ''       |
 | 'query'   | 'mysql_tls_server_key'                          | ''                                                                                                                                                                                                | ''       |
+| 'query'   | 'network_policy_whitelist'                      | ''                                                                                                                                                                                                | ''       |
 | 'query'   | 'num_cpus'                                      | '0'                                                                                                                                                                                               | ''       |
 | 'query'   | 'openai_api_chat_base_url'                      | 'https://api.openai.com/v1/'                                                                                                                                                                      | ''       |
 | 'query'   | 'openai_api_completion_model'                   | 'gpt-3.5-turbo'                                                                                                                                                                                   | ''       |

diff --git a/src/query/users/src/user_mgr.rs b/src/query/users/src/user_mgr.rs
@@ -17,6 +17,7 @@ use core::net::Ipv4Addr;
 use chrono::DateTime;
 use chrono::Utc;
 use cidr::Ipv4Cidr;
+use databend_common_config::GlobalConfig;
 use databend_common_exception::ErrorCode;
 use databend_common_exception::Result;
 use databend_common_management::UserApi;
@@ -91,15 +92,23 @@ impl UserApiProvider {
         client_ip: Option<&str>,
     ) -> Result<()> {
         let ip_addr: Ipv4Addr = match client_ip {
-            Some(client_ip) => client_ip.parse().unwrap(),
+            Some(client_ip) => client_ip.parse()?,
             None => {
                 return Err(ErrorCode::AuthenticateFailure("Unknown client ip"));
             }
         };
 
+        let whitelist = &GlobalConfig::instance().query.network_policy_whitelist;
+        for whitelist_ip in whitelist {
+            let cidr: Ipv4Cidr = whitelist_ip.parse()?;
+            if cidr.contains(&ip_addr) {
+                return Ok(());
+            }
+        }
+
         let network_policy = self.get_network_policy(tenant, policy).await?;
         for blocked_ip in network_policy.blocked_ip_list {
-            let blocked_cidr: Ipv4Cidr = blocked_ip.parse().unwrap();
+            let blocked_cidr: Ipv4Cidr = blocked_ip.parse()?;
             if blocked_cidr.contains(&ip_addr) {
                 return Err(ErrorCode::AuthenticateFailure(format!(
                     "client ip `{}` is blocked",
@@ -109,7 +118,7 @@ impl UserApiProvider {
         }
         let mut allow = false;
         for allowed_ip in network_policy.allowed_ip_list {
-            let allowed_cidr: Ipv4Cidr = allowed_ip.parse().unwrap();
+            let allowed_cidr: Ipv4Cidr = allowed_ip.parse()?;
             if allowed_cidr.contains(&ip_addr) {
                 allow = true;
                 break;