Merge pull request #7 from rustbase/feature/rustbase-0.4.0

feat: support rustbase v0.4.0

Author: peeeuzin

Cargo.lock

@@ -8,7 +8,7 @@ license = "MIT"
 [dependencies]
 tokio = { version = "1.23.1", features = ["macros", "rt-multi-thread", "net", "io-std", "io-util"] }
 bson = "2.4.0"
-rustyline = "10.0.0"
+rustyline = "10.1.1"
 colored = "2.0.0"
 clap = "3.2.19"
 clap_derive = "3.2.18"
@@ -17,4 +17,7 @@ serde = { version = "1.0.147", features = ["derive"] }
 tokio-rustls = "0.23.4"
 rustls-pemfile = "1.0.1"
 rustls = "0.20.7"
-webpki = "0.22.0"
+webpki = "0.22.0"
+scram = "0.6.0"
+unicode-width = "0.1.10"
+rustyline-derive = "0.7.0"

Cargo.toml

@@ -0,0 +1,26 @@
+use scram::ScramClient;
+use tokio::io::{AsyncRead, AsyncWrite};
+
+use super::AuthConfig;
+use super::Rustbase;
+
+pub async fn auth<IO>(auth_config: AuthConfig, client: &mut IO)
+where
+    IO: AsyncRead + AsyncWrite + Unpin,
+{
+    let scram = ScramClient::new(&auth_config.username, &auth_config.password, None);
+    let (scram, client_first) = scram.client_first();
+
+    let server_first = Rustbase::send_and_receive(client, client_first.as_bytes().to_vec()).await;
+    let server_first = String::from_utf8(server_first).unwrap();
+
+    let scram = scram
+        .handle_server_first(&server_first)
+        .expect("Invalid server first message, maybe server is not using scram? or maybe the username is wrong?");
+    let (scram, client_final) = scram.client_final();
+
+    let server_final = Rustbase::send_and_receive(client, client_final.as_bytes().to_vec()).await;
+    let server_final = String::from_utf8(server_final).unwrap();
+
+    scram.handle_server_final(&server_final).unwrap();
+}

src/engine/auth.rs

@@ -11,6 +11,8 @@ use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
 use tokio::net::TcpStream;
 use tokio_rustls::{client::TlsStream, TlsConnector};
 
