kira/nyash-aes-xts256-plain64
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Networking server work, added shared state

Browse Source
This commit is contained in:
Kirill Shakirov
2026-03-07 17:07:54 +01:00
parent fc356982c7
commit 79852c6c3f
2 changed files with 163 additions and 24 deletions
Download Patch File Download Diff File Expand all files Collapse all files
nyash_server/proto/nyash.proto
+2 -1
View File
@@ -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
+161 -23
View File
@@ -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 nyash_server::{WorkRequest, WorkData, WorkReply, work_reply, WorkCommit, CommitReply, ProgressRequest, ProgressReply};
use tonic::{Request, Response, Status, transport::Server};
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,64 +43,187 @@ 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,
progress: f64,
//state: Arc<AsyncRwLock<ServiceState>>,
db: Arc<redb::Database>,
key_found: Arc<AsyncRwLock<bool>>,
progress: Arc<AsyncRwLock<f64>>,
}
#[tonic::async_trait]
impl NyashLuks for NyashService {
async fn request_work(&self, request: Request<WorkRequest>) -> Result<Response<WorkReply>, Status> {
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())),
};
Ok(Response::new(reply))
}
}
}
}
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;
match database::db_commit_job(&self.db, work_id) {
Ok(true) => Ok(Response::new(CommitReply {status_code: 0, progress: self.progress})),
Ok(false) => Ok(Response::new(CommitReply {status_code: 1, progress: self.progress})),
Err(_) => Err(Status::internal("DB error"))
let work_commit = request.into_inner();
let put_k_no_err: bool = match work_commit.result {
Some(work_commit::Result::FoundKey(key_data)) => {
{
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(())
}
}