Lots of networking work

nyash_server/src/config.rs

@@ -0,0 +1,31 @@
+
+
+use std::{ops::Not, path::Path, u64, u128};
+
+use serde::Deserialize;
+
+const DB_DIR: &str = "/var/lib/nyash-aes-xts256-plain64/";
+
+#[derive(Deserialize, Debug)]
+pub struct ServerConfig {
+    pub bind_addr: String,
+    pub listen_port: u16,
+    pub db_dir: Option<String>,
+    //keys: Keys,
+}
+
+// #[derive(Deserialize, Debug)]
+// struct Keys {
+//     github: String,
+//     travis: Option<String>,
+// }
+
+
+pub fn read_config(file_path: &Path) -> Result<ServerConfig, Box<dyn std::error::Error>> {
+    let config_str = std::fs::read_to_string(file_path)?;
+    let mut serer_config: ServerConfig = toml::from_str(&config_str)?;
+    serer_config.db_dir = serer_config.db_dir.or_else(|| Some(DB_DIR.to_string()));
+    println!("IP: {}", serer_config.bind_addr);
+    println!("Listen port: {:?}", serer_config.listen_port);
+    Ok(serer_config)
+}

nyash_server/src/database.rs

@@ -4,7 +4,7 @@ use redb::{
     Database, ReadableDatabase, ReadableTable, ReadableTableMetadata, TableDefinition
 };
 // use std::collections::HashMap;
-use std::{ops::Not,path::Path, u64, u128};
+use std::{ops::Not, path::Path, u64, u128};
 
 //tables defenition
 const DB_FILE: &str = "./data";
@@ -24,13 +24,13 @@ fn get_timestump() -> u64 {
 }
 
 #[derive(Copy, Clone, Debug)]