+mod auth;
+
 #[derive(Debug)]
 pub enum Request {
     Query(String),
@@ -39,42 +41,73 @@ pub enum Status {
     InvalidAuth,
 }
 
-pub struct Rustbase {
+#[derive(Clone)]
+pub struct RustbaseConfig {
+    pub tls: Option<TlsConfig>,
     pub host: String,
+    pub port: String,
+    pub database: String,
+    pub auth: Option<AuthConfig>,
+}
+
+#[derive(Clone)]
+pub struct TlsConfig {
+    pub ca_file: String,
+}
+
+#[derive(Clone)]
+pub struct AuthConfig {
+    pub username: String,
+    pub password: String,
+}
+
+pub struct Rustbase {
     pub client: Option<TcpStream>,
     pub database: String,
     pub tls_client: Option<TlsStream<TcpStream>>,
 }
 
 impl Rustbase {
-    pub async fn connect(host: String, port: String, database: String) -> Rustbase {
-        let client = TcpStream::connect(format!("{}:{}", host, port))
+    pub async fn connect(config: RustbaseConfig) -> Rustbase {
+        if config.clone().tls.is_some() {
+            let mut rb = Rustbase::connect_tls(config.clone()).await;
+
+            if let Some(auth) = config.auth {
+                let client = rb.tls_client.as_mut().unwrap();
+
+                auth::auth(auth, client).await;
+            }
+
+            return rb;
+        }
+
+        let mut client = TcpStream::connect(format!("{}:{}", config.host, config.port))
             .await
             .unwrap();
 
+        if let Some(auth) = config.auth {
+            auth::auth(auth, &mut client).await;
+        }
+
         Rustbase {
             client: Some(client),
-            database,
+            database: config.database,
             tls_client: None,
-            host,
         }
     }
 
-    pub async fn connect_tls(
-        host: String,
-        port: String,
-        database: String,
-        ca_file: String,
-    ) -> Rustbase {
-        if !Path::new(&ca_file).exists() {
+    async fn connect_tls(config: RustbaseConfig) -> Rustbase {
+        let tls = config.tls.unwrap();
+
+        if !Path::new(&tls.ca_file).exists() {
             println!(
                 "{} CA file not found: use --ca_file=<path>",
                 "[Error]".red()
             );
         }
 
         let mut root_cert_store = rustls::RootCertStore::empty();
-        let mut pem = BufReader::new(File::open(ca_file).unwrap());
+        let mut pem = BufReader::new(File::open(tls.ca_file).unwrap());
         let certs = rustls_pemfile::certs(&mut pem).unwrap();
         let trust_anchors = certs.iter().map(|cert| {
             let ta = webpki::TrustAnchor::try_from_cert_der(&cert[..]).unwrap();
@@ -86,28 +119,27 @@ impl Rustbase {
         });
         root_cert_store.add_server_trust_anchors(trust_anchors);
 
-        let config = rustls::ClientConfig::builder()
+        let client_config = rustls::ClientConfig::builder()
             .with_safe_defaults()
             .with_root_certificates(root_cert_store)
             .with_no_client_auth();
 
-        let connector = TlsConnector::from(Arc::new(config));
+        let connector = TlsConnector::from(Arc::new(client_config));
 
-        let domain = rustls::ServerName::try_from(host.as_str())
+        let domain = rustls::ServerName::try_from(config.host.as_str())
             .map_err(|_| std::io::Error::new(std::io::ErrorKind::InvalidInput, "invalid dnsname"))
             .unwrap();
 
-        let client = TcpStream::connect(format!("{}:{}", host, port))
+        let client = TcpStream::connect(format!("{}:{}", config.host, config.port))
             .await
             .unwrap();
 
         let tls_client = connector.connect(domain, client).await.unwrap();
 
         Rustbase {
             client: None,
-            database,
+            database: config.database,
             tls_client: Some(tls_client),
-            host,
         }
     }
 

src/engine/mod.rs

@@ -4,8 +4,36 @@ mod engine;
 mod utils;
 
 use colored::Colorize;
+
+use rustyline::config::Configurer;
 use rustyline::error::ReadlineError;
-use rustyline::{Editor, Result};
+use rustyline::highlight::Highlighter;
+use rustyline::{ColorMode, Editor, Result};
+use rustyline_derive::{Completer, Helper, Hinter, Validator};
+
+use std::borrow::Cow::{self, Borrowed, Owned};
+
+use engine::{AuthConfig, RustbaseConfig, TlsConfig};
+
+#[derive(Completer, Helper, Hinter, Validator)]
+struct MaskingHighlighter {
+    masking: bool,
+}
+
+impl Highlighter for MaskingHighlighter {
+    fn highlight<'l>(&self, line: &'l str, _pos: usize) -> Cow<'l, str> {
+        use unicode_width::UnicodeWidthStr;
+        if self.masking {
+            Owned("*".repeat(line.width()))
+        } else {
+            Borrowed(line)
+        }
+    }
+
+    fn highlight_char(&self, _line: &str, _pos: usize) -> bool {
+        self.masking
+    }
+}
 
 /// A CLI for Rustbase Database Server
 #[derive(clap_derive::Parser)]
@@ -24,14 +52,33 @@ struct Args {
     #[clap(long, default_value = "")]
     ca_file: String,
 
+    /// Use authentication for connection
+    #[clap(long, short = 'a')]
+    use_auth: bool,
+
     #[clap(subcommand)]
     commands: Option<Commands>,
 }
 
+fn prompt<T>(prompt: &str, editor: &mut Editor<T>) -> Result<String>
+where
+    T: rustyline::Helper,
+{
+    let result = loop {
+        let prompt_result = editor.readline(prompt)?;
+
+        if !prompt_result.is_empty() {
+            break prompt_result;
+        }
+    };
+
+    Ok(result)
+}
+
 #[derive(clap_derive::Subcommand, PartialEq)]
 enum Commands {
-    #[clap(about = "Upgrade Rustbase and Rustbase CLI")]
-    Upgrade,
+    #[clap(about = "Update Rustbase CLI")]
+    Update,
     #[clap(about = "Clean repl history")]
     Clean,
 }
@@ -43,7 +90,7 @@ async fn main() -> Result<()> {
     let repl_path = utils::get_current_path().join("repl.history");
 
     match args.commands {
-        Some(Commands::Upgrade) => {
+        Some(Commands::Update) => {
             println!("Not implemented yet");
             return Ok(());
         }
@@ -83,28 +130,51 @@ async fn main() -> Result<()> {
     );
     println!("Press Ctrl+C to exit.");
 
-    let mut rl = Editor::<()>::new()?;
+    let h = MaskingHighlighter { masking: false };
+    let mut rl = Editor::new()?;
+    rl.set_helper(Some(h));
 
     rl.load_history(repl_path.to_str().unwrap()).ok();
     println!();
 
-    let mut database = rl.readline("Database: ")?;
+    let mut database = prompt("Database: ", &mut rl)?;
 
-    loop {
-        if database.is_empty() {
-            database = rl.readline("Database: ")?;
-            continue;
-        } else {
-            break;
-        }
-    }
+    let auth_config = if args.use_auth {
+        let username = prompt("Username: ", &mut rl)?;
+
+        rl.helper_mut().expect("No helper").masking = true;
+        rl.set_color_mode(ColorMode::Forced); // force masking
+        rl.set_auto_add_history(false); // make sure password is not added to history
+
+        let password = prompt("Password: ", &mut rl)?;
 
-    let mut client = if args.tls {
-        engine::Rustbase::connect_tls(args.host, args.port, database.clone(), args.ca_file).await
+        rl.helper_mut().expect("No helper").masking = false;
+        rl.set_color_mode(ColorMode::Disabled);
+        rl.set_auto_add_history(true);
+
+        Some(AuthConfig { username, password })
+    } else {
+        None
+    };
+
+    let tls_config = if args.tls {
+        let ca_file = args.ca_file;
+
+        Some(TlsConfig { ca_file })
     } else {
-        engine::Rustbase::connect(args.host, args.port, database.clone()).await
+        None
+    };
+
+    let config = RustbaseConfig {
+        database: database.clone(),
+        port: args.port,
+        host: args.host,
+        auth: auth_config,
+        tls: tls_config,
     };
 
+    let mut client = engine::Rustbase::connect(config).await;
+
     loop {
         let readline = rl.readline(format!("{}> ", database).as_str());
         match readline {