Romove big files. Wrapping things up.

.gitignore

@@ -19,3 +19,7 @@ target
 #  and can be added to the global gitignore or merged into this file.  For a more nuclear
 #  option (not recommended) you can uncomment the following to ignore the entire idea folder.
 #.idea/
+
+# ignore test data
+*.img
+*.key

nyash_client/Cargo.toml

@@ -2,8 +2,26 @@
 name = "nyash_client"
 version = "0.1.0"
 edition = "2024"
+build = "build.rs"
+
+[[bin]] # Bin to run the HelloWorld gRPC client
+name = "nyash-client"
+path = "src/client.rs"
 
 [dependencies]
-ocl = "0.19"
+flate2 = "1.1.9"
+ocl = { version = "0.19" }
+prost = "0.14.3"
 serde = { version = "1.0.228", features = ["derive"] }
 serde_json = "1.0.149"
+tokio = { version = "1.50.0", features = ["rt", "macros", "signal"] }
+tonic = "0.14.5"
+tonic-prost = "0.14.5"
+
+[build-dependencies]
+tonic-prost-build = "0.14.5"
+ocl-include = "0.6"
+flate2 = "1.1.9"
+
+
+

nyash_client/build.rs

@@ -0,0 +1,55 @@
+use std::env;
+use std::io::Write;
+use std::path::Path;
+use ocl_include;
+use flate2::write::GzEncoder;
+use flate2::Compression;
+
+fn main() -> Result<(), Box<dyn std::error::Error>> {
+    //compile gRPC
+    tonic_prost_build::compile_protos("proto/nyash.proto")?;
+
+
+    //******************* compile ocl into spirv64 *********
+    // let manifest_dir = env::var("CARGO_MANIFEST_DIR").unwrap();
+    // let ocl_src_dir = Path::new(&manifest_dir).join("src/open_cl");
+    // let spirv_build_script = ocl_src_dir.clone().join("build_spirv.sh");
+
+    
+    // let out_dir = env::var("OUT_DIR").unwrap();
+    // let spirv_bin_path = Path::new(&out_dir).join("nyash_aes_xts256_plain.spv");
+    // let str_spirv_bit_p = spirv_bin_path.to_str().expect("Error converting spirv out path to str!");
+    // let _output = Command::new(spirv_build_script)
+    //     .current_dir(ocl_src_dir)
+    //     .arg(str_spirv_bit_p)
+    //     .output()
+    //     .expect("Failed to execute spirv build script!");
+
+    // ***************** concat ocl sources to one file and zip it **********************
+    let manifest_dir = env::var("CARGO_MANIFEST_DIR").unwrap();
+    let ocl_src_dir = Path::new(&manifest_dir).join("src/open_cl");
+    let ocl_src_file = ocl_src_dir.clone().join("nyash_aes_xts256_plain.cl");
+
+    let ocl_parser = ocl_include::Parser::builder()
+        .add_source(
+            ocl_include::source::Fs::builder()
+                .include_dir(ocl_src_dir)
+                .expect("Error adding ocl include dir!")
+                .build(),
+        )
+        .build();
+    println!("Concatenating and compressing ocl sources...");
+    let ocl_node = ocl_parser.parse(&ocl_src_file).expect("Error parsing ocl source!");
+    let full_src = ocl_node.collect().0;
+    let mut gz_encoder = GzEncoder::new(Vec::new(), Compression::best());
+
+    gz_encoder.write_all(full_src.as_bytes()).expect("Error compressing src file!");
+    
+    let compressed_bytes = gz_encoder.finish().expect("Error compressing src!");
+
+    let out_dir = env::var("OUT_DIR").unwrap();
+    let ocl_concat_src_path = Path::new(&out_dir).join("nyash_aes_full.cl.gz");
+    std::fs::write(ocl_concat_src_path, compressed_bytes).expect("Error writing compressed src file!");
+
+    Ok(())
+}

nyash_client/nyash_conf.json

@@ -0,0 +1,12 @@
+{
+  "devices": [
+    {
+      "dev_name": "gfx1103",
+      "platform_name": "AMD Accelerated Parallel Processing",
+      "id": 0,
+      "work_size": 4096,
+      "batch_size": 5412812
+    }
+  ],
+  "dev_fill": 100
+}

nyash_client/proto/nyash.proto

@@ -0,0 +1,80 @@
+syntax = "proto3";
+package nyash_proto;
+
+// serving key brutforcing for aes xts 
+service NyashLuks {
+    // Requesting work, gettin in response work data, error, or message that there is no work currently awaylable.
+    rpc RequestWork (WorkRequest) returns (WorkReply);
+    // Commiting work, with work ID. If key was found, sending also discowered key
+    rpc CommitWork (WorkCommit) returns (CommitReply);
+    // Request owerall key search progress
+    rpc RequestProgress (ProgressRequest) returns (ProgressReply);
+}
+
+message WorkRequest {
+    // Request work from server
+    // Preffered work size in number of keys
+    uint64 pref_work_size = 1;
+}
+
+message WorkReply {
+    oneof result {
+        WorkData work_data = 1;
+        bool no_work = 2;
+        string error = 3;
+    }
+}
+
+message WorkData {
+    // work id
+    uint64 work_id =1;
+    // work size - number of keys to check
+    uint64 work_size= 2;
+    // tweak key is 128 bit. We send it as two uint64 values
+    uint64 tweak_key0 =3;
+    uint64 tweak_key1 =4;
+    // start key is 128 bit. We send it as two uint64 values
+    uint64 start_key0 =5;
+    uint64 start_key1 =6;
+}
+
+message WorkCommit {
+    // work id
+    uint64 work_id =1;
+    oneof result {
+        bool no_key = 2;
+        KeyData found_key = 3;
+    }
+
+}
+
+message KeyData {
+    // tweak key is 128 bit. We send it as two uint64 values
+    uint64 tweak_key0 =3;
+    uint64 tweak_key1 =4;
+    // start key is 128 bit. We send it as two uint64 values
+    uint64 start_key0 =5;
+    uint64 start_key1 =6;
+}
+
+
+message CommitReply {
+    // Status code 
+    // 0 - means ok
+    // 1 - something went wrong on server
+    // 2 - something wrong with request
+    uint32 status_code = 1;
+    // Brutforce progress as ratio from 0 to 1(completed)
+    double progress = 2;
+}
+
+// request total progress
+message ProgressRequest {
+  
+}
+
+message ProgressReply {
+    // Brutforce progress as ratio from 0 to 1(completed)
+    bool key_found = 1;
+    double progress = 2;
+}

nyash_client/src/client.rs

