Networking server work, added shared state

nyash_server/proto/nyash.proto

@@ -1,5 +1,5 @@
 syntax = "proto3";
-package nyash_server;
+package nyash_proto;
 
 service NyashLuks {
     // Requesting work, gettin in response work data, error, or message that there is no work currently awaylable.
@@ -74,5 +74,6 @@ message ProgressRequest {
 
 message ProgressReply {
     // Brutforce progress as ratio from 0 to 1(completed)
+    bool key_found = 1;
     double progress = 2;
 }

nyash_server/src/server.rs

@@ -1,17 +1,32 @@
+use std::ops::DerefMut;
+use std::str::FromStr;
+use std::sync::atomic::AtomicBool;
+use std::{path::PathBuf};
+
 
 use redb;
+use clap::Parser;
 
-use tonic::{transport::Server, Request, Response, Status};
+use std::sync::{Arc, Mutex};
+use tokio::sync::RwLock as AsyncRwLock;
+use tokio::time::{interval, Duration};
 
-use nyash_server::nyash_luks_server::{NyashLuks, NyashLuksServer};
+use tonic::{Request, Response, Status, transport::Server};
-use nyash_server::{WorkRequest, WorkData, WorkReply, work_reply, WorkCommit, CommitReply, ProgressRequest, ProgressReply};
 
-pub mod nyash_server {
+
-    tonic::include_proto!("nyash_server"); // The string specified here must match the proto package name
+pub mod nyash_proto {
+    tonic::include_proto!("nyash_proto"); // The string specified here must match the proto package name
 }
 
+use nyash_proto::nyash_luks_server::{NyashLuks, NyashLuksServer};
+use nyash_proto::{
+    CommitReply, KeyData, ProgressReply, ProgressRequest, WorkCommit, WorkData, WorkReply, WorkRequest, work_commit,
+    work_reply,
+};
+
 mod database;
 mod num_utils;
+mod config;
 
 fn work_rec_to_work_data(work_rec: database::JobRecord) -> WorkData {
     let tw_key = num_utils::u128_to_u64arr(work_rec.tweak_key);
@@ -28,10 +43,48 @@ fn work_rec_to_work_data(work_rec: database::JobRecord) -> WorkData {
     }
 }
 
-#[derive(Debug, Default)]
+// Shared state
+// #[derive(Clone, Debug)]
+// pub struct ServiceState {
+//     pub db: Arc<redb::Database>,
+//     pub key_found: bool,
+//     pub progress: f64,
+// }
+
+// Background task to update state
+// async fn update_service_state(state: Arc<AsyncRwLock<ServiceState>>) {
+//     loop {
+//         tokio::time::sleep(tokio::time::Duration::from_secs(5)).await;
+//         let mut guard = state.write().await;
+//         match database::db_get_progress(&guard.db) {
+//             Ok(p) => {
+//                 guard.progress = p;
+//             },
+//             Err(_) => println!("Error getting progress from DB")
+//         }
+//     }
+// }
+
+// Background task to update progress from DB
+async fn update_service_progress(progress: Arc<AsyncRwLock<f64>>, db: Arc<redb::Database>) {
+    loop {
+        tokio::time::sleep(tokio::time::Duration::from_secs(5)).await;
+        match database::db_get_progress(&db) {
+            Ok(p) => {
+                let mut guard = progress.write().await;
+                *guard = p;
+            },
+            Err(_) => println!("Error getting progress from DB")
+        }
+    }
+}
+
+#[derive(Debug)]
 pub struct NyashService {
-    db: redb::Database,
+    //state: Arc<AsyncRwLock<ServiceState>>,
-    progress: f64,
+    db: Arc<redb::Database>,
+    key_found: Arc<AsyncRwLock<bool>>,
+    progress: Arc<AsyncRwLock<f64>>,
 }
 
 #[tonic::async_trait]
@@ -39,19 +92,20 @@ impl NyashLuks for NyashService {
     async fn request_work(&self, request: Request<WorkRequest>) -> Result<Response<WorkReply>, Status> {
         println!("Got a request: {:?}", request);
         let pref_work_size: u64 = request.into_inner().pref_work_size;
+        
         match database::db_create_job(&self.db, pref_work_size) {
             Ok(Some(work_rec)) => {
                 let reply = WorkReply {
                     result: Some(work_reply::Result::WorkData(work_rec_to_work_data(work_rec))),
                 };
                 Ok(Response::new(reply)) // Send back our formatted greeting
-            },
+            }
             Ok(None) => {
                 let reply = WorkReply {
                     result: Some(work_reply::Result::NoWork(true)),
                 };
                 Ok(Response::new(reply))
-            },
+            }
             Err(_ex) => {
                 let reply = WorkReply {
                     result: Some(work_reply::Result::Error("DB Error!".to_string())),
@@ -63,29 +117,113 @@ impl NyashLuks for NyashService {
 
     async fn commit_work(&self, request: Request<WorkCommit>) -> Result<Response<CommitReply>, Status> {
         println!("Got a request: {:?}", request);
-        let work_id: u64 = request.into_inner().work_id;
+        let work_commit = request.into_inner();
-        match database::db_commit_job(&self.db, work_id) {
+
-            Ok(true) => Ok(Response::new(CommitReply {status_code: 0, progress: self.progress})),
+        let put_k_no_err: bool = match work_commit.result {
-            Ok(false) => Ok(Response::new(CommitReply {status_code: 1, progress: self.progress})),
+            Some(work_commit::Result::FoundKey(key_data)) => {
-            Err(_) => Err(Status::internal("DB error"))
+                {
+                    let mut guard = self.key_found.write().await;
+                    *guard = true;
+                }
+                let tw_k = num_utils::u64arr_to_u128((key_data.tweak_key0, key_data.tweak_key1));
+                let en_k = num_utils::u64arr_to_u128((key_data.start_key0, key_data.start_key1));
+                match database::db_put_found_key(&self.db, tw_k, en_k) {
+                    Ok(_) => true,
+                    Err(_) => false,
+                }
+            }
+            _ => true,
+        };
+
+        let comm_no_err: bool = match database::db_commit_job(&self.db, work_commit.work_id) {
+            Ok(true) => true,
+            Ok(false) => false,
+            Err(_) => false,
+        };
+
+        if (put_k_no_err && comm_no_err) == true {
+            Ok(Response::new(CommitReply {
+                status_code: 0,
+                progress: *self.progress.read().await
+            }))
+        } else {
+            Err(Status::internal("DB error"))
         }
     }
 
     async fn request_progress(&self, _request: Request<ProgressRequest>) -> Result<Response<ProgressReply>, Status> {
-        Ok(Response::new(ProgressReply {progress: self.progress}))
+        Ok(Response::new(ProgressReply {
+            key_found: *self.key_found.read().await,
+            progress: *self.progress.read().await,
+        }))
     }
 }
 
 
+/// Server for AES-XTS distributed bruteforce
+#[derive(clap::Parser, Debug)]
+#[command(version, about, long_about = None)]
+struct ProgArgs {
+    /// Config file name
+    #[arg(short, long, default_value = "/etc/nyash_aes_d.conf")]
+    config: PathBuf,
+}
+
+
+
 #[tokio::main]
 async fn main() -> Result<(), Box<dyn std::error::Error>> {
-    let addr = "[::1]:50051".parse()?;
+
-    let greeter = MyGreeter::default();
+    let args = ProgArgs::parse();
+    log::error!("Arguments parsed: {:?}", args.config);
+    println!("Path: {:?}", args.config);
+    let server_config: config::ServerConfig = config::read_config(&args.config).unwrap();
+
+    let db_dir = server_config.db_dir.unwrap();
+    let db_file = format!("{}data.rdb",db_dir);
+    let db_path = PathBuf::from_str(db_file.as_str()).unwrap();
+    
+    let db = database::db_open(&db_path).expect("Error Opening database");
+
+    // Get the address to bind to
+    let addr = format!("{}:{}",server_config.bind_addr, server_config.listen_port);
+    let addr: std::net::SocketAddr = addr.parse().expect("Invalid address");
+
+    println!("Listening on: {}", addr);
+
+
+    let found_keys = database::db_get_found_keys(&db).expect("Error getting found keys from DB");
+    if found_keys.len() >0 {
+        println!("Key already found!");
+        for key in found_keys {
+            println!("Key found: {:?}", key);
+        }
+        println!("Exit!");
+        return Ok(());
+    }
+
+
+    let progress = database::db_get_progress(&db).expect("Error getting progress from DB");
+    
+    let shared_progress = Arc::new(AsyncRwLock::new(progress));
+    let shared_db = Arc::new(db);
+    let shared_key_found = Arc::new(AsyncRwLock::new(false));
+
+    // let service_state = Arc::new(AsyncRwLock::new(ServiceState {
+    //     db: Arc::new(db),
+    //     key_found: false,
+    //     progress: progress,
+    // }));
+
+    // Spawn the periodic updater
+    tokio::spawn(update_service_progress(shared_progress.clone(), shared_db.clone()));
+    
+    let nyash_service = NyashService{db:shared_db, progress:shared_progress, key_found:shared_key_found};
+
 
     Server::builder()
-        .add_service(GreeterServer::new(greeter))
+        .add_service(NyashLuksServer::new(nyash_service))
         .serve(addr)
         .await?;
-
     Ok(())
 }