Working auth with storing data in pgsql

common/Cargo.toml

@@ -22,6 +22,7 @@ diesel = {version = "2.1", optional = true, default-features = false, features =
 diesel-async = {version = "0.3", optional = true, features = ["postgres", "bb8"]}
 anyhow = "1.0"
 prometheus-client = "0.21.2"
+bcrypt = "0.15.0"
 
 [features]
 default = ["web", "db"]

common/migrations/2023-02-01-210714_init/up.sql

@@ -37,3 +37,9 @@ CREATE TABLE IF NOT EXISTS funcs (
 CREATE UNIQUE INDEX IF NOT EXISTS funcs_db ON funcs(chksum, db_id);
 CREATE INDEX IF NOT EXISTS funcs_ranking ON funcs(chksum,rank);
 CREATE INDEX IF NOT EXISTS func_chksum ON funcs(chksum);
+
+CREATE TABLE IF NOT EXISTS auth_users (
+    id SERIAL PRIMARY KEY,
+    username VARCHAR(32) NOT NULL UNIQUE,
+    password_hash VARCHAR(128) NOT NULL
+);

common/src/db/mod.rs

@@ -3,11 +3,12 @@ use postgres_native_tls::MakeTlsConnector;
 use serde::Serialize;
 use tokio_postgres::{tls::MakeTlsConnect, Socket, NoTls};
 use std::{collections::HashMap};
+use bcrypt::{hash, verify, DEFAULT_COST};
 use crate::async_drop::{AsyncDropper, AsyncDropGuard};
 mod schema_auto;
 pub mod schema;
 
-use diesel::{upsert::excluded, ExpressionMethods, QueryDsl, NullableExpressionMethods, sql_types::{Array, Binary, VarChar, Integer}, query_builder::{QueryFragment, Query}};
+use diesel::{upsert::excluded, ExpressionMethods, QueryDsl, NullableExpressionMethods, sql_types::{Array, Binary, VarChar, Integer}, query_builder::{QueryFragment, Query}, result::Error::NotFound};
 use diesel_async::RunQueryDsl;
 
 pub type DynConfig = dyn crate::config::HasConfig + Send + Sync;
@@ -337,6 +338,76 @@ impl Database {
         Ok(results)
     }
 
+    pub async fn register_user(&self, username: &str, password: &str) -> Result<(), anyhow::Error> {
+        let conn = &mut self.diesel.get().await?;
+
+        // Hash the password
+        let hashed_password = hash(password, DEFAULT_COST)?;
+
+        // Insert new user
+        diesel::insert_into(schema::auth_users::table)
+            .values((
+                schema::auth_users::username.eq(username),
+                schema::auth_users::password_hash.eq(hashed_password),
+            ))
+            .execute(conn)
+            .await?;
+
+        Ok(())
+    }
+
+    pub async fn change_user_password(&self, username: &str, new_password: &str) -> Result<(), anyhow::Error> {
+        let conn = &mut self.diesel.get().await?;
+
+        // Hash the new password
+        let hashed_password = hash(new_password, DEFAULT_COST)?;
+
+        // Update the user's password
+        diesel::update(schema::auth_users::table.filter(schema::auth_users::username.eq(username)))
+            .set(schema::auth_users::password_hash.eq(hashed_password))
+            .execute(conn)
+            .await?;
+
+        Ok(())
+    }
+
+    pub async fn remove_user(&self, username: &str) -> Result<(), anyhow::Error> {
+        let conn = &mut self.diesel.get().await?;
+
+        // Execute the delete query
+        diesel::delete(schema::auth_users::table.filter(schema::auth_users::username.eq(username)))
+            .execute(conn)
+            .await?;
+
+        Ok(())
+    }
+
+    pub async fn auth_user(&self, login: &str, password: &str) -> Result<bool, anyhow::Error> {
+        let conn = &mut self.diesel.get().await?;
+
+        match schema::auth_users::table
+            .select((schema::auth_users::username, schema::auth_users::password_hash))
+            .filter(schema::auth_users::username.eq(login))
+            .first::<(String, String)>(conn)
+            .await {
+                Ok((_, password_hash)) => {
+                    // If user is found, verify the password
+                    match verify(password, &password_hash) {
+                        Ok(valid) => Ok(valid),
+                        Err(e) => Err(anyhow::Error::new(e)),
+                    }
+                },
+                Err(diesel::result::Error::NotFound) => {
+                    // If user is not found, return false
+                    Ok(false)
+                },
+                Err(e) => {
+                    // For all other errors, return the error
+                    Err(anyhow::Error::new(e))
+                }
+            }
+    }
+
     pub async fn get_files_with_func(&self, func: &[u8]) -> Result<Vec<Vec<u8>>, anyhow::Error> {
         let conn = &mut self.diesel.get().await?;
 

common/src/db/schema_auto.rs

@@ -1,5 +1,13 @@
 // @generated automatically by Diesel CLI.
 
+diesel::table! {
+    auth_users (id) {
+        id -> Int4,
+        username -> Varchar,
+        password_hash -> Varchar,
+    }
+}
+
 diesel::table! {
     dbs (id) {
         id -> Int4,
@@ -46,6 +54,7 @@ diesel::joinable!(dbs -> users (user_id));
 diesel::joinable!(funcs -> dbs (db_id));
 
 diesel::allow_tables_to_appear_in_same_query!(
+    auth_users,
     dbs,
     files,
     funcs,

lumen/src/main.rs

@@ -200,15 +200,25 @@ async fn handle_client<S: AsyncRead + AsyncWrite + Unpin>(state: &SharedState, m
     }).inc();
 
     if let Some(ref creds) = creds {
-        if creds.username != "guest" {
+
+        let auth_state = state.db.auth_user(creds.username, creds.password).await;
+
+        if !auth_state.is_ok() || !auth_state.unwrap() {
             // Only allow "guest" to connect for now.
             rpc::RpcMessage::Fail(rpc::RpcFail {
                 code: 1,
-                message: &format!("{server_name}: invalid username or password. Try logging in with `guest` instead."),
+                message: &format!("{server_name}: invalid username or password."),
             }).async_write(&mut stream).await?;
             return Ok(());
         }
     }
+    else {
+        rpc::RpcMessage::Fail(rpc::RpcFail {
+            code: 1,
+            message: &format!("{server_name}: username and password should be specified."),
+        }).async_write(&mut stream).await?;
+        return Ok(());
+    }
 
     let resp = match hello.protocol_version {
         0..=4 => rpc::RpcMessage::Ok(()),
@@ -329,6 +339,33 @@ fn main() {
                 .default_value("config.toml")
                 .help("Configuration file path")
         )
+        .subcommand(
+            clap::Command::new("add_user")
+            .about("Adds a new user")
+            .arg(Arg::new("username")
+                .help("The username for the new user")
+                .required(true))
+            .arg(Arg::new("password")
+                .help("The password for the new user")
+                .required(true))
+        )
+        .subcommand(
+            clap::Command::new("change_user_pass")
+            .about("Changes a user's password")
+            .arg(Arg::new("username")
+                .help("The username of the user")
+                .required(true))
+            .arg(Arg::new("new_password")
+                .help("The new password for the user")
+                .required(true))
+        )
+        .subcommand(
+            clap::Command::new("remove_user")
+            .about("Removes user")
+            .arg(Arg::new("username")
+                .help("The username of the user")
+                .required(true))
+        )
         .get_matches();
 
     let config = {
@@ -347,7 +384,7 @@ fn main() {
             exit(1);
         },
     };
-    
+
     let db = rt.block_on(async {
         match Database::open(&config.database).await {
             Ok(v) => v,
@@ -367,8 +404,46 @@ fn main() {
         metrics: common::metrics::Metrics::default(),
     });
 
-    let tls_acceptor;
+    let subcommand_future = async {
+        match matches.subcommand() {
+            Some(("add_user", sub_m)) => {
+                let username = sub_m.get_one::<String>("username").unwrap();
+                let password = sub_m.get_one::<String>("password").unwrap();
+
+                match state.db.register_user(username, password).await {
+                    Ok(_) => println!("User added successfully"),
+                    Err(e) => eprintln!("Error adding user: {}", e),
+                }
+                exit(0);
+            },
+            Some(("change_user_pass", sub_m)) => {
+                let username = sub_m.get_one::<String>("username").unwrap();
+                let new_password = sub_m.get_one::<String>("new_password").unwrap();
 
+                match state.db.change_user_password(username, new_password).await {
+                    Ok(_) => println!("User password changed successfully"),
+                    Err(e) => eprintln!("Error changing user password: {}", e),
+                }
+                exit(0);
+            },
+            Some(("remove_user", sub_m)) => {
+                let username = sub_m.get_one::<String>("username").unwrap();
+
+                match state.db.remove_user(username).await {
+                    Ok(_) => println!("User removed successfully"),
+                    Err(e) => eprintln!("Error removing user: {}", e),
+                }
+                exit(0);
+            },
+            _ => {
+
+            }
+        }
+    };
+
+    rt.block_on(subcommand_future);
+
+    let tls_acceptor;
     if state.config.lumina.use_tls.unwrap_or_default() {
         let cert_path = &state.config.lumina.tls.as_ref().expect("tls section is missing").server_cert;
         let mut crt = match std::fs::read(cert_path) {
@@ -427,7 +502,7 @@ fn main() {
         };
 
         info!("listening on {:?} secure={}", server.local_addr().unwrap(), tls_acceptor.is_some());
-    
+
         serve(server, tls_acceptor, state, exit_signal_rx).await;
     };