-struct JobRecord {
-    id: u64,
-    range_id: u16,
-    tweak_key: u128,
-    start_key: u128,
-    len: u64,
-    start_time: u64,
+pub struct JobRecord {
+    pub id: u64,
+    pub range_id: u16,
+    pub tweak_key: u128,
+    pub start_key: u128,
+    pub len: u64,
+    pub start_time: u64,
 }
 
 impl JobRecord {
@@ -344,7 +344,7 @@ impl<'a> RangesTable<'a> {
 
 }
 
-fn db_create_job(db: &Database, pref_job_size: u64) -> Result<Option<JobRecord>, redb::Error> {
+pub fn db_create_job(db: &Database, pref_job_size: u64) -> Result<Option<JobRecord>, redb::Error> {
     let transct = db.begin_write()?;
 
     // 1. look for abandonent jobs
@@ -410,7 +410,7 @@ fn db_create_job(db: &Database, pref_job_size: u64) -> Result<Option<JobRecord>,
 
 
 // returns false if no such job found in DB
-fn db_commit_job(db: &Database, job_id: u64) -> Result<bool, redb::Error> {
+pub fn db_commit_job(db: &Database, job_id: u64) -> Result<bool, redb::Error> {
     let transct: redb::WriteTransaction = db.begin_write()?;
 
     let job_found = {
@@ -444,7 +444,7 @@ fn db_commit_job(db: &Database, job_id: u64) -> Result<bool, redb::Error> {
 }
 
 
-fn db_get_progress(db: &Database) -> Result<f64, redb::Error> {
+pub fn db_get_progress(db: &Database) -> Result<f64, redb::Error> {
     let trx: redb::ReadTransaction = db.begin_read()?;
     let progress = {
         let ranges_t = trx.open_table(RANGES)?;
@@ -479,7 +479,7 @@ fn db_init(db: &Database) -> Result<(), redb::Error> {
     Ok(())
 }
 
-fn db_open(db_file_path: &Path) -> Result<Database, redb::Error> {
+pub fn db_open(db_file_path: &Path) -> Result<Database, redb::Error> {
     let db = Database::create(db_file_path)?;
     let trx = db.begin_read()?;
     let ranges_len = {

nyash_server/src/main.rs

@@ -1,35 +1,309 @@
+use clap::Parser;
+use tokio_tungstenite::tungstenite::Utf8Bytes;
+use std::str::FromStr;
+use std::{path::PathBuf};
+use log;
+
+use redb;
+use tokio::net::{TcpListener, TcpStream};
+use tokio_tungstenite::{accept_async, tungstenite::protocol::Message};
+use futures::{StreamExt, SinkExt};
+use std::env;
+use std::net::SocketAddr;
+
+
+mod config;
+
 
 mod num_utils;
 mod database;
 
 
-fn read_config() -> (){
-    
+/// 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,
 }
 
-fn main() {
-    println!("Hello, world!");
-    let a = 1855454745u128.to_le_bytes();
-    let b = 1u32.to_le_bytes();
-    //let c = u32::from_le_bytes(a[0..4].try_into().unwrap());
-    println!("{:?}", a);
-    println!("{:?}", b);
 
-    let (c,r) = a.as_chunks::<4>();
-    println!("{:?}", c);
-    println!("{:?}", r);
-    let a1 = 128u32;
-    let e = a1.min(u32::MAX);
-    println!("e = {}",e);
 
-    for i in 0..=u16::MAX {
-        println!("{}", i);
+#[derive(Debug)]
+enum ProtocolParseError {
+    ParseIntError(std::num::ParseIntError),
+    WrongFormat(usize),
+}
+
+impl From<std::num::ParseIntError> for ProtocolParseError {
+    fn from(err: std::num::ParseIntError) -> Self {
+        ProtocolParseError::ParseIntError(err)
+    }
+}
+
+impl From<usize> for ProtocolParseError {
+    fn from(err: usize) -> Self {
+        ProtocolParseError::WrongFormat(err)
+    }
+}
+
+#[derive(Copy, Clone, Debug)]
+struct NyashWorkRequest {
+    pref_job_size: u64,
+}
+
+impl NyashWorkRequest {
+    pub const HEADER: &str  = "WORK_REQUEST";
+
+    pub fn to_msg(&self) -> String {
+        return format!("{} {}", Self::HEADER, self.pref_job_size);
+    }
+
+    pub fn from_msg(mut it: std::str::Split<'_, &str>) -> Result<Self, ProtocolParseError> {
+        Ok(Self {
+            pref_job_size: it.next().ok_or(0)?.parse()?,
+        })
+    }
+}
+
+#[derive(Copy, Clone, Debug)]
+struct NyashWorkDone {
+    job_id: u64,
+}
+
+impl NyashWorkDone {
+    pub const HEADER: &str  = "WORK_DONE";
+
+    pub fn to_msg(&self) -> String {
+        return format!("{} {}", Self::HEADER, self.job_id);
+    }
+
+    pub fn from_msg(mut it: std::str::Split<'_, &str>) -> Result<Self, ProtocolParseError> {
+        Ok(Self {
+            job_id: it.next().ok_or(0)?.parse()?,
+        })
+    }
+}
+
+enum NyashReq {
+    NyashWorkRequest(NyashWorkRequest),
+    NyashWorkDone(NyashWorkDone),
+}
+
+
+#[derive(Copy, Clone, Debug)]
+struct NyashWorkResp {
+    job_id: u64,
+    tweak_key: u128,
+    start_key: u128,
+    len: u64,
+}
+
+impl NyashWorkResp {
+
+    pub const HEADER: &str  = "WORK_RESP";
+
+    pub fn from_job_record(jr: &database::JobRecord) -> Self {
+        Self {
+            job_id: jr.id,
+            tweak_key: jr.tweak_key,
+            start_key: jr.start_key,
+            len: jr.len,
+        }
+    }
+
+    pub fn to_msg(&self) -> String {
+        return format!("{} {} {} {} {}", Self::HEADER, self.job_id, self.tweak_key, self.start_key, self.len);
+    }
+
+    pub fn from_msg(mut it: std::str::Split<'_, &str>) -> Result<Self, ProtocolParseError> {
+        Ok(Self {
+            job_id: it.next().ok_or(0)?.parse()?,
+            tweak_key: it.next().ok_or(0)?.parse()?,
+            start_key: it.next().ok_or(0)?.parse()?,
+            len: it.next().ok_or(0)?.parse()?,
+        })
+    }
+}
+
+#[derive(Copy, Clone, Debug)]
+struct NyashErrorResp {
+    code: u32,
+}
+
+impl NyashErrorResp {
+    pub const HEADER: &str  = "ERROR";
+
+    pub fn to_msg(&self) -> String {
+        return format!("{} {}", Self::HEADER, self.code);
+    }
+
+    pub fn from_msg(mut it: std::str::Split<'_, &str>) -> Result<Self, ProtocolParseError> {
+        Ok(Self {
+            code: it.next().ok_or(0)?.parse()?,
+        })
+    }
+}
+
+#[derive(Copy, Clone, Debug)]
+struct NyashOkResp {}
+impl NyashOkResp {
+    pub const HEADER: &str  = "OK";
+
+    pub fn to_msg(&self) -> String {
+        return format!("{}", Self::HEADER);
+    }
+}
+
+#[derive(Copy, Clone, Debug)]
+struct NyashNoWorkResp {
+    code: u32,
+}
+
+impl NyashNoWorkResp {
+    pub const HEADER: &str  = "NO_WORK";
+
+    pub fn to_msg(&self) -> String {
+        return format!("{} {}", Self::HEADER, self.code);
+    }
+    pub fn from_msg(mut it: std::str::Split<'_, &str>) -> Result<Self, ProtocolParseError> {
+        Ok(Self {
+            code: it.next().ok_or(0)?.parse()?,
+        })
+    }
+}
+
+
+
+
+enum NyashResp {
+    NyashWorkResp(NyashWorkResp),
+    NyashErrorResp(NyashErrorResp),
+    NyashNoWorkResp(NyashNoWorkResp),
+}
+
+fn parse_message(mut it: std::str::Split<'_, &str>) -> Result<NyashReq, ProtocolParseError> {
+    let code: &str = it.next().ok_or(0)?;
+    match code {
+        NyashWorkRequest::HEADER => Ok(NyashReq::NyashWorkRequest(NyashWorkRequest::from_msg(it)?)),
+        NyashWorkDone::HEADER => Ok(NyashReq::NyashWorkDone(NyashWorkDone::from_msg(it)?)),
+        _ => Err(ProtocolParseError::WrongFormat(0))
+    }
+}
+
+async fn handle_connection(stream: TcpStream, db: &redb::Database) {
+    // Accept the WebSocket connection
+    let ws_stream = match accept_async(stream).await {
+        Ok(ws) => ws,
+        Err(e) => {
+            println!("Error during the websocket handshake: {}", e);
+            return;
+        }
+    };
+
+    // Split the WebSocket stream into a sender and receiver
+    let (mut sender, mut receiver) = ws_stream.split();
+
+    // Handle incoming messages
+    while let Some(msg) = receiver.next().await {
+        match msg {
+            Ok(Message::Text(text)) => {
+                // Reverse the received string and send it back
+                let str_text = text.to_string();
+                println!("Got text: {}", str_text);
+                //let mesg = parse_message(str_text.split(" "));
+                match parse_message(str_text.split(" ")) {
+                    Ok(NyashReq::NyashWorkDone(wrk_done)) => {
+                        let res = database::db_commit_job(db, wrk_done.job_id);
+                        match res {
+                            Ok(_) => {
+                                if let Err(e) = sender.send(Message::Text(NyashOkResp::HEADER.into())).await {
+                                    println!("Error sending message: {}", e);
+                                }
+                            },
+                            Err(ex) => {
+                                println!("Error comiting work: {}", ex);
+                                let rest_str = NyashErrorResp{code: 0}.to_msg();
+                                if let Err(e) = sender.send(Message::Text(rest_str.into())).await {
+                                    println!("Error sending message: {}", e);
+                                }
+                            }
+                        }
+                    },
+                    Ok(NyashReq::NyashWorkRequest(wrk_req)) => {
+                        let res = database::db_create_job(db, wrk_req.pref_job_size);
+                        match res {
+                            Ok(Some(job_record)) => {
+                                let rest_str = NyashWorkResp::from_job_record(&job_record).to_msg();
+                                if let Err(e) = sender.send(Message::Text(rest_str.into())).await {
+                                    println!("Error sending message: {}", e);
+                                }   
+                            },
+                            Ok(None) => {
+                                let rest_str = NyashNoWorkResp{code: 0}.to_msg();
+                                if let Err(e) = sender.send(Message::Text(rest_str.into())).await {
+                                    println!("Error sending message: {}", e);
+                                }
+                            },
+                            Err(ex) => {
+                                println!("Error creating job: {}", ex);
+                                let rest_str = NyashErrorResp{code: 0}.to_msg();
+                                if let Err(e) = sender.send(Message::Text(rest_str.into())).await {
+                                    println!("Error sending message: {}", e);
+                                }
+                            }
+                        }
+                    },
+                    Err(ex) => println!("Error sending message: {:?}", ex)
+                }
+                
+            },
+            Ok(Message::Close(_)) => break,
+            Ok(_) => (),
+            Err(e) => {
+                println!("Error processing message: {}", e);
+                break;
+            }
+        }
+    }
+}
+
+
+async fn accept_loop(listener: TcpListener, db: redb::Database) {
+    while let Ok((stream, _)) = listener.accept().await {
+        // Spawn a new task for each connection
+        handle_connection(stream, &db).await;
+    }
+
+}
+// Так. Нам нужно база данныъ с дипазонами всей фигни. И сервить её через tls или что-то типо того.
+
+#[tokio::main]
+async fn main() {
+    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: SocketAddr = addr.parse().expect("Invalid address");
+
+    // Create the TCP listener
+    let listener = TcpListener::bind(&addr).await.expect("Failed to bind");
+
+    println!("Listening on: {}", addr);
+    let res = tokio::spawn(accept_loop(listener, db)).await;
+    if res.is_err() {
+        print!("Something gone wrong! {:?}", res.unwrap_err());
     }
 
 }
 
-
-// Так. Нам нужно база данныъ с дипазонами всей фигни. И сервить её через tls или что-то типо того.
-
-
-

nyash_server/src/nyash_aes_d.conf

@@ -0,0 +1,4 @@
+# Nya
+
+bind_addr = "127.0.0.1"
+listen_port = 37939