Add groups to key server

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 }

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<String> {
@@ -54,7 +55,7 @@ fn get(file_name: &str) -> std::io::Result<Option<String>> {
 
 #[derive(Clone)]
 struct AppState {
-    lock: Arc<Mutex<()>>,
+    lock: Arc<tokio::sync::Mutex<()>>,
     file_name: String,
 }
 
@@ -83,7 +84,7 @@ pub struct CliOptions {
 }
 
 async fn add_to_queue(data: web::Data<AppState>, 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<AppState>, item: String) -> impl Responder
 }
 
 async fn count(data: web::Data<AppState>) -> 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<AppState>) -> 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<AppState>) -> 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::<PathBuf>();
+
+    format!("{}_{}", sanitized_group, sanitized_file_name.display())
+}
+
+async fn add_to_queue_group(
+    data: web::Data<AppState>,
+    path: web::Path<String>,
+    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<AppState>, path: web::Path<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 results = read_results(&file_name);
+    HttpResponse::Ok().body(results.len().to_string())
+}
+
+async fn get_from_queue_group(
+    data: web::Data<AppState>,
+    path: web::Path<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);
+    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)