@@ -0,0 +1,344 @@
+use std::time::Duration;
+
+use tokio;
+use tonic;
+use tonic::transport::Channel;
+//use tonic::{Request, Response, Status};
+pub mod nyash_proto {
+    tonic::include_proto!("nyash_proto"); // The string specified here must match the proto package name
+}
+
+use nyash_proto::nyash_luks_client::NyashLuksClient;
+use nyash_proto::{
+    CommitReply, KeyData, ProgressReply, ProgressRequest, WorkCommit, WorkReply,
+    WorkRequest, work_commit, work_reply,
+};
+
+use std::sync::{Arc};
+use tokio::sync::RwLock as AsyncRwLock;
+
+use crate::ocl_utils::ExecData;
+
+pub mod client_config;
+pub mod num_utils;
+pub mod ocl_utils;
+mod search_params;
+//mod test_cl;
+
+const S_ADDR: &str = "http://127.0.0.1:37939";
+// const SRC_PATH: &str = "src/open_cl/nyash_aes_xts256_plain.cl";
+// const OCL_COMP_OPT: &str = "-I src/open_cl";
+const CONF_RILE_NAME: &str = "nyash_conf.json";
+
+
+async fn shutdown_signal() {
+    let ctrl_c = async {
+        tokio::signal::ctrl_c()
+            .await
+            .expect("Failed to install Ctrl+C handler");
+    };
+
+    #[cfg(unix)]
+    let terminate = async {
+        tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate())
+            .expect("Failed to install SIGTERM handler")
+            .recv()
+            .await;
+    };
+
+    #[cfg(not(unix))]
+    let terminate = std::future::pending::<()>();
+
+    tokio::select! {
+        _ = ctrl_c => println!("Received Ctrl+C, initiating shutdown"),
+        _ = terminate => println!("Received SIGTERM, initiating shutdown"),
+    }
+}
+
+fn key_dat_from_exec_dat(ex_dat: &ExecData) -> KeyData {
+    let t_k = num_utils::u128_to_u64arr(num_utils::u32arr_to_u128(num_utils::vec_to_u32_4arr(
+        &ex_dat.tweak_key,
+        0,
+    )));
+    let e_k = num_utils::u128_to_u64arr(ex_dat.get_found_key());
+
+    KeyData {
+        start_key0: e_k[0],
+        start_key1: e_k[1],
+        tweak_key0: t_k[0],
+        tweak_key1: t_k[1],
+    }
+}
+
+async fn get_progress(chanel: Channel) -> Result<ProgressReply, tonic::Status> {
+    let mut client = NyashLuksClient::new(chanel);
+
+    let request = tonic::Request::new(ProgressRequest {});
+    let response = client.request_progress(request).await?;
+    return Ok(response.into_inner());
+}
+
+async fn get_work(chanel: Channel, work_size: u64) -> Result<WorkReply, tonic::Status> {
+    let mut client = NyashLuksClient::new(chanel);
+
+    println!("Requesting work {} keys...", work_size);
+    let request = tonic::Request::new(WorkRequest {
+        pref_work_size: work_size,
+    });
+    let response = client.request_work(request).await?;
+    return Ok(response.into_inner());
+}
+
+async fn commit_work(
+    chanel: Channel,
+    commit_data: WorkCommit,
+) -> Result<CommitReply, tonic::Status> {
+    let mut client = NyashLuksClient::new(chanel);
+
+    let request = tonic::Request::new(commit_data);
+    let response = client.commit_work(request).await?;
+    return Ok(response.into_inner());
+}
+
+fn benchmark(exec_context: &mut ocl_utils::ExecContext) -> (u64, usize) {
+    let mut nyan_exec_dat = ocl_utils::ExecData {
+        start_key: vec![1u32; 4],
+        tweak_key: vec![2u32; 4],
+        uenc_data: vec![0u32; 4],
+        target_data: vec![0u32; 4],
+        tweak_i: 0,
+        tweak_j: 0,
+        key_found: vec![0u32; 5],
+        batch_size: 10000000,
+        work_size: 128,
+    };
+
+    let total_work: u64 = 1280000000;
+    let work_sizes: [usize; 8] = [128, 256, 512, 1024, 2048, 4096, 8192, 16384];
+    let mut work_time = [0f64; 8];
+
+    ocl_utils::set_target_data(exec_context, &mut nyan_exec_dat).expect("Error set target data!");
+
+    let mut preffered_work_size: usize = work_sizes[0];
+    let mut preffered_batch_size: u64 = 0;
+    for i in 0..8 {
+        let test_work_s = work_sizes[i];
+        let batch_size: u64 = total_work / test_work_s as u64;
+        nyan_exec_dat.work_size = test_work_s;
+        nyan_exec_dat.batch_size = batch_size;
+        println!("Benchmarking work size {}", test_work_s);
+        for _j in 0..3 {
+            let (_, exec_time) =
+                ocl_utils::do_work(exec_context, &mut nyan_exec_dat).expect("Error running tests!");
+            work_time[i] += exec_time;
+        }
+        work_time[i] = work_time[i] / 3.0;
+        println!("Average time {}", work_time[i]);
+        if i > 0 {
+            //giving 5% error for speed mesure
+            if (work_time[i]*1.05) > work_time[i - 1] {
+                break;
+            }
+        }
+        preffered_work_size = work_sizes[i];
+        // calculate batch size so it correspond to 30 sec job
+        preffered_batch_size = (batch_size as f64 * (20.0 / work_time[i])) as u64;
+        println!("batch_size {}, work_time {}, preffered_batch_size {}, preffered_work_size {}",
+        batch_size,
+        work_time[i],
+        preffered_batch_size,
+        preffered_work_size);
+    }
+
+    return (preffered_batch_size, preffered_work_size);
+}
+
+#[tokio::main(flavor = "current_thread")]
+async fn main() -> Result<(), Box<dyn std::error::Error>> {
+    const OCL_NYAS_GZ_SRC: &[u8] =
+        include_bytes!(concat!(env!("OUT_DIR"), "/nyash_aes_full.cl.gz"));
+
+    println!("OCL bin src lenght {}", OCL_NYAS_GZ_SRC.len());
+    println!("Hello, world nya!");
+    //use ocl::{Buffer, Context, Device, Kernel, Platform, Program, Queue, flags};
+    let (devices, mut app_conf) =
+        client_config::get_devices_conf(CONF_RILE_NAME).expect("Error loading config!");
+
+    let nyash_dev = devices[0].0;
+    let nyash_plt = devices[0].1;
+//
+    println!(
+        "Platform: {:?}, Device: {:?}",
+        nyash_plt.name().unwrap(),
+        nyash_dev.name().unwrap()
+    );
+
+    // reading ocl program sources
+    //let prog_src = std::fs::read_to_string(SRC_PATH).expect("Error reading program sources!");
+
+    let mut exec_context: ocl_utils::ExecContext = ocl_utils::ExecContext::new(
+        nyash_dev,
+        nyash_plt,
+        OCL_NYAS_GZ_SRC,
+        1024
+    )
+    .expect("Error creating exec context!");
+
+    // need to train in order to learn optimal params
+   
+    // need to train in order to learn optimal params
+    if (app_conf.devices[0].batch_size == 0) || (app_conf.devices[0].work_size == 0) {
+        println!("Performing banchmark to determine optimal GPU parameters...");
+        let (batch_size, work_size) = benchmark(&mut exec_context);
+        println!("batch_size {}, work_size {}", batch_size, work_size);
+        app_conf.devices[0].batch_size = batch_size;
+        app_conf.devices[0].work_size = work_size;
+        client_config::save_config(CONF_RILE_NAME, &app_conf).expect("Error saving config!");
+    }
+
+    let nyash_dev_cfg = &app_conf.devices[0];
+    println!(
+        "Preffered Work size: {}, Batch size {}",
+        nyash_dev_cfg.work_size, nyash_dev_cfg.batch_size
+    );
+
+    let (_, _, encrypted_data) = search_params::get_params();
+    //setting data
+    let mut nyan_exec_dat = ocl_utils::ExecData {
+        start_key: Vec::new(),
+        tweak_key: Vec::new(),
+        uenc_data: vec![0u32; 4],
+        target_data: encrypted_data.to_vec(),
+        tweak_i: 0,
+        tweak_j: 0,
+        key_found: vec![0u32; 5],
+        batch_size: nyash_dev_cfg.batch_size,
+        work_size: nyash_dev_cfg.work_size,
+    };
+     println!(
+        "nyan_exec_dat Work size: {}, Batch size {}",
+        nyan_exec_dat.work_size, nyan_exec_dat.batch_size
+    );
+
+
+    ocl_utils::set_target_data(&mut exec_context, &mut nyan_exec_dat)
+        .expect("Error setting target data!");
+
+    // Don't keep connection alive when idle
+    let nya_channel: tonic::transport::Channel = tonic::transport::Endpoint::from_static(S_ADDR)
+        .keep_alive_while_idle(false)
+        .keep_alive_timeout(Duration::from_secs(10))
+        .connect_timeout(Duration::from_secs(10))
+        .connect()
+        .await
+        .expect("Error connecting to server!");
+
+    let key_found = match get_progress(nya_channel.clone()).await {
+        Err(_) => {
+            println!("Error getting progress!");
+            false
+        }
+        Ok(p_r) => {
+            println!("Current progress {:.8}%", p_r.progress * 100.0);
+            p_r.key_found
+        }
+    };
+
+    let mut giga_keys_per_second: f64 = 0f64;
+    let req_work_size: u64 = nyash_dev_cfg.batch_size as u64 * nyash_dev_cfg.work_size as u64;
+
+
+    let shared_key_found = Arc::new(AsyncRwLock::new(key_found));
+    // handling program termination
+    let sh_k_f_clone = shared_key_found.clone();
+    tokio::spawn(async move {
+        shutdown_signal().await;
+        //signaling that we should stop
+        let mut guard = sh_k_f_clone.write().await;
+        *guard = true;
+    });
+
+    while *shared_key_found.read().await == false {
+        let mut work = get_work(nya_channel.clone(), req_work_size).await;
+        while work.is_err() {
+            println!("Error getting work, waiting 10 seconds and trying again...");
+            tokio::time::sleep(Duration::from_secs(10)).await;
+            work = get_work(nya_channel.clone(), req_work_size).await;
+        }
+        let work = work.expect("Error getting work!");
+
+        let work_data = match work.result.expect("Error! Expected WorkResult!") {
+            work_reply::Result::NoWork(_) => {
+                println!("No work right now, try again later...");
+                continue;
+            }
+            work_reply::Result::Error(ex) => {
+                println!("Erro getting work: {}", ex);
+                continue;
+            }
+            work_reply::Result::WorkData(wd) => wd,
+        };
+        println!("Got work, {} keys...", work_data.work_size);
+        nyan_exec_dat.start_key = num_utils::u128_to_u32arr(num_utils::u64arr_to_u128([
+            work_data.start_key0,
+            work_data.start_key1,
+        ]))
+        .to_vec();
+        nyan_exec_dat.tweak_key = num_utils::u128_to_u32arr(num_utils::u64arr_to_u128([
+            work_data.tweak_key0,
+            work_data.tweak_key1,
+        ]))
+        .to_vec();
+
+        let mut batch_size = work_data.work_size / nyan_exec_dat.work_size as u64;
+        if (work_data.work_size % nyan_exec_dat.work_size as u64) != 0 {
+            batch_size += 1;
+        }
+        println!("Setting batch size to {}", batch_size);
+        nyan_exec_dat.batch_size = batch_size;
+
+        println!("Crunching numbers...");
+        match ocl_utils::do_work(&mut exec_context, &mut nyan_exec_dat) {
+            Err(_) => println!("Error doing work!"),
+            Ok((k_f, work_time)) => {
+                let mut w_k = WorkCommit {
+                    work_id: work_data.work_id,
+                    result: Some(work_commit::Result::NoKey(true)),
+                };
+
+                let g_k_p_s = (work_data.work_size as f64 / work_time) / 1000000000.0;
+                if giga_keys_per_second != 0f64 {
+                    giga_keys_per_second -= giga_keys_per_second / 10.0;
+                    giga_keys_per_second += g_k_p_s / 10.0;
+                } else {
+                    giga_keys_per_second = g_k_p_s;
+                }
+                println!("Average speed: {:.3}GigaKeys/Sec", giga_keys_per_second);
+                if k_f == true {
+                    println!(
+                        "We found the key! {:?} {:?}",
+                        nyan_exec_dat.key_found, nyan_exec_dat.tweak_key
+                    );
+
+                    //signaling that key found
+                    let mut guard = shared_key_found.write().await;
+                    *guard = true;
+
+                    w_k.result = Some(work_commit::Result::FoundKey(key_dat_from_exec_dat(
+                        &nyan_exec_dat,
+                    )));
+                }
+
+                let resp = commit_work(nya_channel.clone(), w_k).await;
+                match resp {
+                    Ok(c_r) => println!("Work commited. Progress: {:.8}%", c_r.progress * 100.0),
+                    Err(_) => println!("Error commiting work..."),
+                };
+            }
+        };
+    }
+
+    println!("Exiting!");
+
+    Ok(())
+}

nyash_client/src/client_config.rs

@@ -1,7 +1,8 @@
-use ocl::{Device, Platform};
+use ocl::{Device, Platform, DeviceType, flags};
 use serde::{Deserialize, Serialize};
 use serde_json;
 use std::error::Error;
