From 21fbdd39ec5a7e405e6112ce2daa28be5b838091 Mon Sep 17 00:00:00 2001 From: scx1332 Date: Sun, 29 Sep 2024 22:43:54 +0200 Subject: [PATCH] Add groups to key server --- Cargo.lock | 3 +- crates/web3_account_server/Cargo.toml | 3 +- crates/web3_account_server/src/main.rs | 80 ++++++++++++++++++++++++-- 3 files changed, 78 insertions(+), 8 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index fef1b726..fee89df6 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4401,7 +4401,7 @@ dependencies = [ [[package]] name = "web3_account_server" -version = "0.4.11" +version = "0.4.12" dependencies = [ "actix-cors", "actix-web", @@ -4412,6 +4412,7 @@ dependencies = [ "serde", "serde_json", "structopt", + "tokio", ] [[package]] diff --git a/crates/web3_account_server/Cargo.toml b/crates/web3_account_server/Cargo.toml index b4b07c4f..0e9a13dc 100644 --- a/crates/web3_account_server/Cargo.toml +++ b/crates/web3_account_server/Cargo.toml @@ -5,7 +5,7 @@ edition = "2021" name = "web3_account_server" license = "MIT" repository = "https://github.com/scx1332/" -version = "0.4.11" +version = "0.4.12" [dependencies] actix-cors = { workspace = true } @@ -17,3 +17,4 @@ log = { workspace = true } serde = { workspace = true } serde_json = { workspace = true } structopt = { workspace = true } +tokio = { workspace = true } diff --git a/crates/web3_account_server/src/main.rs b/crates/web3_account_server/src/main.rs index a3d9ac53..234c117a 100644 --- a/crates/web3_account_server/src/main.rs +++ b/crates/web3_account_server/src/main.rs @@ -5,7 +5,8 @@ use actix_web_httpauth::middleware::HttpAuthentication; use std::env; use std::fs::OpenOptions; use std::io::{BufReader, BufWriter}; -use std::sync::{Arc, Mutex}; +use std::path::{Component, Path, PathBuf}; +use std::sync::Arc; use structopt::StructOpt; fn read_results(file_name: &str) -> Vec { @@ -54,7 +55,7 @@ fn get(file_name: &str) -> std::io::Result> { #[derive(Clone)] struct AppState { - lock: Arc>, + lock: Arc>, file_name: String, } @@ -83,7 +84,7 @@ pub struct CliOptions { } async fn add_to_queue(data: web::Data, item: String) -> impl Responder { - let _lock = data.lock.lock().unwrap(); + let _lock = data.lock.lock().await; let Ok(private_key) = hex::decode(item.replace("0x", "")) else { return HttpResponse::BadRequest().body("Invalid item type"); }; @@ -101,14 +102,14 @@ async fn add_to_queue(data: web::Data, item: String) -> impl Responder } async fn count(data: web::Data) -> impl Responder { - let _lock = data.lock.lock().unwrap(); + let _lock = data.lock.lock().await; let file_name = &data.file_name; let results = read_results(file_name); HttpResponse::Ok().body(results.len().to_string()) } async fn get_from_queue(data: web::Data) -> impl Responder { - let _lock = data.lock.lock().unwrap(); + let _lock = data.lock.lock().await; match get(&data.file_name) { Ok(Some(item)) => HttpResponse::Ok().body(item), Ok(None) => HttpResponse::BadRequest().body("Queue is empty"), @@ -119,6 +120,70 @@ async fn get_from_queue(data: web::Data) -> impl Responder { } } +fn get_file_name_from_filename_and_group(file_name: &str, group: &str) -> String { + let sanitized_group: String = group + .chars() + .filter(|c| c.is_alphanumeric() || *c == '_') // Allow only alphanumeric and underscores + .collect(); + + let sanitized_file_name = Path::new(file_name) + .components() + .filter(|comp| matches!(comp, Component::Normal(_))) // Allow only valid file names + .collect::(); + + format!("{}_{}", sanitized_group, sanitized_file_name.display()) +} + +async fn add_to_queue_group( + data: web::Data, + path: web::Path, + item: String, +) -> impl Responder { + let _lock = data.lock.lock().await; + let group = path.into_inner(); + let file_name = get_file_name_from_filename_and_group(&data.file_name, &group); + let Ok(private_key) = hex::decode(item.replace("0x", "")) else { + return HttpResponse::BadRequest().body("Invalid item type"); + }; + if private_key.len() != 32 { + return HttpResponse::BadRequest().body("Invalid item length"); + } + match add(hex::encode(private_key), &file_name) { + Ok(true) => HttpResponse::Ok().body("Added to the queue"), + Ok(false) => HttpResponse::Ok().body("Item already in the queue"), + Err(e) => { + log::error!("Error adding item: {}", e); + HttpResponse::InternalServerError().finish() + } + } +} + +async fn count_group(data: web::Data, path: web::Path) -> impl Responder { + let _lock = data.lock.lock().await; + let group = path.into_inner(); + let file_name = get_file_name_from_filename_and_group(&data.file_name, &group); + let results = read_results(&file_name); + HttpResponse::Ok().body(results.len().to_string()) +} + +async fn get_from_queue_group( + data: web::Data, + path: web::Path, +) -> impl Responder { + let _lock = data.lock.lock().await; + + let group = path.into_inner(); + let file_name = get_file_name_from_filename_and_group(&data.file_name, &group); + match get(&file_name) { + Ok(Some(item)) => HttpResponse::Ok().body(item), + Ok(None) => HttpResponse::BadRequest().body("Queue is empty"), + Err(e) => { + log::error!("Error getting item: {}", e); + HttpResponse::InternalServerError().finish() + } + } +} + fn get_env_access_token() -> String { env::var("BEARER_KEY").unwrap_or("change_me".to_string()) } @@ -152,7 +217,7 @@ async fn main() -> std::io::Result<()> { // Load the queue from file or create a new one let app_state = AppState { - lock: Arc::new(Mutex::new(())), + lock: Arc::new(tokio::sync::Mutex::new(())), file_name: args.file_name, }; @@ -167,6 +232,9 @@ async fn main() -> std::io::Result<()> { .route("/count", web::get().to(count)) .route("/add", web::post().to(add_to_queue)) .route("/get", web::get().to(get_from_queue)) + .route("/count/{group}", web::get().to(count_group)) + .route("/add/{group}", web::post().to(add_to_queue_group)) + .route("/get/{group}", web::get().to(get_from_queue_group)) }) .bind(format!("{}:{}", args.http_addr, args.http_port))? .workers(1)