From cb2b4da0f3e835c99c8baac57aef078ef94314b1 Mon Sep 17 00:00:00 2001
From: scx1332 <sieciech.czajka@golem.network>
Date: Thu, 26 Sep 2024 00:00:54 +0200
Subject: [PATCH] Fix issues with private key server

---
 Cargo.lock                             |  3 +-
 crates/web3_account_server/Cargo.toml  |  3 +-
 crates/web3_account_server/src/main.rs | 70 +++++++++++++-------------
 3 files changed, 39 insertions(+), 37 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 6884bd59..fef1b726 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -4401,8 +4401,9 @@ dependencies = [
 
 [[package]]
 name = "web3_account_server"
-version = "0.4.8"
+version = "0.4.11"
 dependencies = [
+ "actix-cors",
  "actix-web",
  "actix-web-httpauth",
  "env_logger",
diff --git a/crates/web3_account_server/Cargo.toml b/crates/web3_account_server/Cargo.toml
index 1bca2ba3..b4b07c4f 100644
--- a/crates/web3_account_server/Cargo.toml
+++ b/crates/web3_account_server/Cargo.toml
@@ -5,9 +5,10 @@ edition = "2021"
 name = "web3_account_server"
 license = "MIT"
 repository = "https://github.com/scx1332/"
-version = "0.4.8"
+version = "0.4.11"
 
 [dependencies]
+actix-cors = { workspace = true }
 actix-web = { workspace = true }
 actix-web-httpauth = { workspace = true }
 env_logger =  { workspace = true }
diff --git a/crates/web3_account_server/src/main.rs b/crates/web3_account_server/src/main.rs
index d41b805b..a3d9ac53 100644
--- a/crates/web3_account_server/src/main.rs
+++ b/crates/web3_account_server/src/main.rs
@@ -8,17 +8,19 @@ use std::io::{BufReader, BufWriter};
 use std::sync::{Arc, Mutex};
 use structopt::StructOpt;
 
-fn add(item: String, file_name: &str) -> bool {
-    let mut results: Vec<String> = {
-        if let Ok(file) = OpenOptions::new().read(true).open(file_name) {
-            let reader = BufReader::new(file);
-            serde_json::from_reader(reader).unwrap_or_else(|_| Vec::new())
-        } else {
-            Vec::new()
-        }
-    };
+fn read_results(file_name: &str) -> Vec<String> {
+    if let Ok(file) = OpenOptions::new().read(true).open(file_name) {
+        let reader = BufReader::new(file);
+        serde_json::from_reader(reader).unwrap_or_else(|_| Vec::new())
+    } else {
+        Vec::new()
+    }
+}
+
+fn add(item: String, file_name: &str) -> std::io::Result<bool> {
+    let mut results = read_results(file_name);
     if results.contains(&item) {
-        return false;
+        return Ok(false);
     }
     results.push(item);
     let file = OpenOptions::new()
@@ -26,21 +28,17 @@ fn add(item: String, file_name: &str) -> bool {
         .create(true)
         .truncate(true)
         .open(file_name)
-        .unwrap();
+        .inspect_err(|e| log::error!("Error opening file: {}", e))?;
     let writer = BufWriter::new(file);
     serde_json::to_writer(writer, &results).unwrap();
-    true
+    Ok(true)
 }
 
-fn get(file_name: &str) -> Option<String> {
-    let mut results: Vec<String> = {
-        let file = OpenOptions::new().read(true).open(file_name).unwrap();
-        let reader = BufReader::new(file);
-        serde_json::from_reader(reader).unwrap_or_else(|_| Vec::new())
-    };
+fn get(file_name: &str) -> std::io::Result<Option<String>> {
+    let mut results = read_results(file_name);
     // get first item
     if results.is_empty() {
-        return None;
+        return Ok(None);
     }
     let item = results.remove(0);
 
@@ -48,11 +46,10 @@ fn get(file_name: &str) -> Option<String> {
     let file = OpenOptions::new()
         .write(true)
         .truncate(true)
-        .open(file_name)
-        .unwrap();
+        .open(file_name)?;
     let writer = BufWriter::new(file);
     serde_json::to_writer(writer, &results).unwrap();
-    Some(item)
+    Ok(Some(item))
 }
 
 #[derive(Clone)]
@@ -93,30 +90,32 @@ async fn add_to_queue(data: web::Data<AppState>, item: String) -> impl Responder
     if private_key.len() != 32 {
         return HttpResponse::BadRequest().body("Invalid item length");
     }
-    if add(hex::encode(private_key), &data.file_name) {
-        HttpResponse::Ok().body("Added to the queue")
-    } else {
-        HttpResponse::Ok().body("Item already in the queue")
+    match add(hex::encode(private_key), &data.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(data: web::Data<AppState>) -> impl Responder {
     let _lock = data.lock.lock().unwrap();
     let file_name = &data.file_name;
-    let results: Vec<String> = {
-        let file = OpenOptions::new().read(true).open(file_name).unwrap();
-        let reader = BufReader::new(file);
-        serde_json::from_reader(reader).unwrap_or_else(|_| Vec::new())
-    };
+    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();
-    if let Some(item) = get(&data.file_name) {
-        HttpResponse::Ok().body(item)
-    } else {
-        HttpResponse::BadRequest().body("Queue is empty")
+    match get(&data.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()
+        }
     }
 }
 
@@ -164,6 +163,7 @@ async fn main() -> std::io::Result<()> {
             .app_data(web::Data::new(app_state.clone()))
             .wrap(actix_web::middleware::Logger::default())
             .wrap(auth)
+            .wrap(actix_cors::Cors::permissive())
             .route("/count", web::get().to(count))
             .route("/add", web::post().to(add_to_queue))
             .route("/get", web::get().to(get_from_queue))