+use std::{io};
 
 #[derive(Clone, Deserialize, Serialize, Debug)]
 pub struct DevConf {
@@ -9,7 +10,7 @@ pub struct DevConf {
     pub platform_name: String,
     pub id: usize,
     pub work_size: usize,
-    pub batch_size: u32,
+    pub batch_size: u64,
 }
 
 impl DevConf {
@@ -68,7 +69,130 @@ pub fn load_config(file_name: &str) -> Result<AppConfig, Box<dyn Error>> {
 }
 
 pub fn save_config(file_name: &str, app_conf: &AppConfig) -> Result<(), Box<dyn Error>> {
-    let conf_str = serde_json::to_string(app_conf)?;
+    let conf_str = serde_json::to_string_pretty(app_conf)?;
     std::fs::write(file_name, conf_str)?;
     return Ok(());
 }
+
+fn dev_type_from_str(s: &str) -> Result<flags::DeviceType, ()> {
+    match s {
+        "CPU" => Ok(flags::DeviceType::CPU),
+        "GPU" => Ok(flags::DeviceType::GPU),
+        "ALL" => Ok(flags::DeviceType::ALL),
+        "CUSTOM" => Ok(flags::DeviceType::CUSTOM),
+        "ACCELERATOR" => Ok(flags::DeviceType::ACCELERATOR),
+        "DEFAULT" => Ok(flags::DeviceType::DEFAULT),
+        _ => Err(()),
+    }
+}
+
+fn str_or_empty(r: ocl::error::Result<String>) -> String {
+    match r {
+        Ok(s) => s,
+        Err(_) => "".to_string(),
+    }
+}
+
+fn print_devices(dev_list: &Vec<(Device, Platform)>) {
+    let mut i = 0;
+    for (dev, plt) in dev_list.iter() {
+        let dev_name = str_or_empty(dev.name());
+        let plt_name = str_or_empty(plt.name());
+        println!("({i}) device: \"{dev_name}\" ----- platorm: \"{plt_name}\"");
+        i += 1;
+    }
+}
+
+fn choose_devices(devices_num: usize) -> Result<Vec<usize>, String> {
+    println!("Please input desired device to use as a number and press Enter.");
+    let mut result: Vec<usize> = Vec::new();
+
+    let mut s_devs_nums = String::new();
+
+    io::stdin()
+        .read_line(&mut s_devs_nums)
+        .expect("Failed to read line");
+
+    for s_dev_num in s_devs_nums.split(' ') {
+        let dev_num: usize = match s_dev_num.trim().parse() {
+            Ok(num) => num,
+            Err(_) => return Err("You must input a number from device list.".to_string()),
+        };
+        if dev_num >= devices_num {
+            return Err("You must input a number from device list.".to_string());
+        };
+        result.push(dev_num);
+    }
+    return Ok(result);
+}
+
+fn list_devices(dev_type: DeviceType) -> Vec<(Device, Platform)> {
+    let platforms = Platform::list();
+    let mut devices: Vec<(Device, Platform)> = Vec::new();
+    for plt in platforms.iter() {
+        //let plat_name = str_or_empty(plt.name());
+        let list_res = Device::list(plt, Some(dev_type));
+        match list_res {
+            Ok(dev_l) => devices.extend(dev_l.iter().map(|dev| (*dev, plt.clone()))),
+            Err(_) => {}
+        }
+    }
+    return devices;
+}
+
+
+fn dev_sel_dialog(all_devices: &Vec<(Device, Platform)>) -> Vec<usize> {
+    let devs_nums = loop {
+        print_devices(&all_devices);
+        match choose_devices(all_devices.len()) {
+            Ok(value) => break value,
+            Err(exc) => {
+                println!("Error! {exc}\n")
+            }
+        }
+    };
+    return devs_nums;
+}
+
+pub fn get_devices_conf(file_name: &str) -> Result<(Vec<(Device, Platform)>, AppConfig), String> {
+    let dev_type = dev_type_from_str("GPU").expect("Unexpected device type!");
+
+    // Get devices to be used for key search
+    let all_devices: Vec<(Device, Platform)> = list_devices(dev_type);
+    if all_devices.len() == 0 {
+        return Err("Cannot find any usable devices.".to_string());
+    };
+
+    let app_conf = match load_config(file_name) {
+        Ok(readed_config) => {
+            let dev_found = all_devices
+                .iter()
+                .filter(|dp| readed_config.device_exist(&dp.0))
+                .count();
+            if dev_found < readed_config.devices.len() {
+                println!("Devices from config not found in the system!");
+                let devs_nums = dev_sel_dialog(&all_devices);
+                let res = AppConfig::from_dev_list(&all_devices, devs_nums);
+                save_config(file_name, &res).expect("Error saving config!");
+                res
+            } else {
+                readed_config
+            }
+        }
+        Err(_) => {
+            println!("Cannot find config file {}", file_name);
+            let devs_nums = dev_sel_dialog(&all_devices);
+            let res = AppConfig::from_dev_list(&all_devices, devs_nums);
+            save_config(file_name, &res).expect("Error saving config!");
+            res
+        }
+    };
+
+    let selected_devs = all_devices
+        .iter()
+        .filter(|dp| app_conf.device_exist(&dp.0))
+        .cloned()
+        .collect();
+
+    return Ok((selected_devs, app_conf));
+}

nyash_client/src/num_utils.rs

@@ -3,265 +3,334 @@
 //     i: [u32; 2],
 // }
 
-fn add_u32_to_u256(a: &[u32; 8], b: u32) -> ([u32; 8], bool) {
+pub fn vec_to_u32_4arr(in_v: &Vec<u32>, start_idx: usize) -> [u32; 4] {
-    let mut res: [u32; 8] = [0; 8];
+    let mut u32_arr_k = [0u32; 4];
-    let mut carry = false;
+    for i in 0..4 {
-    (res[0], carry) = a[0].carrying_add(b, carry);
+        u32_arr_k[0] = in_v[start_idx + i];
-
-    for idx in 1..8 {
-        (res[idx], carry) = a[idx].carrying_add(0, carry);
     }
-
+    return u32_arr_k;
-    return (res, carry);
 }
 
-fn add_u32_to_u256_(a: &mut [u32; 8], b: u32) -> bool {
+pub fn u128_to_u64arr(a: u128) -> [u64; 2] {
-    let mut carry = false;
+    let mut res = [0u64; 2];
-    (a[0], carry) = a[0].carrying_add(b, carry);
+    let a_bytes = a.to_le_bytes();
-
+    let chunks = a_bytes.as_chunks::<8>().0;
-    for idx in 1..8 {
+    for i in 0..2 {
-        (a[idx], carry) = a[idx].carrying_add(0, carry);
+        res[i] = u64::from_le_bytes(chunks[i]);
     }
-
-    return carry;
-}
-
-fn bytes_from_chars(chars_chunk: &[char]) -> [u8; 4] {
-    let mut res: [u8; 4] = [0; 4];
-    let mut idx: usize = 0;
-    chars_chunk.chunks_exact(2).for_each(|b_c| {
-        if idx < 4 {
-            match u8::from_str_radix(&b_c.iter().collect::<String>(), 16) {
-                Ok(n) => res[idx] = n,
-                Err(_) => (),
-            }
-            idx += 1;
-        }
-    });
     return res;
 }
 
-fn bignum_from_hex(hex: &str) -> [u32; 8] {
+pub fn u64arr_to_u128(a: [u64; 2]) -> u128 {
-    let mut res: [u32; 8] = [0; 8];
+    let mut bytes_data: [u8; 16] = [0u8; 16];
-    let mut idx: usize = 0;
+    bytes_data[0..8].copy_from_slice(a[0].to_le_bytes().as_slice());
-    let chars_hex = hex.chars().collect::<Vec<char>>();
+    bytes_data[8..].copy_from_slice(a[1].to_le_bytes().as_slice());
-    chars_hex
+
-        .chunks_exact(8)
+    return u128::from_le_bytes(bytes_data);
-        .rev()
-        .map(|chunk| bytes_from_chars(chunk))
-        .for_each(|b_arr| {
-            if idx < 8 {
-                res[idx] = u32::from_be_bytes(b_arr);
-                idx += 1;
-            }
-        });
-    return res;
 }
 
-fn hex_fmt_byte(n: u32) -> String {
+pub fn u128_to_u32arr(a: u128) -> [u32; 4] {
-    let res: String = n
+    let mut res = [0u32; 4];
-        .to_be_bytes()
+    let a_bytes = a.to_le_bytes();
-        .iter()
+    let chunks = a_bytes.as_chunks::<4>().0;
-        .map(|b| format!("{:02x}", b))
+    for i in 0..4 {
-        .collect();
+        res[i] = u32::from_le_bytes(chunks[i]);
-    return res;
-}
-
-fn bignum_to_hex(a: &[u32; 8]) -> String {
-    let res: String = a
-        .iter()
-        .rev()
-        .map(|n| hex_fmt_byte(*n))
-        .collect::<Vec<String>>()
-        .join("");
-    return res;
-}
-
-#[cfg(test)]
-mod num_utils_tests {
-    use std::io::Read;
-
-    use super::*;
-
-    #[test]
-    fn test_add() {
-        use std::io::{BufRead, BufReader};
-        use std::process::{Command, Stdio};
-
-        let test_gen_cmd = "/home/kira/Development/Rust/nyash-aes-xts256-plain64/nyash_client/src/tests/gen_test_data.py";
-        let mut child = Command::new(test_gen_cmd)
-            .stdout(Stdio::piped())
-            .spawn()
-            .unwrap();
-
-        let gen_stdout = child
-            .stdout
-            .take()
-            .ok_or("Failed to capture stdout")
-            .unwrap();
-        let gen_reader = BufReader::new(gen_stdout);
-
-        for r_line in gen_reader.lines() {
-            let test_line: String = r_line.unwrap(); // Handle any I/O errors
-            let test_data_line = test_line.split(' ').collect::<Vec<&str>>();
-            let num_to_add = u32::from_str_radix(test_data_line[0], 10).unwrap();
-            let t0 = bignum_from_hex(test_data_line[1]);
-
-            let t1_test = add_u32_to_u256(&t0, 1).0;
-            let t2_test = add_u32_to_u256(&t0, num_to_add).0;
-
-            let res_actual = format!(
-                "{} {} {} {}",
-                num_to_add,
-                bignum_to_hex(&t0),
-                bignum_to_hex(&t1_test),
-                bignum_to_hex(&t2_test)
-            );
-            assert_eq!(test_line, res_actual);
-        }
-
-        let _ = child.wait().unwrap();
-
-    }
-
-    #[test]
-    fn test_cl_add() {
-        extern crate ocl;
-        use ocl::{
-            Buffer, Context, Device, DeviceType, Kernel, Platform, Program, Queue, SpatialDims,
-            flags,
-        };
-        use std::fs::File;
-        use std::io::{BufRead, BufReader};
-        use std::process::{Command, Stdio};
-
-        let cl_test_path = "/home/kira/Development/Rust/nyash-aes-xts256-plain64/nyash_client/src/open_cl/test_num_utils.cl";
-        let cl_include_opt =
-            "-I /home/kira/Development/Rust/nyash-aes-xts256-plain64/nyash_client/src/open_cl";
-        let mut cl_src = String::new();
-        // read ocl source
-        BufReader::new(File::open(cl_test_path).unwrap()).read_to_string(&mut cl_src);
-
-        const G_WORK_SIZE: usize = 4096;
-        
-        let cl_platform = Platform::default();
-        let cl_device = Device::first(cl_platform).unwrap();
-        let cl_context = Context::builder()
-            .platform(cl_platform)
-            .devices(cl_device.clone())
-            .build()
-            .unwrap();
-        let cl_program = Program::builder()
-            .devices(cl_device)
-            .src(cl_src)
-            .cmplr_opt(cl_include_opt)
-            .build(&cl_context)
-            .unwrap();
-        let cl_queue = Queue::new(&cl_context, cl_device, None).unwrap();
-        
-        let cl_buffer_num = Buffer::<u32>::builder()
-            .queue(cl_queue.clone())
-            .flags(flags::MEM_READ_ONLY)
-            .len(G_WORK_SIZE)
-            .fill_val(0u32)
-            .build()
-            .unwrap();
-
-        let cl_buffer_t0 = Buffer::<u32>::builder()
-            .queue(cl_queue.clone())
-            .flags(flags::MEM_READ_ONLY)
-            .len(G_WORK_SIZE * 8)
-            .fill_val(0u32)
-            .build()
-            .unwrap();
-
-        let cl_buffer_t1 = Buffer::<u32>::builder()
-            .queue(cl_queue.clone())
-            .flags(flags::MEM_WRITE_ONLY)
-            .len(G_WORK_SIZE * 8)
-            .fill_val(0u32)
-            .build()
-            .unwrap();
-
-        let cl_buffer_t2 = Buffer::<u32>::builder()
-            .queue(cl_queue.clone())
-            .flags(flags::MEM_WRITE_ONLY)
-            .len(G_WORK_SIZE * 8)
-            .fill_val(0u32)
-            .build()
-            .unwrap();
-
-        // (3) Create a kernel with arguments matching those in the source above:
-        let kernel = Kernel::builder()
-            .program(&cl_program)
-            .name("test_add")
-            .queue(cl_queue.clone())
-            .global_work_size(G_WORK_SIZE)
-            .arg(&cl_buffer_num)
-            .arg(&cl_buffer_t0)
-            .arg(&cl_buffer_t1)
-            .arg(&cl_buffer_t2)
-            .build()
-            .unwrap();
-
-        let test_gen_cmd = "/home/kira/Development/Rust/nyash-aes-xts256-plain64/nyash_client/src/tests/gen_test_data.py";
-        let mut child = Command::new(test_gen_cmd)
-            .stdout(Stdio::piped())
-            .spawn()
-            .unwrap();
-
-        let gen_stdout = child
-            .stdout
-            .take()
-            .ok_or("Failed to capture stdout")
-            .unwrap();
-        let gen_reader = BufReader::new(gen_stdout);
-
-
-        let mut buffer_num: Vec<u32> = vec![0u32; G_WORK_SIZE];
-        let mut buffer_t0: Vec<u32> = vec![0u32; G_WORK_SIZE*8];
-        let mut exp_buffer_t1: Vec<u32> = vec![0u32; G_WORK_SIZE*8];
-        let mut exp_buffer_t2: Vec<u32> = vec![0u32; G_WORK_SIZE*8];
-
-        let mut act_buffer_t1: Vec<u32> = vec![0u32; G_WORK_SIZE*8];
-        let mut act_buffer_t2: Vec<u32> = vec![0u32; G_WORK_SIZE*8];
-
-        let mut w_id: usize = 0;
-
-        for r_line in gen_reader.lines() {
-            let test_line: String = r_line.unwrap(); // Handle any I/O errors
-            let test_data_line = test_line.split(' ').collect::<Vec<&str>>();
-            let num_to_add = u32::from_str_radix(test_data_line[0], 10).unwrap();
-            buffer_num[w_id] = num_to_add;
-            let slise_id = w_id*8;
-            buffer_t0[slise_id..slise_id+8].copy_from_slice(&bignum_from_hex(test_data_line[1]));
-            exp_buffer_t1[slise_id..slise_id+8].copy_from_slice(&bignum_from_hex(test_data_line[2]));
-            exp_buffer_t2[slise_id..slise_id+8].copy_from_slice(&bignum_from_hex(test_data_line[3]));
-
-            w_id += 1;
-            if w_id >= G_WORK_SIZE {
-                w_id = 0; // reset counter
-                cl_buffer_num.cmd().queue(&cl_queue).offset(0).write(&buffer_num).enq().unwrap();
-                cl_buffer_t0.cmd().queue(&cl_queue).offset(0).write(&buffer_t0).enq().unwrap();
-                
-                // (4) Run the kernel
-                unsafe {
-                    kernel
-                        .cmd()
-                        .queue(&cl_queue)
-                        .global_work_size(G_WORK_SIZE)
-                        .enq().unwrap();
-                }
-
-                cl_buffer_t1.cmd().queue(&cl_queue).offset(0).read(&mut act_buffer_t1).enq().unwrap();
-                cl_buffer_t2.cmd().queue(&cl_queue).offset(0).read(&mut act_buffer_t2).enq().unwrap();
-
-                assert_eq!(exp_buffer_t1, act_buffer_t1);
-                assert_eq!(exp_buffer_t2, act_buffer_t2);
-            }
-
-        }
-
-        let _ = child.wait().unwrap();
     }
+    return res;
 }
+
+pub fn u32arr_to_u128(a: [u32; 4]) -> u128 {
+    let mut bytes_data: [u8; 16] = [0u8; 16];
+    bytes_data[0..4].copy_from_slice(a[0].to_le_bytes().as_slice());
+    bytes_data[4..8].copy_from_slice(a[1].to_le_bytes().as_slice());
+    bytes_data[8..12].copy_from_slice(a[2].to_le_bytes().as_slice());
+    bytes_data[12..16].copy_from_slice(a[3].to_le_bytes().as_slice());
+
+    return u128::from_le_bytes(bytes_data);
+}
+
+// fn add_u32_to_u256(a: &[u32; 8], b: u32) -> ([u32; 8], bool) {
+//     let mut res: [u32; 8] = [0; 8];
+//     let mut carry = false;
+//     (res[0], carry) = a[0].carrying_add(b, carry);
+
+//     for idx in 1..8 {
+//         (res[idx], carry) = a[idx].carrying_add(0, carry);
+//     }
+
+//     return (res, carry);
+// }
+
+// fn add_u32_to_u256_(a: &mut [u32; 8], b: u32) -> bool {
+//     let mut carry = false;
+//     (a[0], carry) = a[0].carrying_add(b, carry);
+
+//     for idx in 1..8 {
+//         (a[idx], carry) = a[idx].carrying_add(0, carry);
+//     }
+
+//     return carry;
+// }
+
+// fn bytes_from_chars(chars_chunk: &[char]) -> [u8; 4] {
+//     let mut res: [u8; 4] = [0; 4];
+//     let mut idx: usize = 0;
+//     chars_chunk.chunks_exact(2).for_each(|b_c| {
+//         if idx < 4 {
+//             match u8::from_str_radix(&b_c.iter().collect::<String>(), 16) {
+//                 Ok(n) => res[idx] = n,
+//                 Err(_) => (),
+//             }
+//             idx += 1;
+//         }
+//     });
+//     return res;
+// }
+
+// fn bignum_from_hex(hex: &str) -> [u32; 8] {
+//     let mut res: [u32; 8] = [0; 8];
+//     let mut idx: usize = 0;
+//     let chars_hex = hex.chars().collect::<Vec<char>>();
+//     chars_hex
+//         .chunks_exact(8)
+//         .rev()
+//         .map(|chunk| bytes_from_chars(chunk))
+//         .for_each(|b_arr| {
+//             if idx < 8 {
+//                 res[idx] = u32::from_be_bytes(b_arr);
+//                 idx += 1;
+//             }
+//         });
+//     return res;
+// }
+
+// fn hex_fmt_byte(n: u32) -> String {
+//     let res: String = n
+//         .to_be_bytes()
+//         .iter()
+//         .map(|b| format!("{:02x}", b))
+//         .collect();
+//     return res;
+// }
+
+// fn bignum_to_hex(a: &[u32; 8]) -> String {
+//     let res: String = a
+//         .iter()
+//         .rev()
+//         .map(|n| hex_fmt_byte(*n))
+//         .collect::<Vec<String>>()
+//         .join("");
+//     return res;
+// }
+
+// #[cfg(test)]
+// mod num_utils_tests {
+//     use std::io::Read;
+
+//     use super::*;
+
+//     #[test]
+//     fn test_add() {
+//         use std::io::{BufRead, BufReader};
+//         use std::process::{Command, Stdio};
+
+//         let test_gen_cmd = "/home/kira/Development/Rust/nyash-aes-xts256-plain64/nyash_client/src/tests/gen_test_data.py";
+//         let mut child = Command::new(test_gen_cmd)
+//             .stdout(Stdio::piped())
+//             .spawn()
+//             .unwrap();
+
+//         let gen_stdout = child
+//             .stdout
+//             .take()
+//             .ok_or("Failed to capture stdout")
+//             .unwrap();
+//         let gen_reader = BufReader::new(gen_stdout);
+
+//         for r_line in gen_reader.lines() {
+//             let test_line: String = r_line.unwrap(); // Handle any I/O errors
+//             let test_data_line = test_line.split(' ').collect::<Vec<&str>>();
+//             let num_to_add = u32::from_str_radix(test_data_line[0], 10).unwrap();
+//             let t0 = bignum_from_hex(test_data_line[1]);
+
+//             let t1_test = add_u32_to_u256(&t0, 1).0;
+//             let t2_test = add_u32_to_u256(&t0, num_to_add).0;
+
+//             let res_actual = format!(
+//                 "{} {} {} {}",
+//                 num_to_add,
+//                 bignum_to_hex(&t0),
+//                 bignum_to_hex(&t1_test),
+//                 bignum_to_hex(&t2_test)
+//             );
+//             assert_eq!(test_line, res_actual);
+//         }
+
+//         let _ = child.wait().unwrap();
+//     }
+
+//     #[test]
+//     fn test_cl_add() {
+//         extern crate ocl;
+//         use ocl::{Buffer, Context, Device, Kernel, Platform, Program, Queue, flags};
+//         use std::fs::File;
+//         use std::io::{BufRead, BufReader};
+//         use std::process::{Command, Stdio};
+
+//         let cl_test_path = "/home/kira/Development/Rust/nyash-aes-xts256-plain64/nyash_client/src/open_cl/test_num_utils.cl";
+//         let cl_include_opt =
+//             "-I /home/kira/Development/Rust/nyash-aes-xts256-plain64/nyash_client/src/open_cl";
+//         let mut cl_src = String::new();
+//         // read ocl source
+//         BufReader::new(File::open(cl_test_path).unwrap())
+//             .read_to_string(&mut cl_src)
+//             .expect("Error reading cl_src!");
+
+//         const G_WORK_SIZE: usize = 4096;
+
+//         let cl_platform = Platform::default();
+//         let cl_device = Device::first(cl_platform).unwrap();
+//         let cl_context = Context::builder()
+//             .platform(cl_platform)
+//             .devices(cl_device.clone())
+//             .build()
+//             .unwrap();
+//         let cl_program = Program::builder()
+//             .devices(cl_device)
+//             .src(cl_src)
+//             .cmplr_opt(cl_include_opt)
+//             .build(&cl_context)
+//             .unwrap();
+//         let cl_queue = Queue::new(&cl_context, cl_device, None).unwrap();
+
+//         let cl_buffer_num = Buffer::<u32>::builder()
+//             .queue(cl_queue.clone())
+//             .flags(flags::MEM_READ_ONLY)
+//             .len(G_WORK_SIZE)
+//             .fill_val(0u32)
+//             .build()
+//             .unwrap();
+
+//         let cl_buffer_t0 = Buffer::<u32>::builder()
+//             .queue(cl_queue.clone())
+//             .flags(flags::MEM_READ_ONLY)
+//             .len(G_WORK_SIZE * 8)
+//             .fill_val(0u32)
+//             .build()
+//             .unwrap();
+
+//         let cl_buffer_t1 = Buffer::<u32>::builder()
+//             .queue(cl_queue.clone())
+//             .flags(flags::MEM_WRITE_ONLY)
+//             .len(G_WORK_SIZE * 8)
+//             .fill_val(0u32)
+//             .build()
+//             .unwrap();
+
+//         let cl_buffer_t2 = Buffer::<u32>::builder()
+//             .queue(cl_queue.clone())
+//             .flags(flags::MEM_WRITE_ONLY)
+//             .len(G_WORK_SIZE * 8)
+//             .fill_val(0u32)
+//             .build()
+//             .unwrap();
+
+//         // (3) Create a kernel with arguments matching those in the source above:
+//         let kernel = Kernel::builder()
+//             .program(&cl_program)
+//             .name("test_add")
+//             .queue(cl_queue.clone())
+//             .global_work_size(G_WORK_SIZE)
+//             .arg(&cl_buffer_num)
+//             .arg(&cl_buffer_t0)
+//             .arg(&cl_buffer_t1)
+//             .arg(&cl_buffer_t2)
+//             .build()
+//             .unwrap();
+
+//         let test_gen_cmd = "/home/kira/Development/Rust/nyash-aes-xts256-plain64/nyash_client/src/tests/gen_test_data.py";
+//         let mut child = Command::new(test_gen_cmd)
+//             .stdout(Stdio::piped())
+//             .spawn()
+//             .unwrap();
+
+//         let gen_stdout = child
+//             .stdout
+//             .take()
+//             .ok_or("Failed to capture stdout")
+//             .unwrap();
+//         let gen_reader = BufReader::new(gen_stdout);
+
+//         let mut buffer_num: Vec<u32> = vec![0u32; G_WORK_SIZE];
+//         let mut buffer_t0: Vec<u32> = vec![0u32; G_WORK_SIZE * 8];
+//         let mut exp_buffer_t1: Vec<u32> = vec![0u32; G_WORK_SIZE * 8];
+//         let mut exp_buffer_t2: Vec<u32> = vec![0u32; G_WORK_SIZE * 8];
+
+//         let mut act_buffer_t1: Vec<u32> = vec![0u32; G_WORK_SIZE * 8];
+//         let mut act_buffer_t2: Vec<u32> = vec![0u32; G_WORK_SIZE * 8];
+
+//         let mut w_id: usize = 0;
+
+//         for r_line in gen_reader.lines() {
+//             let test_line: String = r_line.unwrap(); // Handle any I/O errors
+//             let test_data_line = test_line.split(' ').collect::<Vec<&str>>();
+//             let num_to_add = u32::from_str_radix(test_data_line[0], 10).unwrap();
+//             buffer_num[w_id] = num_to_add;
+//             let slise_id = w_id * 8;
+//             buffer_t0[slise_id..slise_id + 8].copy_from_slice(&bignum_from_hex(test_data_line[1]));
+//             exp_buffer_t1[slise_id..slise_id + 8]
+//                 .copy_from_slice(&bignum_from_hex(test_data_line[2]));
+//             exp_buffer_t2[slise_id..slise_id + 8]
+//                 .copy_from_slice(&bignum_from_hex(test_data_line[3]));
+
+//             w_id += 1;
+//             if w_id >= G_WORK_SIZE {
+//                 w_id = 0; // reset counter
+//                 cl_buffer_num
+//                     .cmd()
+//                     .queue(&cl_queue)
+//                     .offset(0)
+//                     .write(&buffer_num)
+//                     .enq()
+//                     .unwrap();
+//                 cl_buffer_t0
+//                     .cmd()
+//                     .queue(&cl_queue)
+//                     .offset(0)
+//                     .write(&buffer_t0)
+//                     .enq()
+//                     .unwrap();
+
+//                 // (4) Run the kernel
+//                 unsafe {
+//                     kernel
+//                         .cmd()
+//                         .queue(&cl_queue)
+//                         .global_work_size(G_WORK_SIZE)
+//                         .enq()
+//                         .unwrap();
+//                 }
+
+//                 cl_buffer_t1
+//                     .cmd()
+//                     .queue(&cl_queue)
+//                     .offset(0)
+//                     .read(&mut act_buffer_t1)
+//                     .enq()
+//                     .unwrap();
+//                 cl_buffer_t2
+//                     .cmd()
+//                     .queue(&cl_queue)
+//                     .offset(0)
+//                     .read(&mut act_buffer_t2)
+//                     .enq()
+//                     .unwrap();
+
+//                 assert_eq!(exp_buffer_t1, act_buffer_t1);
+//                 assert_eq!(exp_buffer_t2, act_buffer_t2);
+//             }
+//         }
+
+//         let _ = child.wait().unwrap();
+//     }
+// }

nyash_client/src/ocl_utils.rs

@@ -0,0 +1,320 @@
+
+use std::io::Read;
+
+use ocl::{Buffer, Context, Device, Kernel, Platform, Program, Queue, flags};
+
+use crate::num_utils;
+
+pub struct CtxBuffers {
+    tweak_params: Buffer<u32>,
+    batch_size: Buffer<u64>,
+    start_key: Buffer<u32>,
+    tweak_key: Buffer<u32>,
+    uenc_data: Buffer<u32>,
+    target_data: Buffer<u32>,
+    key_found: Buffer<u32>,
+}
+
+pub struct ExecData {
+    pub start_key: Vec<u32>,
+    pub tweak_key: Vec<u32>,
+    pub uenc_data: Vec<u32>,
+    pub target_data: Vec<u32>,
+    pub tweak_i: u64,
+    pub tweak_j: u32,
+    pub key_found: Vec<u32>,
+    pub batch_size: u64,
+    pub work_size: usize,
+}
+ impl ExecData {
+    // g_params[uint4]
+    // g_params[0-1] - ulog g_Ti
+    // g_params[2] -   g_Tj
+    pub fn tweak_params(&self) -> Vec<u32> {
+        let mut res: Vec<u32> = Vec::with_capacity(4);
+
+        // the sector number (S) is first converted into a little-endian byte array 
+        // before being encrypted using the second AES key (K₂)
+        let tweak_i_b: [u8;8] = self.tweak_i.to_le_bytes();
+        let (tweak_i_cnk,_) = tweak_i_b.as_chunks::<4>();
+        res.push(u32::from_le_bytes(tweak_i_cnk[0]));
+        res.push(u32::from_le_bytes(tweak_i_cnk[1]));
+
+        //last enc block number (tweak_j)
+        res.push(self.tweak_j);
+
+        return res;
+    }
+
+    pub fn get_found_key(&self) -> u128 {
+        let mut u32_arr_k = [0u32;4];
+        u32_arr_k[0] = self.key_found[1];
+        u32_arr_k[1] = self.key_found[2];
+        u32_arr_k[2] = self.key_found[3];
+        u32_arr_k[3] = self.key_found[4];
+        return num_utils::u32arr_to_u128(u32_arr_k);
+    }
+     
+ }
+
+pub struct ExecContext {
+    _ctx: Context,
+    kernel: Kernel,
+    _prog: Program,
+    queue: Queue,
+    buffers: CtxBuffers,
+}
+
+pub fn init_program(
+    cl_device: Device,
+    cl_platform: Platform,
+    cl_src_gz_bytes: &[u8],
+) -> Result<(Context, Program, Queue), ocl::Error> {
+    use flate2::read::GzDecoder;
+    let mut gz_decoder = GzDecoder::new(cl_src_gz_bytes);
+    let mut decompressed_src = String::new();
+    gz_decoder.read_to_string(&mut decompressed_src).expect("Error decompressing OCL sources!");
+
+    let cl_context = Context::builder()
+        .platform(cl_platform)
+        .devices(cl_device.clone())
+        .build()?;
+
+    let cl_program = Program::builder()
+        .devices(cl_device)
+        .src(decompressed_src)
+        .build(&cl_context)
+        .unwrap();
+
+    let cl_queue: Queue = Queue::new(&cl_context, cl_device, None)?;
+
+    return Ok((cl_context, cl_program, cl_queue));
+}
+
+pub fn init_buffers(cl_queue: &Queue) -> Result<CtxBuffers, ocl::Error> {
+    let cl_buffer_tweak_params = Buffer::<u32>::builder()
+        .queue(cl_queue.clone())
+        .flags(flags::MEM_READ_ONLY)
+        .len(3)
+        .fill_val(0u32)
+        .build()?;
+
+    let cl_buffer_batch_size = Buffer::<u64>::builder()
+        .queue(cl_queue.clone())
+        .flags(flags::MEM_READ_ONLY)
+        .len(1)
+        .fill_val(0u64)
+        .build()?;
+
+    let cl_buffer_start_key = Buffer::<u32>::builder()
+        .queue(cl_queue.clone())
+        .flags(flags::MEM_READ_ONLY)
+        .len(4)
+        .fill_val(0u32)
+        .build()?;
+
+    let cl_buffer_tweak_key = Buffer::<u32>::builder()
+        .queue(cl_queue.clone())
+        .flags(flags::MEM_READ_ONLY)
+        .len(4)
+        .fill_val(0u32)
+        .build()?;
+
+    let cl_buffer_uenc_data = Buffer::<u32>::builder()
+        .queue(cl_queue.clone())
+        .flags(flags::MEM_READ_ONLY)
+        .len(4)
+        .fill_val(0u32)
+        .build()?;
+
+    let cl_buffer_target_data = Buffer::<u32>::builder()
+        .queue(cl_queue.clone())
+        .flags(flags::MEM_READ_ONLY)
+        .len(4)
+        .fill_val(0u32)
+        .build()?;
+
+    let cl_buffer_key_found = Buffer::<u32>::builder()
+        .queue(cl_queue.clone())
+        .flags(flags::MEM_WRITE_ONLY)
+        .len(5)
+        .fill_val(0u32)
+        .build()?;
+
+
+    Ok(CtxBuffers {
+        tweak_params: cl_buffer_tweak_params,
+        batch_size: cl_buffer_batch_size,
+        start_key: cl_buffer_start_key,
+        tweak_key: cl_buffer_tweak_key,
+        uenc_data: cl_buffer_uenc_data,
+        target_data: cl_buffer_target_data,
+        key_found: cl_buffer_key_found,
+    })
+}
+
+fn init_kernel(
+    work_size: usize,
+    cl_program: &Program,
+    cl_queue: &Queue,
+    buffs: &CtxBuffers,
+) -> Result<Kernel, ocl::Error> {
+    Kernel::builder()
+        .program(cl_program)
+        .name("search_key")
+        .queue(cl_queue.clone())
+        .global_work_size(work_size)
+        .arg(&buffs.tweak_params)
+        .arg(&buffs.batch_size)
+        .arg(&buffs.start_key)
+        .arg(&buffs.tweak_key)
+        .arg(&buffs.uenc_data)
+        .arg(&buffs.target_data)
+        .arg(&buffs.key_found)
+        .build()
+}
+
+impl ExecContext {
+    // Constructor with parameters
+    pub fn new(
+        cl_device: Device,
+        cl_platform: Platform,
+        cl_src_gz_bytes: &[u8],
+        global_work_size: usize,
+    ) -> Result<Self, ocl::Error> {
+        let (nya_cl_context, nya_cl_program, nya_cl_queue) =
+            init_program(cl_device, cl_platform, cl_src_gz_bytes)?;
+
+        let nya_cl_buffers = init_buffers(&nya_cl_queue)?;
+
+        let nya_cl_kernel = init_kernel(
+            global_work_size,
+            &nya_cl_program,
+            &nya_cl_queue,
+            &nya_cl_buffers,
+        )?;
+        Ok(Self {
+            _ctx: nya_cl_context,
+            kernel: nya_cl_kernel,
+            _prog: nya_cl_program,
+            queue: nya_cl_queue,
+            buffers: nya_cl_buffers,
+        })
+    }
+}
+
+pub fn set_target_data(ex_ctx: &mut ExecContext, ex_data: &mut ExecData) -> Result<(), ocl::Error> {
+    // transfer tweaks
+    let t_p = ex_data.tweak_params();
+    ex_ctx
+        .buffers
+        .tweak_params
+        .cmd()
+        .queue(&ex_ctx.queue)
+        .offset(0)
+        .write(&t_p)
+        .enq()?;
+
+    // transfen unencrypted data to device
+    ex_ctx
+        .buffers
+        .uenc_data
+        .cmd()
+        .queue(&ex_ctx.queue)
+        .offset(0)
+        .write(&ex_data.uenc_data)
+        .enq()?;
+
+    // transfet target data
+    ex_ctx
+        .buffers
+        .target_data
+        .cmd()
+        .queue(&ex_ctx.queue)
+        .offset(0)
+        .write(&ex_data.target_data)
+        .enq()?;
+
+    ex_ctx.queue.finish()?;
+
+    return Ok(());
+}
+
+pub fn do_work(ex_ctx: &mut ExecContext, ex_data: &mut ExecData) -> Result<(bool, f64), ocl::Error> {
+
+    let b_s = vec![ex_data.batch_size];
+
+    let start_time = std::time::Instant::now();
+
+    // tranfer batch_size
+    ex_ctx
+        .buffers
+        .batch_size
+        .cmd()
+        .queue(&ex_ctx.queue)
+        .offset(0)
+        .write(&b_s)
+        .enq()?;
+
+
+    // transfer start key to device
+    ex_ctx
+        .buffers
+        .start_key
+        .cmd()
+        .queue(&ex_ctx.queue)
+        .offset(0)
+        .write(&ex_data.start_key)
+        .enq()?;
+
+    // transfet tweak key
+    ex_ctx
+        .buffers
+        .tweak_key
+        .cmd()
+        .queue(&ex_ctx.queue)
+        .offset(0)
+        .write(&ex_data.tweak_key)
+        .enq()?;
+
+    // zero out key_found buffer
+    // ex_ctx
+    //     .buffers
+    //     .key_found
+    //     .cmd()
+    //     .queue(&ex_ctx.queue)
+    //     .offset(0)
+    //     .fill(0u32, None)
+    //     .enq()?;
+
+    
+    // (4) Run the kernel
+    unsafe {
+        ex_ctx
+            .kernel
+            .cmd()
+            .queue(&ex_ctx.queue)
+            .global_work_size(ex_data.work_size)
+            .enq()?;
+    }
+
+    // read key_foun buffer
+    ex_ctx
+        .buffers
+        .key_found
+        .cmd()
+        .queue(&ex_ctx.queue)
+        .offset(0)
+        .read(&mut ex_data.key_found)
+        .enq()?;
+
+    let exec_duration = start_time.elapsed().as_secs_f64();
+
+    //ex_ctx.queue.finish()?;
+
+    if ex_data.key_found[0] == 0 {
+        Ok((false, exec_duration))
+    } else {
+        Ok((true, exec_duration))
+    }
+}

nyash_client/src/old_main.rs

@@ -12,32 +12,7 @@ use crate::client_config::{AppConfig, DevConf};
 
 mod client_config;
 mod num_utils;
-/// Exploded version. Boom!
-///
-/// The functions above use `ProQue` and other abstractions to greatly reduce
-/// the amount of boilerplate and configuration necessary to do basic work.
-/// Many tasks, however, will require more configuration and will necessitate
-/// doing away with `ProQue` altogether. Enqueuing kernels and reading/writing
-/// from buffers and images usually requires a more explicit interface.
-///
-/// The following function performs the exact same steps that the above
-/// functions did, with many of the convenience abstractions peeled away.
-///
-/// See the function below this to take things a step deeper...
-///
-// trait FromStr {
-//     fn from_str(&self);
-// }
-// // Define a trait with a constructor method
-// trait NewFile {
-//     fn new<P: AsRef<Path>>(path: P) -> std::io::Result<Self> where Self: Sized;
-// }
 
-// impl Foo for ocl::flags::DeviceType {
-//     fn foo(&self) {
-//         println!("foo");
-//     }
-// }
 
 fn dev_type_from_str(s: &str) -> Result<flags::DeviceType, ()> {
     match s {

nyash_client/src/open_cl/build_spirv.sh

@@ -0,0 +1,8 @@
+#!/bin/sh
+
+out_file="$1"
+
+clang -c -target spir64 -O0 -finclude-default-header -I ./ -emit-llvm -o nyash_aes_xts256_plain.bc nyash_aes_xts256_plain.cl
+
+#llc -march=spir64 nyash_aes_xts256_plain.bc -filetype=obj -o $out_file
+llvm-spirv -o $out_file nyash_aes_xts256_plain.bc

nyash_client/src/open_cl/nyash_aes_xts256_plain.cl

@@ -32,15 +32,25 @@
 // g_key_found uint[9] - 0 element - flag that sets to 1 if key found.
 // Other 8 elements is found key
 
-__kernel void search_key(const uint batch_size, const ulong g_Ti, const uint g_Tj, 
+
-                              __global const uint8* g_start_enc_key,
+// in current implementation tweak key is not changing by kernel
+// So changing it a bit
+
+// g_params[uint4]
+// g_params[0] -   batch_size
+// g_params[1-2] - ulong g_Ti
+// g_params[3] -   g_Tj
+__kernel void search_key(__global const uint* g_tweak_params, 
+                              __global const ulong* g_batch_size,
+                              __global const uint4* g_start_enc_key,
+                              __global const uint4* g_tweak_key,
                               __global const uint4* g_uenc_data,
                               __global const uint4* g_target_data,
                               __global uint* g_key_found)
 {
   const uint g_id = get_global_id(0);
 
-  uint enc_key[8];
+  uint enc_key[4];
   uint tweak[4];
   uint uenc_data[4];
   uint4 target_data = *g_target_data;
@@ -48,32 +58,34 @@ __kernel void search_key(const uint batch_size, const ulong g_Ti, const uint g_T
   uint d_ks[44]; // data expanded key
   uint t_ks[44]; // tweak expanded key
 
+  //set batch_size
+  ulong batch_size = g_batch_size[0];
+
   // set disk sector number
   uint sec_n[4] = {0};
-  sec_n[0] = ((uint*)&g_Ti)[0];
+  sec_n[0] = g_tweak_params[0];
-  sec_n[1] = ((uint*)&g_Ti)[1];
+  sec_n[1] = g_tweak_params[1];
 
-  uint Tj = g_Tj; // AES block number
+  uint Tj = g_tweak_params[2]; // AES block number
 
   vstore4(*g_uenc_data, 0, uenc_data);
-  vstore8(*g_start_enc_key, 0, enc_key);
+  vstore4(*g_start_enc_key, 0, enc_key);
-
+  vstore4(*g_tweak_key, 0, tweak);
 
   // Set initial start key for every work thread
   uint k_data_carry = add_uint_to_bigint4_ (enc_key, (g_id*batch_size));
-  // uint k_tweak_carry = add_uint_to_bigint4_ (&enc_key[4], k_data_carry);
+  if (k_data_carry != 0u) return; // if reached max key value exit thread
-  // No need to store tweak carry
-  if (add_uint_to_bigint4_ (&enc_key[4], k_data_carry) != 0u) return; // if reached max key value exit thread
 
   // Generate tweak
-  aes128_set_encrypt_key (t_ks, &enc_key[4]);
+  aes128_set_encrypt_key (t_ks, tweak);
   aes_xts256_gen_tweak (t_ks, sec_n, Tj, tweak);
+  //if (g_id == 0) g_key_found[1] = 1;
 
-  for (uint batch_id = 0u; (batch_id < batch_size); batch_id++)
+  for (ulong batch_id = 0ul; batch_id < batch_size; batch_id++)
   {
-    // Data encrypt key always changing because we increment from 0 index to 8
+    //if (g_id == 0) g_key_found[1] = 2;
+    // Set encrypt key
     aes128_set_encrypt_key (d_ks, enc_key);
-
     // encrypt data
     aes_xts256_enc_block (d_ks, tweak, uenc_data, (uint*)&enc_data);
 
@@ -81,29 +93,17 @@ __kernel void search_key(const uint batch_size, const ulong g_Ti, const uint g_T
     if (all(enc_data==target_data)) 
     {
       g_key_found[0] = 1;
-      vstore8(vload8(0, enc_key), 0, &g_key_found[1]);
+      g_key_found[1] = enc_key[0];
+      g_key_found[2] = enc_key[1];
+      g_key_found[3] = enc_key[2];
+      g_key_found[4] = enc_key[3];
       return;
     }
 
-    // Increment data key part by 1.
+    // Increment data key by 1.
     k_data_carry = add_one_to_bigint4_ (enc_key);
-
+    // if reached max key value exit thread
-    // Tweak changes only once in 2^128 times
+    if (k_data_carry != 0u) return; 
-    if (k_data_carry != 0u) {
-
-      // Increment tweak part
-      // k_tweak_carry = add_one_to_bigint4_ (&enc_key[4]);
-      add_one_to_bigint4_ (&enc_key[4]); // no need to store tweak carry
-
-      // *** I commented next line because, its a really!! lol! rare event,
-      // *** and in a worse case we just do a bit of noneed work.
-      // *** But additional check on every itaration actually mesurable cost.
-      // if (k_tweak_carry != 0u) return; // if reached max key value exit thread
-
-      // Gen new tweak
-      aes128_set_encrypt_key (t_ks, &enc_key[4]);
-      aes_xts256_gen_tweak (t_ks, sec_n, Tj, tweak);
-    }
   }
 }
 

nyash_client/src/search_params.rs

@@ -0,0 +1,38 @@
+
+
+pub fn get_params() ->([u32; 4],[u32; 4],[u32; 4]) {
+    use crate::num_utils;
+
+    const ENCRYPTED_DATA: [u8; 16] = [
+            198, 255, 55, 185, 15, 226, 223, 174, 119, 8, 36, 239, 242, 89, 126, 230
+    ];
+
+    const KEY_DATA: [u8; 32] = [
+        206, 193, 83, 54, 46, 234, 185, 41, 146, 244, 130, 6, 212, 68, 106, 162, 165, 97, 188,
+        218, 39, 111, 141, 236, 67, 159, 157, 157, 166, 79, 89, 134
+    ];
+
+    // let key_bytes_reversed: Vec<u8> = KEY_DATA.iter().rev().map(|e| *e).collect();
+    // let data_bytes_reversed: Vec<u8> = ENCRYPTED_DATA.iter().rev().map(|e| *e).collect();
+
+    let mut tweak_key_b: [u8;16] = [0u8;16];
+    let mut data_key_b: [u8;16] = [0u8;16];
+
+    data_key_b.copy_from_slice(&KEY_DATA[0..16]);
+    tweak_key_b.copy_from_slice(&KEY_DATA[16..32]);
+
+    // getting keys
+    let data_key = u128::from_le_bytes(data_key_b);
+    let tweak_key = u128::from_le_bytes(tweak_key_b);
+    let data_key = num_utils::u128_to_u32arr(data_key);
+    let tweak_key = num_utils::u128_to_u32arr(tweak_key);
+
+    // converting bytes raw data to u32 arr
+    let mut encrypted_data: [u32; 4] = [0u32; 4];
+    let (enc_dat_bytes_chunks, _) = ENCRYPTED_DATA.as_chunks::<4>();
+    for i in 0..4 {
+        encrypted_data[i] = u32::from_le_bytes(enc_dat_bytes_chunks[i]);
+    }
+
+    return (data_key, tweak_key, encrypted_data);
+}

nyash_client/src/test_cl.rs

@@ -0,0 +1,79 @@
+
+#[cfg(test)]
+mod test_cl {
+    #[test]
+    fn test_encryption() {
+        use ocl::{Device, Platform};
+        use crate::ocl_utils;
+        use crate::num_utils;
+        
+
+        const SRC_PATH: &str = "src/open_cl/nyash_aes_xts256_plain.cl";
+        const OCL_COMP_OPT: &str = "-I src/open_cl";
+        const ENCRYPTED_DATA: [u8; 16] = [
+            198, 255, 55, 185, 15, 226, 223, 174, 119, 8, 36, 239, 242, 89, 126, 230
+        ];
+        const KEY_DATA: [u8; 32] = [
+            206, 193, 83, 54, 46, 234, 185, 41, 146, 244, 130, 6, 212, 68, 106, 162, 165, 97, 188,
+            218, 39, 111, 141, 236, 67, 159, 157, 157, 166, 79, 89, 134
+        ];
+
+        // let key_bytes_reversed: Vec<u8> = KEY_DATA.iter().rev().map(|e| *e).collect();
+        // let data_bytes_reversed: Vec<u8> = ENCRYPTED_DATA.iter().rev().map(|e| *e).collect();
+
+        let mut tweak_key_b: [u8;16] = [0u8;16];
+        let mut data_key_b: [u8;16] = [0u8;16];
+
+        data_key_b.copy_from_slice(&KEY_DATA[0..16]);
+        tweak_key_b.copy_from_slice(&KEY_DATA[16..32]);
+
+        // getting keys
+        let data_key = u128::from_le_bytes(data_key_b);
+        let tweak_key = u128::from_le_bytes(tweak_key_b);
+        let data_key = num_utils::u128_to_u32arr(data_key);
+        let tweak_key = num_utils::u128_to_u32arr(tweak_key);
+
+
+        // converting bytes raw data to u32 arr
+        let mut encrypted_data: [u32; 4] = [0u32; 4];
+        let (enc_dat_bytes_chunks, _) = ENCRYPTED_DATA.as_chunks::<4>();
+        for i in 0..4 {
+            encrypted_data[i] = u32::from_le_bytes(enc_dat_bytes_chunks[i]);
+        }
+
+        // init devices
+        let platform = Platform::first().expect("Error getting platform!");
+        let device = Device::first(platform).expect("Error getting device!");
+
+        println!("Platform: {:?}, Device: {:?}", platform.name().unwrap(), device.name().unwrap());
+
+        // reading ocl program sources
+        let prog_src = std::fs::read_to_string(SRC_PATH).expect("Error reading program sources!");
+
+        let mut nyan_context =
+            ocl_utils::ExecContext::new(device, platform, prog_src.as_str(), OCL_COMP_OPT, 256)
+                .expect("Error creating execution nyan context!");
+
+
+        //setting data
+        let mut nyan_exec_dat = ocl_utils::ExecData {
+            start_key: data_key.to_vec(),
+            tweak_key: tweak_key.to_vec(),
+            uenc_data: vec![0u32;4],
+            target_data: encrypted_data.to_vec(),
+            tweak_i: 0,
+            tweak_j: 0,
+            key_found: vec![0u32;5],
+            batch_size: 1000000,
+            work_size: 256,
+        };
+
+        println!("Set target data");
+        ocl_utils::set_target_data(&mut nyan_context, &mut nyan_exec_dat).expect("Error set target data!");
+        
+        let found_flag = ocl_utils::do_work(&mut nyan_context, &mut nyan_exec_dat).expect("Error do work!");
+        println!("Found?: {}", found_flag);
+        println!("Key found: {:?}", nyan_exec_dat.key_found);
+        assert_eq!(true, found_flag);
+    }
+}

nyash_client/src/test_data/create_test.sh

@@ -0,0 +1,16 @@
+#!/bin/bash
+
+truncate -s 255M test_btrfs.img
+dd if=/dev/urandom of=master.key bs=32 count=1
+
+cryptsetup luksFormat --type=luks2 --sector-size 512 --pbkdf=pbkdf2 --pbkdf-force-iterations=1000 --hash=sha256 --key-size=256 --cipher=aes-xts-plain64 --master-key-file ./test_master.key ./test.img
+
+cryptsetup luksFormat --type=luks2 --pbkdf=pbkdf2 --pbkdf-force-iterations=1000 --hash=sha256 --key-size=256 --cipher=aes-xts-plain64 ./luks-container.img
+
+sudo cryptsetup luksOpen ./test.img luks-container-crypt
+
+sudo mkfs.btrfs /dev/mapper/luks-container-crypt
+
+sudo dd if=/dev/mapper/luks-container-crypt of=./test_btrfs_luks_unencrypt.img bs=1M count=255
+
+sudo cryptsetup close luks-container-crypt

nyash_client/src/test_data/extract_encrypted_block.py

@@ -0,0 +1,44 @@
+from utils import read_metadata
+
+
+# Init logger
+
+LUKS_FILE_NAME = "test.img"
+KEY_FILE_NAME = "master.key"
+
+
+
+
+def main():
+
+    metadat = read_metadata(LUKS_FILE_NAME)
+    print(f"metadata:\n{metadat}")
+
+
+    segments_offset_bytes = int(metadat["segments"]["0"]["offset"])
+    superblock_start_bytes = 0x00010000
+    superblock_start_sector = superblock_start_bytes//512
+    magic_offset = 0x40
+    superblock_lenght_bytes = 0x1000
+    sector_size = 512
+
+    with open(LUKS_FILE_NAME, 'rb') as luks_file:
+        luks_file.seek(segments_offset_bytes)
+        enc_data = luks_file.read(16)
+
+
+    print("ENC DATA:")
+    print("[" + ",".join([format(a, 'd') for a in enc_data])+"]")
+
+    print()
+    print("KEY DATA:")
+    with open(KEY_FILE_NAME, 'rb') as key_file:
+        key_data = key_file.read(32)
+        print("[" + ",".join([format(a, 'd') for a in key_data])+"]")
+
+
+
+
+
+if __name__ == '__main__':
+    main()

nyash_client/src/test_data/utils.py

@@ -0,0 +1,25 @@
+import subprocess
+import json
+
+def read_metadata(file_name: str) -> dict:
+    #cryptsetup luksDump --dump-json-metadata /dev/loop0
+    luks_cmd: list[str] = ["cryptsetup", "luksDump", "--dump-json-metadata", file_name]
+
+    result = subprocess.run(luks_cmd, capture_output=True, encoding="UTF-8")
+    if result.returncode == 0 and result.stdout is not None:
+        metadata = json.loads(result.stdout)
+        return metadata
+    else:
+        raise Exception(f"Error executing 'cryptsetup' binary! {result.stderr}")
+
+
+def read_encrypted_key(f_name: str, metadata: dict, keyslot: int) -> bytes:
+    stripes = metadata["keyslots"][str(keyslot)]["af"]["stripes"]
+    offset = int(metadata["keyslots"][str(keyslot)]["area"]["offset"])
+    # size = int(metadata["keyslots"][str(keyslot)]["area"]["size"])
+    key_size = metadata["keyslots"][str(keyslot)]["area"]["key_size"]
+    with open(f_name, 'rb') as luks_file:
+        luks_file.seek(offset)
+        data = luks_file.read(key_size*stripes)
+    return data
+

nyash_client/src/test_nums.rs

@@ -1,9 +1,6 @@
 
 
 
-
-
-
 #[cfg(test)]
 mod cl_num_utils_tests {
     use super::*;

nyash_server/Cargo.lock

@@ -581,15 +581,6 @@ version = "0.11.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "df1d3c3b53da64cf5760482273a98e575c651a67eec7f77df96b5b642de8f039"
 
-[[package]]
-name = "lock_api"
-version = "0.4.14"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "224399e74b87b5f3557511d98dff8b14089b3dadafcab6bb93eab67d3aace965"
-dependencies = [
- "scopeguard",
-]
-
 [[package]]
 name = "log"
 version = "0.4.29"
@@ -662,29 +653,6 @@ version = "1.70.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "384b8ab6d37215f3c5301a95a4accb5d64aa607f1fcb26a11b5303878451b4fe"
 
-[[package]]
-name = "parking_lot"
-version = "0.12.5"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "93857453250e3077bd71ff98b6a65ea6621a19bb0f559a85248955ac12c45a1a"
-dependencies = [
- "lock_api",
- "parking_lot_core",
-]
-
-[[package]]
-name = "parking_lot_core"
-version = "0.9.12"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "2621685985a2ebf1c516881c026032ac7deafcda1a2c9b7850dc81e3dfcb64c1"
-dependencies = [
- "cfg-if",
- "libc",
- "redox_syscall",
- "smallvec",
- "windows-link",
-]
-
 [[package]]
 name = "percent-encoding"
 version = "2.3.2"
@@ -888,15 +856,6 @@ dependencies = [
  "libc",
 ]
 
-[[package]]
-name = "redox_syscall"
-version = "0.5.18"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "ed2bf2547551a7053d6fdfafda3f938979645c44812fbfcda098faae3f1a362d"
-dependencies = [
- "bitflags",
-]
-
 [[package]]
 name = "regex"
 version = "1.12.3"
@@ -939,12 +898,6 @@ dependencies = [
  "windows-sys 0.61.2",
 ]
 
-[[package]]
-name = "scopeguard"
-version = "1.2.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49"
-
 [[package]]
 name = "serde"
 version = "1.0.228"
@@ -1092,7 +1045,6 @@ dependencies = [
  "bytes",
  "libc",
  "mio",
- "parking_lot",
  "pin-project-lite",
  "signal-hook-registry",
  "socket2",

nyash_server/Cargo.toml

@@ -17,7 +17,7 @@ prost = "0.14.3"
 rand = "0.9.2"
 redb = "3.1.0"
 serde = { version = "1.0.228", features = ["derive"] }
-tokio = { version = "1.49.0", features = ["full"] }
+tokio = { version = "1.49.0", features = ["rt", "macros", "signal"] }
 tokio-tungstenite = "0.28.0"
 toml = "1.0.1"
 tonic = "0.14.5"

nyash_server/src/database.rs

@@ -10,7 +10,7 @@ use std::{path::Path, u64, u128};
 // const RANGE_STATUS: MultimapTableDefinition<bool, u128> =
 //     MultimapTableDefinition::new("range_status");
 
-const JOBS_TABLE: TableDefinition<u64, (u16, u128, u128, u64, u64)> = TableDefinition::new("jobs");
+const JOBS_TABLE: TableDefinition<u64, (u32, u128, u128, u64, u64)> = TableDefinition::new("jobs");
 const JOBS_FREE_IDS: TableDefinition<u64, ()> = TableDefinition::new("jobs_free_ids");
 
 const FOUND_KEYS_TABLE: TableDefinition<(u128, u128), u64> = TableDefinition::new("found_keys");
@@ -24,7 +24,7 @@ fn get_timestump() -> u64 {
 #[derive(Copy, Clone, Debug)]
 pub struct JobRecord {
     pub id: u64,
-    pub range_id: u16,
+    pub range_id: u32,
     pub tweak_key: u128,
     pub start_key: u128,
     pub len: u64,
@@ -32,7 +32,7 @@ pub struct JobRecord {
 }
 
 impl JobRecord {
-    pub fn from_data(k: u64, v: &(u16, u128, u128, u64, u64)) -> Self {
+    pub fn from_data(k: u64, v: &(u32, u128, u128, u64, u64)) -> Self {
         Self {
             id: k,
             range_id: v.0,
@@ -47,7 +47,7 @@ impl JobRecord {
         item: Result<
             (
                 redb::AccessGuard<'_, u64>,
-                redb::AccessGuard<'_, (u16, u128, u128, u64, u64)>,
+                redb::AccessGuard<'_, (u32, u128, u128, u64, u64)>,
             ),
             redb::StorageError,
         >,
@@ -56,7 +56,7 @@ impl JobRecord {
         Ok(Self::from_data(ag.0.value(), &ag.1.value()))
     }
 
-    pub fn get_value(&self) -> (u16, u128, u128, u64, u64) {
+    pub fn get_value(&self) -> (u32, u128, u128, u64, u64) {
         return (
             self.range_id,
             self.tweak_key,
@@ -69,7 +69,7 @@ impl JobRecord {
     // search for abandoned job
     // returns first job older than 3 hours
     pub fn get_staled_job(
-        jobs_table: &mut redb::Table<'_, u64, (u16, u128, u128, u64, u64)>,
+        jobs_table: &mut redb::Table<'_, u64, (u32, u128, u128, u64, u64)>,
         timeout: u64,
     ) -> Result<Option<JobRecord>, redb::Error> {
         let current_time = get_timestump();
@@ -103,7 +103,7 @@ impl JobRecord {
     // What we want is to check max and min ids
     // and return id lower than min or grater than max
     fn get_new_job_id(
-        jobs_table: &redb::Table<'_, u64, (u16, u128, u128, u64, u64)>,
+        jobs_table: &redb::Table<'_, u64, (u32, u128, u128, u64, u64)>,
     ) -> Result<u64, redb::Error> {
         // If table is empty return 0 ID
         if jobs_table.len()? == 0 {
@@ -132,8 +132,8 @@ impl JobRecord {
     }
 }
 
-const RANGES: TableDefinition<u16, (u128, u128, u128, u128)> = TableDefinition::new("ranges");
+const RANGES: TableDefinition<u32, (u128, u128, u128, u128)> = TableDefinition::new("ranges");
-const RANGES_TO_USE: TableDefinition<u16, ()> = TableDefinition::new("ranges_to_use");
+const RANGES_TO_USE: TableDefinition<u32, ()> = TableDefinition::new("ranges_to_use");
  
 // to not owerflow u128 we calculate it in range 0..max and then substruct
 // const SUB_DIV: u128 = 0x1000100010001000100010001;
@@ -147,7 +147,7 @@ const TWEAKS_PER_RANGE_ID_VICE: u128 = TWEAKS_PER_RANGE - 1;
 
 #[derive(Copy, Clone, Debug)]
 struct RangeRecord {
-    id: u16,
+    id: u32,
 
     tweak_current: u128,
     tweak_end: u128,
@@ -157,7 +157,7 @@ struct RangeRecord {
 
 impl RangeRecord {
 
-    fn from_value(k: u16, v: &(u128, u128, u128, u128)) -> Self {
+    fn from_value(k: u32, v: &(u128, u128, u128, u128)) -> Self {
         Self {
             id: k,
             tweak_current: v.0,
@@ -167,7 +167,7 @@ impl RangeRecord {
         }
     }
 
-    fn new(id: u16) -> Self {
+    fn new(id: u32) -> Self {
 
         let tweak_start: u128 = id as u128 * TWEAKS_PER_RANGE;
         Self {
@@ -259,8 +259,8 @@ impl RangeRecord {
 
 #[derive(Debug)]
 struct RangesTable<'a> {
-    ranges: redb::Table<'a, u16, (u128, u128, u128, u128)>,
+    ranges: redb::Table<'a, u32, (u128, u128, u128, u128)>,
-    ranges_to_use: redb::Table<'a, u16, ()>,
+    ranges_to_use: redb::Table<'a, u32, ()>,
 }
 
 impl<'a> RangesTable<'a> {
@@ -281,7 +281,7 @@ impl<'a> RangesTable<'a> {
         })
     }
 
-    fn get_range(&self, k: u16) -> Result<Option<RangeRecord>, redb::Error> {
+    fn get_range(&self, k: u32) -> Result<Option<RangeRecord>, redb::Error> {
         Ok(self
             .ranges
             .get(k)?
@@ -501,30 +501,43 @@ pub fn db_get_progress(db: &Database) -> Result<f64, redb::Error> {
 
 
 fn db_init(db: &Database) -> Result<(), redb::Error> {
-    let transct: redb::WriteTransaction = db.begin_write()?;
+    println!("Start DB init.");
-    {
-        let mut ranges = transct.open_table(RANGES)?;
-        let mut ranges_to_use = transct.open_table(RANGES_TO_USE)?;
     
-        // create new ranges
+    let trx: redb::WriteTransaction = db.begin_write()?;
-        for r_id in 0..=u16::MAX {
+    {
-            let r = RangeRecord::new(r_id);
+        let mut ranges = trx.open_table(RANGES)?;
-            ranges.insert(r_id, r.value())?;
+        let mut ranges_to_use = trx.open_table(RANGES_TO_USE)?;
-            ranges_to_use.insert(r_id, ())?;
+        for k in 0..=u16::MAX {
+            // create new ranges
+            let r = RangeRecord::new(k as u32);
+            ranges.insert(k as u32, r.value())?;
+            ranges_to_use.insert(k as u32, ())?;
         }
+
+        println!("ensure that all tables is created");
+        // ensure that all tables is created
+        trx.open_table(FOUND_KEYS_TABLE)?.len()?;
+        trx.open_table(JOBS_TABLE)?.len()?;
+        trx.open_table(JOBS_FREE_IDS)?.len()?;
     }
-    transct.commit()?;
+    println!("Commit transaction");
+    trx.commit()?;
+    println!("Done!");
+    
     Ok(())
 }
 
 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 trx = db.begin_write()?;
     let ranges_len = {
+        println!("Opening RANGES to get len");
         let ranges_table = trx.open_table(RANGES)?;
+        println!("Try to get len");
         ranges_table.len()?
     };
-    trx.close()?;
+    println!("Commit transaction");
+    trx.commit()?;
     if ranges_len == 0 {
         db_init(&db)?;
     }

nyash_server/src/num_utils.rs

@@ -4,15 +4,15 @@
 // }
 
 
-fn u128_to_u32arr(a: u128) -> [u32;4] {
+// fn u128_to_u32arr(a: u128) -> [u32;4] {
-    let mut res = [0u32;4];
+//     let mut res = [0u32;4];
-    let a_bytes = a.to_le_bytes();
+//     let a_bytes = a.to_le_bytes();
-    let chunks = a_bytes.as_chunks::<4>().0;
+//     let chunks = a_bytes.as_chunks::<4>().0;
-    for i in 0..4 {
+//     for i in 0..4 {
-        res[i] = u32::from_le_bytes(chunks[i]);
+//         res[i] = u32::from_le_bytes(chunks[i]);
-    }
+//     }
-    return res;
+//     return res;
-}
+// }
 
 pub fn u128_to_u64arr(a: u128) -> [u64;2] {
     let mut res = [0u64;2];
@@ -32,143 +32,143 @@ pub fn u64arr_to_u128(a:(u64,u64)) -> u128 {
     return u128::from_le_bytes(bytes_data);
 }
 
-pub fn add_u128_to_u256(a: &[u32; 8], b: u128) -> ([u32; 8], bool) {
+// pub fn add_u128_to_u256(a: &[u32; 8], b: u128) -> ([u32; 8], bool) {
-    let mut res: [u32; 8] = [0; 8];
+//     let mut res: [u32; 8] = [0; 8];
-    let mut carry = false;
+//     let mut carry = false;
 
-    // convert b u128 value to bytes then to u32 and then add it to a
+//     // convert b u128 value to bytes then to u32 and then add it to a
-    b.to_le_bytes()
+//     b.to_le_bytes()
-        .as_chunks::<4>()
+//         .as_chunks::<4>()
-        .0
+//         .0
-        .iter()
+//         .iter()
-        .enumerate()
+//         .enumerate()
-        .for_each(|(i, sl)| {
+//         .for_each(|(i, sl)| {
-            (res[i], carry) = a[i].carrying_add(u32::from_le_bytes(*sl), carry);
+//             (res[i], carry) = a[i].carrying_add(u32::from_le_bytes(*sl), carry);
-        });
+//         });
 
-    // propagate carry till the end of a
+//     // propagate carry till the end of a
-    for idx in 4..8 {
+//     for idx in 4..8 {
-        (res[idx], carry) = a[idx].carrying_add(0, carry);
+//         (res[idx], carry) = a[idx].carrying_add(0, carry);
-    }
+//     }
 
-    return (res, carry);
+//     return (res, carry);
-}
+// }
 
-fn add_u32_to_u256(a: &[u32; 8], b: u32) -> ([u32; 8], bool) {
+// fn add_u32_to_u256(a: &[u32; 8], b: u32) -> ([u32; 8], bool) {
-    let mut res: [u32; 8] = [0; 8];
+//     let mut res: [u32; 8] = [0; 8];
-    let mut carry = false;
+//     let mut carry = false;
-    (res[0], carry) = a[0].carrying_add(b, carry);
+//     (res[0], carry) = a[0].carrying_add(b, carry);
 
-    for idx in 1..8 {
+//     for idx in 1..8 {
-        (res[idx], carry) = a[idx].carrying_add(0, carry);
+//         (res[idx], carry) = a[idx].carrying_add(0, carry);
-    }
+//     }
 
-    return (res, carry);
+//     return (res, carry);
-}
+// }
 
-fn add_u32_to_u256_(a: &mut [u32; 8], b: u32) -> bool {
+// fn add_u32_to_u256_(a: &mut [u32; 8], b: u32) -> bool {
-    let mut carry = false;
+//     let mut carry = false;
-    (a[0], carry) = a[0].carrying_add(b, carry);
+//     (a[0], carry) = a[0].carrying_add(b, carry);
 
-    for idx in 1..8 {
+//     for idx in 1..8 {
-        (a[idx], carry) = a[idx].carrying_add(0, carry);
+//         (a[idx], carry) = a[idx].carrying_add(0, carry);
-    }
+//     }
 
-    return carry;
+//     return carry;
-}
+// }
 
-fn bytes_from_chars(chars_chunk: &[char]) -> [u8; 4] {
+// fn bytes_from_chars(chars_chunk: &[char]) -> [u8; 4] {
-    let mut res: [u8; 4] = [0; 4];
+//     let mut res: [u8; 4] = [0; 4];
-    let mut idx: usize = 0;
+//     let mut idx: usize = 0;
-    chars_chunk.chunks_exact(2).for_each(|b_c| {
+//     chars_chunk.chunks_exact(2).for_each(|b_c| {
-        if idx < 4 {
+//         if idx < 4 {
-            match u8::from_str_radix(&b_c.iter().collect::<String>(), 16) {
+//             match u8::from_str_radix(&b_c.iter().collect::<String>(), 16) {
-                Ok(n) => res[idx] = n,
+//                 Ok(n) => res[idx] = n,
-                Err(_) => (),
+//                 Err(_) => (),
-            }
+//             }
-            idx += 1;
+//             idx += 1;
-        }
+//         }
-    });
+//     });
-    return res;
+//     return res;
-}
+// }
 
-fn bignum_from_hex(hex: &str) -> [u32; 8] {
+// fn bignum_from_hex(hex: &str) -> [u32; 8] {
-    let mut res: [u32; 8] = [0; 8];
+//     let mut res: [u32; 8] = [0; 8];
-    let mut idx: usize = 0;
+//     let mut idx: usize = 0;
-    let chars_hex = hex.chars().collect::<Vec<char>>();
+//     let chars_hex = hex.chars().collect::<Vec<char>>();
-    chars_hex
+//     chars_hex
-        .chunks_exact(8)
+//         .chunks_exact(8)
-        .rev()
+//         .rev()
-        .map(|chunk| bytes_from_chars(chunk))
+//         .map(|chunk| bytes_from_chars(chunk))
-        .for_each(|b_arr| {
+//         .for_each(|b_arr| {
-            if idx < 8 {
+//             if idx < 8 {
-                res[idx] = u32::from_be_bytes(b_arr);
+//                 res[idx] = u32::from_be_bytes(b_arr);
-                idx += 1;
+//                 idx += 1;
-            }
+//             }
-        });
+//         });
-    return res;
+//     return res;
-}
+// }
 
-fn hex_fmt_byte(n: u32) -> String {
+// fn hex_fmt_byte(n: u32) -> String {
-    let res: String = n
+//     let res: String = n
-        .to_be_bytes()
+//         .to_be_bytes()
-        .iter()
+//         .iter()
-        .map(|b| format!("{:02x}", b))
+//         .map(|b| format!("{:02x}", b))
-        .collect();
+//         .collect();
-    return res;
+//     return res;
-}
+// }
 
-fn bignum_to_hex(a: &[u32; 8]) -> String {
+// fn bignum_to_hex(a: &[u32; 8]) -> String {
-    let res: String = a
+//     let res: String = a
-        .iter()
+//         .iter()
-        .rev()
+//         .rev()
-        .map(|n| hex_fmt_byte(*n))
+//         .map(|n| hex_fmt_byte(*n))
-        .collect::<Vec<String>>()
+//         .collect::<Vec<String>>()
-        .join("");
+//         .join("");
-    return res;
+//     return res;
-}
+// }
 
-#[cfg(test)]
+// #[cfg(test)]
-mod num_utils_tests {
+// mod num_utils_tests {
-    use super::*;
+//     use super::*;
 
-    #[test]
+//     #[test]
-    fn test_add() {
+//     fn test_add() {
-        use std::io::{BufRead, BufReader};
+//         use std::io::{BufRead, BufReader};
-        use std::process::{Command, Stdio};
+//         use std::process::{Command, Stdio};
 
-        let test_gen_cmd = "/home/kira/Development/Rust/nyash-aes-xts256-plain64/nyash_client/src/tests/gen_test_data.py";
+//         let test_gen_cmd = "/home/kira/Development/Rust/nyash-aes-xts256-plain64/nyash_client/src/tests/gen_test_data.py";
-        let mut child = Command::new(test_gen_cmd)
+//         let mut child = Command::new(test_gen_cmd)
-            .stdout(Stdio::piped())
+//             .stdout(Stdio::piped())
-            .spawn()
+//             .spawn()
-            .unwrap();
+//             .unwrap();
 
-        let gen_stdout = child
+//         let gen_stdout = child
-            .stdout
+//             .stdout
-            .take()
+//             .take()
-            .ok_or("Failed to capture stdout")
+//             .ok_or("Failed to capture stdout")
-            .unwrap();
+//             .unwrap();
-        let gen_reader = BufReader::new(gen_stdout);
+//         let gen_reader = BufReader::new(gen_stdout);
 
-        for r_line in gen_reader.lines() {
+//         for r_line in gen_reader.lines() {
-            let test_line: String = r_line.unwrap(); // Handle any I/O errors
+//             let test_line: String = r_line.unwrap(); // Handle any I/O errors
-            let test_data_line = test_line.split(' ').collect::<Vec<&str>>();
+//             let test_data_line = test_line.split(' ').collect::<Vec<&str>>();
-            let num_to_add = u32::from_str_radix(test_data_line[0], 10).unwrap();
+//             let num_to_add = u32::from_str_radix(test_data_line[0], 10).unwrap();
-            let t0 = bignum_from_hex(test_data_line[1]);
+//             let t0 = bignum_from_hex(test_data_line[1]);
 
-            let t1_test = add_u32_to_u256(&t0, 1).0;
+//             let t1_test = add_u32_to_u256(&t0, 1).0;
-            let t2_test = add_u32_to_u256(&t0, num_to_add).0;
+//             let t2_test = add_u32_to_u256(&t0, num_to_add).0;
 
-            let res_actual = format!(
+//             let res_actual = format!(
-                "{} {} {} {}",
+//                 "{} {} {} {}",
-                num_to_add,
+//                 num_to_add,
-                bignum_to_hex(&t0),
+//                 bignum_to_hex(&t0),
-                bignum_to_hex(&t1_test),
+//                 bignum_to_hex(&t1_test),
-                bignum_to_hex(&t2_test)
+//                 bignum_to_hex(&t2_test)
-            );
+//             );
-            assert_eq!(test_line, res_actual);
+//             assert_eq!(test_line, res_actual);
-        }
+//         }
 
-        let _ = child.wait().unwrap();
+//         let _ = child.wait().unwrap();
-    }
+//     }
-}
+// }

nyash_server/src/server.rs

@@ -25,6 +25,31 @@ mod database;
 mod num_utils;
 mod config;
 
+async fn shutdown_signal() {
+    let ctrl_c = async {
+        tokio::signal::ctrl_c()
+            .await
+            .expect("Failed to install Ctrl+C handler");
+    };
+
+    #[cfg(unix)]
+    let terminate = async {
+        tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate())
+            .expect("Failed to install SIGTERM handler")
+            .recv()
+            .await;
+    };
+
+    #[cfg(not(unix))]
+    let terminate = std::future::pending::<()>();
+
+    tokio::select! {
+        _ = ctrl_c => println!("Received Ctrl+C, initiating shutdown"),
+        _ = terminate => println!("Received SIGTERM, initiating shutdown"),
+    }
+}
+
+
 fn work_rec_to_work_data(work_rec: database::JobRecord) -> WorkData {
     let tw_key = num_utils::u128_to_u64arr(work_rec.tweak_key);
     let st_key = num_utils::u128_to_u64arr(work_rec.start_key);
@@ -42,8 +67,8 @@ fn work_rec_to_work_data(work_rec: database::JobRecord) -> WorkData {
 
 
 // Background task to update progress from DB
-async fn update_service_progress(progress: Arc<AsyncRwLock<f64>>, db: Arc<redb::Database>) {
+async fn update_service_progress(progress: Arc<AsyncRwLock<f64>>, db: Arc<redb::Database>, term: Arc<AsyncRwLock<bool>>) {
-    loop {
+    while *term.read().await == false {
         tokio::time::sleep(tokio::time::Duration::from_secs(5)).await;
         match database::db_get_progress(&db) {
             Ok(p) => {
@@ -66,7 +91,11 @@ pub struct NyashService {
 #[tonic::async_trait]
 impl NyashLuks for NyashService {
     async fn request_work(&self, request: Request<WorkRequest>) -> Result<Response<WorkReply>, Status> {
-        println!("Got a request: {:?}", request);
+        let rem_addr = match request.remote_addr() {
+            Some(v) => v.to_string(),
+            None => "None".to_string()
+        };
+        println!("Remoute: {}, request_work.", rem_addr);
 
 
         if *self.key_found.read().await == true {
@@ -98,7 +127,12 @@ impl NyashLuks for NyashService {
     }
 
     async fn commit_work(&self, request: Request<WorkCommit>) -> Result<Response<CommitReply>, Status> {
-        println!("Got a request: {:?}", request);
+        let rem_addr = match request.remote_addr() {
+            Some(v) => v.to_string(),
+            None => "None".to_string()
+        };
+        println!("Remoute: {}, commit_work.", rem_addr);
+
         let work_commit = request.into_inner();
 
         let put_k_no_err: bool = match work_commit.result {
@@ -133,7 +167,12 @@ impl NyashLuks for NyashService {
         }
     }
 
-    async fn request_progress(&self, _request: Request<ProgressRequest>) -> Result<Response<ProgressReply>, Status> {
+    async fn request_progress(&self, request: Request<ProgressRequest>) -> Result<Response<ProgressReply>, Status> {
+        let rem_addr = match request.remote_addr() {
+            Some(v) => v.to_string(),
+            None => "None".to_string()
+        };
+        println!("Remoute: {}, request_progress.", rem_addr);
         Ok(Response::new(ProgressReply {
             key_found: *self.key_found.read().await,
             progress: *self.progress.read().await,
@@ -153,7 +192,7 @@ struct ProgArgs {
 
 
 
-#[tokio::main]
+#[tokio::main(flavor = "current_thread")]
 async fn main() -> Result<(), Box<dyn std::error::Error>> {
 
     let args = ProgArgs::parse();
@@ -165,7 +204,9 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
     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");
+    let mut db = database::db_open(&db_path).expect("Error Opening database");
+    println!("Compact db");
+    db.compact()?;
 
     // Get the address to bind to
     let addr = format!("{}:{}",server_config.bind_addr, server_config.listen_port);
@@ -190,22 +231,33 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
     let shared_progress = Arc::new(AsyncRwLock::new(progress));
     let shared_db = Arc::new(db);
     let shared_key_found = Arc::new(AsyncRwLock::new(false));
-
+    let shared_termination: Arc<AsyncRwLock<bool>> = Arc::new(AsyncRwLock::new(false));
     // let service_state = Arc::new(AsyncRwLock::new(ServiceState {
     //     db: Arc::new(db),
     //     key_found: false,
     //     progress: progress,
     // }));
 
+    let term_clone = shared_termination.clone();
+    let (shutdown_tx, shutdown_rx) = tokio::sync::oneshot::channel();
+    tokio::spawn(async move {
+        shutdown_signal().await;
+        //signaling that we should stop
+        let mut guard = term_clone.write().await;
+        *guard = true;
+        let _ = shutdown_tx.send(());
+    });
+
     // Spawn the periodic updater
-    tokio::spawn(update_service_progress(shared_progress.clone(), shared_db.clone()));
+    tokio::spawn(update_service_progress(shared_progress.clone(), shared_db.clone(), shared_termination.clone()));
     
     let nyash_service = NyashService{db:shared_db, progress:shared_progress, key_found:shared_key_found};
 
 
     Server::builder()
         .add_service(NyashLuksServer::new(nyash_service))
-        .serve(addr)
+        .serve_with_shutdown(addr, async { shutdown_rx.await.ok(); })
         .await?;
+    
     Ok(())
 }