feat: make password & dsn sensitive (#404)

* feat: make password & dsn sensitive

* chore: tmp disable geometry test

cli/Cargo.toml

@@ -11,9 +11,11 @@ authors = { workspace = true }
 repository = { workspace = true }
 
 [dependencies]
+databend-client = { workspace = true }
 databend-driver = { workspace = true, features = ["rustls", "flight-sql"] }
 
 anyhow = "1.0"
+async-recursion = "1.1.0"
 async-trait = "0.1"
 chrono = { version = "0.4.35", default-features = false, features = ["clock"] }
 clap = { version = "4.4", features = ["derive", "env"] }
@@ -42,7 +44,6 @@ toml = "0.8"
 tracing-appender = "0.2"
 unicode-segmentation = "1.10"
 url = { version = "2.5", default-features = false }
-async-recursion = "1.1.0"
 
 [build-dependencies]
 vergen = { version = "8.2", features = ["build", "git", "gix"] }

cli/src/args.rs

@@ -15,13 +15,14 @@
 use std::collections::BTreeMap;
 
 use anyhow::{anyhow, Result};
+use databend_client::auth::SensitiveString;
 
 #[derive(Debug, Clone, PartialEq, Default)]
 pub struct ConnectionArgs {
     pub host: String,
     pub port: Option<u16>,
     pub user: String,
-    pub password: Option<String>,
+    pub password: SensitiveString,
     pub database: Option<String>,
     pub flight: bool,
     pub args: BTreeMap<String, String>,
@@ -33,9 +34,7 @@ impl ConnectionArgs {
         dsn.set_host(Some(&self.host))?;
         _ = dsn.set_port(self.port);
         _ = dsn.set_username(&self.user);
-        if let Some(password) = self.password {
-            _ = dsn.set_password(Some(&password))
-        };
+        _ = dsn.set_password(Some(self.password.inner()));
         if let Some(database) = self.database {
             dsn.set_path(&database);
         }
@@ -64,7 +63,7 @@ impl ConnectionArgs {
         let host = u.host_str().ok_or(anyhow!("missing host"))?.to_string();
         let port = u.port();
         let user = u.username().to_string();
-        let password = u.password().map(|s| s.to_string());
+        let password = SensitiveString::from(u.password().unwrap_or_default());
         let database = u.path().strip_prefix('/').map(|s| s.to_string());
         Ok(Self {
             host,
@@ -90,7 +89,7 @@ mod test {
             host: "app.databend.com".to_string(),
             port: None,
             user: "username".to_string(),
-            password: Some("3a@SC(nYE1k={{R".to_string()),
+            password: SensitiveString::from("3a@SC(nYE1k={{R"),
             database: Some("test".to_string()),
             flight: false,
             args: {
@@ -113,7 +112,7 @@ mod test {
             host: "app.databend.com".to_string(),
             port: Some(443),
             user: "username".to_string(),
-            password: Some("3a@SC(nYE1k={{R".to_string()),
+            password: SensitiveString::from("3a@SC(nYE1k={{R"),
             database: Some("test".to_string()),
             flight: false,
             args: {

cli/src/display.rs

@@ -14,23 +14,20 @@
 
 use std::collections::HashSet;
 use std::fmt::Write;
-use unicode_segmentation::UnicodeSegmentation;
 
 use anyhow::{anyhow, Result};
 use comfy_table::{Cell, CellAlignment, Table};
-use terminal_size::{terminal_size, Width};
-
 use databend_driver::{Row, RowStatsIterator, RowWithStats, SchemaRef, ServerStats};
+use indicatif::{HumanBytes, ProgressBar, ProgressState, ProgressStyle};
 use rustyline::highlight::Highlighter;
+use terminal_size::{terminal_size, Width};
 use tokio::time::Instant;
 use tokio_stream::StreamExt;
+use unicode_segmentation::UnicodeSegmentation;
 
-use indicatif::{HumanBytes, ProgressBar, ProgressState, ProgressStyle};
-
-use crate::config::OutputQuoteStyle;
 use crate::{
     ast::format_query,
-    config::{ExpandMode, OutputFormat, Settings},
+    config::{ExpandMode, OutputFormat, OutputQuoteStyle, Settings},
     helper::CliHelper,
     session::QueryKind,
 };

cli/src/main.rs

@@ -27,14 +27,17 @@ use std::{
     io::{stdin, IsTerminal},
 };
 
-use crate::args::ConnectionArgs;
-use crate::config::OutputQuoteStyle;
 use anyhow::{anyhow, Result};
 use clap::{ArgAction, CommandFactory, Parser, ValueEnum};
-use config::{Config, OutputFormat, Settings, TimeOption};
+use databend_client::auth::SensitiveString;
 use log::info;
 use once_cell::sync::Lazy;
 
+use crate::{
+    args::ConnectionArgs,
+    config::{Config, OutputFormat, OutputQuoteStyle, Settings, TimeOption},
+};
+
 static VERSION: Lazy<String> = Lazy::new(|| {
     let version = option_env!("CARGO_PKG_VERSION").unwrap_or("unknown");
     let sha = option_env!("VERGEN_GIT_SHA").unwrap_or("dev");
@@ -106,32 +109,45 @@ struct Args {
     #[clap(long, help = "Print help information")]
     help: bool,
 
-    #[clap(long, help = "Using flight sql protocol")]
+    #[clap(long, help = "Using flight sql protocol, ignored when --dsn is set")]
     flight: bool,
 
-    #[clap(long, help = "Enable TLS")]
+    #[clap(long, help = "Enable TLS, ignored when --dsn is set")]
     tls: bool,
 
-    #[clap(short = 'h', long, help = "Databend Server host, Default: 127.0.0.1")]
+    #[clap(
+        short = 'h',
+        long,
+        help = "Databend Server host, Default: 127.0.0.1, ignored when --dsn is set"
+    )]
     host: Option<String>,
 
-    #[clap(short = 'P', long, help = "Databend Server port, Default: 8000")]
+    #[clap(
+        short = 'P',
+        long,
+        help = "Databend Server port, Default: 8000, ignored when --dsn is set"
+    )]
     port: Option<u16>,
 
-    #[clap(short = 'u', long, help = "Default: root")]
+    #[clap(short = 'u', long, help = "Default: root, overrides username in DSN")]
     user: Option<String>,
 
-    #[clap(short = 'p', long, env = "BENDSQL_PASSWORD")]
-    password: Option<String>,
+    #[clap(
+        short = 'p',
+        long,
+        env = "BENDSQL_PASSWORD",
+        help = "Password, overrides password in DSN"
+    )]
+    password: Option<SensitiveString>,
 
-    #[clap(short = 'D', long, help = "Database name")]
+    #[clap(short = 'D', long, help = "Database name, overrides database in DSN")]
     database: Option<String>,
 
-    #[clap(long, value_parser = parse_key_val::<String, String>, help = "Settings")]
+    #[clap(long, value_parser = parse_key_val::<String, String>, help = "Settings, ignored when --dsn is set")]
     set: Vec<(String, String)>,
 
     #[clap(long, env = "BENDSQL_DSN", help = "Data source name")]
-    dsn: Option<String>,
+    dsn: Option<SensitiveString>,
 
     #[clap(short = 'n', long, help = "Force non-interactive mode")]
     non_interactive: bool,
@@ -219,15 +235,6 @@ pub async fn main() -> Result<()> {
             if args.port.is_some() {
                 eprintln!("warning: --port is ignored when --dsn is set");
             }
-            if args.user.is_some() {
-                eprintln!("warning: --user is ignored when --dsn is set");
-            }
-            if args.password.is_some() {
-                eprintln!("warning: --password is ignored when --dsn is set");
-            }
-            if args.role.is_some() {
-                eprintln!("warning: --role is ignored when --dsn is set");
-            }
             if !args.set.is_empty() {
                 eprintln!("warning: --set is ignored when --dsn is set");
             }
@@ -237,7 +244,7 @@ pub async fn main() -> Result<()> {
             if args.flight {
                 eprintln!("warning: --flight is ignored when --dsn is set");
             }
-            ConnectionArgs::from_dsn(&dsn)?
+            ConnectionArgs::from_dsn(dsn.inner())?
         }
         None => {
             if let Some(host) = args.host {
@@ -246,9 +253,6 @@ pub async fn main() -> Result<()> {
             if let Some(port) = args.port {
                 config.connection.port = Some(port);
             }
-            if let Some(user) = args.user {
-                config.connection.user = user;
-            }
             for (k, v) in args.set {
                 config.connection.args.insert(k, v);
             }
@@ -258,14 +262,11 @@ pub async fn main() -> Result<()> {
                     .args
                     .insert("sslmode".to_string(), "disable".to_string());
             }
-            if let Some(role) = args.role {
-                config.connection.args.insert("role".to_string(), role);
-            }
             ConnectionArgs {
                 host: config.connection.host.clone(),
                 port: config.connection.port,
                 user: config.connection.user.clone(),
-                password: args.password,
+                password: SensitiveString::from(""),
                 database: config.connection.database.clone(),
                 flight: args.flight,
                 args: config.connection.args.clone(),
@@ -276,6 +277,18 @@ pub async fn main() -> Result<()> {
     if args.database.is_some() {
         conn_args.database = args.database;
     }
+    // override user if specified in command line
+    if let Some(user) = args.user {
+        config.connection.user = user;
+    }
+    // override password if specified in command line
+    if let Some(password) = args.password {
+        conn_args.password = password;
+    }
+    // override role if specified in command line
+    if let Some(role) = args.role {
+        config.connection.args.insert("role".to_string(), role);
+    }
     let dsn = conn_args.get_dsn()?;
 
     let mut settings = Settings::default();

cli/tests/00-base.result

@@ -20,5 +20,5 @@ Asia/Shanghai
 NULL	{'k1':'v1','k2':'v2'}	(2,NULL)
 1	NULL	1	ab
 NULL	v1	2	NULL
-{'k1':'v1','k2':'v2'}	[6162,78797A]	('[1,2]','SRID=4326;POINT(1 2)','2024-04-10')
+{'k1':'v1','k2':'v2'}	[6162,78797A]	('[1,2]','2024-04-10')
 bye

cli/tests/00-base.sql

@@ -46,7 +46,7 @@ insert into test_nested values([1,2,3], null, (1, 'ab')), (null, {'k1':'v1', 'k2
 select * from test_nested;
 select a[1], b['k1'], c:x, c:y from test_nested;
 
-select {'k1':'v1','k2':'v2'}, [to_binary('ab'), to_binary('xyz')], (parse_json('[1,2]'), st_geometryfromwkt('SRID=4326;POINT(1.0 2.0)'), to_date('2024-04-10'));
+select {'k1':'v1','k2':'v2'}, [to_binary('ab'), to_binary('xyz')], (parse_json('[1,2]'), to_date('2024-04-10'));
 
 select 'bye';
 drop table test;

core/src/auth.rs

@@ -25,19 +25,22 @@ pub trait Auth: Sync + Send {
 #[derive(Clone)]
 pub struct BasicAuth {
     username: String,
-    password: Option<String>,
+    password: SensitiveString,
 }
 
 impl BasicAuth {
-    pub fn new(username: String, password: Option<String>) -> Self {
-        Self { username, password }
+    pub fn new(username: impl ToString, password: impl ToString) -> Self {
+        Self {
+            username: username.to_string(),
+            password: SensitiveString(password.to_string()),
+        }
     }
 }
 
 #[async_trait::async_trait]
 impl Auth for BasicAuth {
     async fn wrap(&self, builder: RequestBuilder) -> Result<RequestBuilder> {
-        Ok(builder.basic_auth(&self.username, self.password.as_deref()))
+        Ok(builder.basic_auth(&self.username, Some(self.password.inner())))
     }
 
     fn username(&self) -> String {
@@ -47,19 +50,21 @@ impl Auth for BasicAuth {
 
 #[derive(Clone)]
 pub struct AccessTokenAuth {
-    token: String,
+    token: SensitiveString,
 }
 
 impl AccessTokenAuth {
-    pub fn new(token: String) -> Self {
-        Self { token }
+    pub fn new(token: impl ToString) -> Self {
+        Self {
+            token: SensitiveString::from(token.to_string()),
+        }
     }
 }
 
 #[async_trait::async_trait]
 impl Auth for AccessTokenAuth {
     async fn wrap(&self, builder: RequestBuilder) -> Result<RequestBuilder> {
-        Ok(builder.bearer_auth(&self.token))
+        Ok(builder.bearer_auth(self.token.inner()))
     }
 
     fn username(&self) -> String {
@@ -73,7 +78,8 @@ pub struct AccessTokenFileAuth {
 }
 
 impl AccessTokenFileAuth {
-    pub fn new(token_file: String) -> Self {
+    pub fn new(token_file: impl ToString) -> Self {
+        let token_file = token_file.to_string();
         Self { token_file }
     }
 }
@@ -96,3 +102,68 @@ impl Auth for AccessTokenFileAuth {
         "token".to_string()
     }
 }
+
+#[derive(::serde::Deserialize, ::serde::Serialize)]
+#[serde(from = "String", into = "String")]
+#[derive(Clone, Default, PartialEq, Eq)]
+pub struct SensitiveString(String);
+
+impl From<String> for SensitiveString {
+    fn from(value: String) -> Self {
+        Self(value)
+    }
+}
+
+impl From<&str> for SensitiveString {
+    fn from(value: &str) -> Self {
+        Self(value.to_string())
+    }
+}
+
+impl From<SensitiveString> for String {
+    fn from(value: SensitiveString) -> Self {
+        value.0
+    }
+}
+
+impl std::fmt::Display for SensitiveString {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(f, "**REDACTED**")
+    }
+}
+
+impl std::fmt::Debug for SensitiveString {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        // we keep the double quotes here to keep the String behavior
+        write!(f, "\"**REDACTED**\"")
+    }
+}
+
+impl SensitiveString {
+    #[must_use]
+    pub fn inner(&self) -> &str {
+        self.0.as_str()
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn serialization() {
+        let json_value = "\"foo\"";
+        let value: SensitiveString = serde_json::from_str(json_value).unwrap();
+        let result: String = serde_json::to_string(&value).unwrap();
+        assert_eq!(result, json_value);
+    }
+
+    #[test]
+    fn hide_content() {
+        let value = SensitiveString("hello world".to_string());
+        let display = format!("{value}");
+        assert_eq!(display, "**REDACTED**");
+        let debug = format!("{value:?}");
+        assert_eq!(debug, "\"**REDACTED**\"");
+    }
+}