From 1f5a82ff918bfc77c30f81f0e13dc97ef7bd937e Mon Sep 17 00:00:00 2001 From: nganhkhoa Date: Thu, 5 May 2022 23:54:41 +0700 Subject: [PATCH] format code and add strip command --- src/main.rs | 179 +++++++++++++++++++++++++++++----------------------- 1 file changed, 100 insertions(+), 79 deletions(-) diff --git a/src/main.rs b/src/main.rs index c4440a9..dcb09d6 100644 --- a/src/main.rs +++ b/src/main.rs @@ -2,35 +2,32 @@ extern crate rocket; use std::collections::HashMap; -use std::fs; use std::error::Error; use std::path::Path; use std::sync::{Arc, Mutex}; -use std::time::Duration; -use std::thread; -use rocket::tokio as tokio; -use rocket::serde as serde; +use rocket::serde; +use rocket::tokio; -use rocket::State; use rocket::http::Status; +use rocket::State; use rocket_download_response::DownloadResponse; -use serde::{Serialize, Deserialize, json::Json}; +use serde::{json::Json, Deserialize, Serialize}; -use tokio::sync::mpsc; -use tokio::fs as async_fs; +use tokio::fs; use tokio::process::Command; +use tokio::sync::mpsc; use uuid::Uuid; #[derive(Clone, Debug)] struct RequestBuildArgs { - token: String, - vendor_name: String, - vendor_metamask_account: String, - aead_key: String, - ssl_private_key: String + token: String, + vendor_name: String, + vendor_metamask_account: String, + aead_key: String, + ssl_private_key: String, } #[derive(Debug)] @@ -42,9 +39,6 @@ enum Request { } async fn perform_build(args: &RequestBuildArgs) -> Result<(), Box> { - println!("build with token: {}", args.token); - println!("python setup.py {} {} {} {}", args.vendor_name, args.vendor_metamask_account, args.aead_key, args.ssl_private_key); - Command::new("python3") .current_dir("/home/r00t/work/ollvm/ssm-for-test/ssm/dynamic/") .env("VENDOR_NAME", &args.vendor_name) @@ -58,36 +52,61 @@ async fn perform_build(args: &RequestBuildArgs) -> Result<(), Box> { Command::new("/home/r00t/work/ollvm/llvm-install/bin/go") .current_dir("/home/r00t/work/ollvm/ssm-for-test/ssm/") - .env("LD_LIBRARY_PATH", "/home/r00t/work/ollvm/llvm-install/lib64") + .env( + "LD_LIBRARY_PATH", + "/home/r00t/work/ollvm/llvm-install/lib64", + ) .args([ - "build", "-a", - "-o", &format!("ssm-{}", args.token), - "-gccgoflags=\"-static-libgo -Wl,--version-script=ssm.version\"", - "-mllvm", "-obfuscation=gvo", - "-mllvm", "-obfuscation=sub", - "-mllvm", "-obfuscation=flatten", - "-mllvm", "-obfuscation=idr-branch", + "build", + "-a", + "-o", + &format!("ssm-{}", args.token), + "-gccgoflags", + "-static-libgo -Wl,--version-script=ssm.version", + "-mllvm", + "-obfuscation=gvo", + "-mllvm", + "-obfuscation=sub", + "-mllvm", + "-obfuscation=flatten", + "-mllvm", + "-obfuscation=idr-branch", ]) .spawn()? .wait() .await?; - async_fs::create_dir("out").await.ok(); - async_fs::rename( + Command::new("strip") + .current_dir("/home/r00t/work/ollvm/ssm-for-test/ssm/") + .args([ + "-s", + &format!("ssm-{}", args.token) + ]) + .spawn()? + .wait() + .await?; + + fs::create_dir("out").await.ok(); + fs::rename( format!("/home/r00t/work/ollvm/ssm-for-test/ssm/ssm-{}", args.token), - format!("./out/ssm-{}", args.token)) - .await - .ok(); + format!("./out/ssm-{}", args.token), + ) + .await + .ok(); println!("build complete out/ssm-{}", args.token); Ok(()) } -async fn process_request(request: Request, token_list: TokenList, build_sender: mpsc::Sender) { +async fn process_request( + request: Request, + token_list: TokenList, + build_sender: mpsc::Sender, +) { match request { Request::InitToken(token) => { let mut token_list = token_list.lock().unwrap(); token_list.insert(token.clone(), TokenStatus::Initialized); - }, + } Request::SetToken(token, new_status) => { println!("set token {:?} {:?}", token, new_status); let mut token_list = token_list.lock().unwrap(); @@ -97,7 +116,7 @@ async fn process_request(request: Request, token_list: TokenList, build_sender: } Request::Build(args) => { build_sender.send(args).await.ok(); - }, + } Request::Delete(token) => { fs::remove_file(format!("./out/ssm-{}", token)).await.ok(); let mut token_list = token_list.lock().unwrap(); @@ -118,21 +137,19 @@ enum TokenStatus { } type TokenList = Arc>>; -// type BuildQueue = Arc>>; struct AppState { pub sender: mpsc::Sender, pub token_list: TokenList, } -#[derive(Debug)] -#[derive(Deserialize)] +#[derive(Debug, Deserialize)] #[serde(crate = "rocket::serde")] struct BuildArgs { vendor_name: String, vendor_metamask_account: String, aead_key: String, - ssl_private_key: String + ssl_private_key: String, } #[derive(Serialize)] @@ -141,7 +158,7 @@ struct BuildResponse { token: String, } -#[post("/build", data="")] +#[post("/build", data = "")] async fn build(state: &State, args: Json) -> Json { let args = &*args; println!("build with args={:?}", args); @@ -149,46 +166,36 @@ async fn build(state: &State, args: Json) -> Json, args: Json) -> String { let args = &*args; println!("args: {:?}", args); let token_list = state.token_list.lock().unwrap(); let result = match token_list.get(&args.token) { - Some(state) => { - format!("{:?}", state).to_string() - }, - _ => { - "token non exist".to_string() - } + Some(state) => format!("{:?}", state).to_string(), + _ => "token non exist".to_string(), }; println!("{}", result); result @@ -198,51 +205,62 @@ async fn status(state: &State, args: Json) -> String { async fn download(state: &State, token: String) -> Result { println!("download with token={}", token); - { // in a seperate block to prevent mutex not implement Sync error + { + // in a seperate block to prevent mutex not implement Sync error let token_list = state.token_list.lock().unwrap(); let file_ready = match token_list.get(&token) { Some(TokenStatus::Finished) => true, - _ => false + _ => false, }; if !file_ready { return Err(Status::NotFound); } } + // load the content into memory, so that we can delete file let path = Path::join(Path::new("out"), format!("ssm-{}", token)); - let content = async_fs::read(path) - .await - .map_err(|_| Status::NotFound)?; + let content = fs::read(path).await.map_err(|_| Status::NotFound)?; - state.sender.send(Request::Delete(token)) - .await - .ok(); + state.sender.send(Request::Delete(token)).await.ok(); Ok(DownloadResponse::from_vec(content, Some("ssm"), None)) } -async fn setup_request_receiver(mut receiver: mpsc::Receiver, token_list: TokenList, builder_sender: mpsc::Sender) { +async fn setup_request_receiver( + mut receiver: mpsc::Receiver, + token_list: TokenList, + builder_sender: mpsc::Sender, +) { while let Some(request) = receiver.recv().await { println!("requesting {:?}", request); process_request(request, token_list.clone(), builder_sender.clone()).await; } } -async fn setup_builder(mut receiver: mpsc::Receiver, sender: mpsc::Sender) { +async fn setup_builder( + mut receiver: mpsc::Receiver, + sender: mpsc::Sender, +) { while let Some(args) = receiver.recv().await { println!("requesting build {:?}", args); - sender.send(Request::SetToken(args.token.clone(), TokenStatus::Building)) + sender + .send(Request::SetToken(args.token.clone(), TokenStatus::Building)) .await .ok(); let build_result = perform_build(&args).await.ok(); match build_result { Some(_) => { - sender.send(Request::SetToken(args.token.clone(), TokenStatus::Finished)) + sender + .send(Request::SetToken(args.token.clone(), TokenStatus::Finished)) .await .ok(); } _ => { - sender.send(Request::SetToken(args.token.clone(), TokenStatus::Error("Build is not success".into()))) + sender + .send(Request::SetToken( + args.token.clone(), + TokenStatus::Error("Build is not success".into()), + )) .await .ok(); } @@ -252,8 +270,8 @@ async fn setup_builder(mut receiver: mpsc::Receiver, sender: m #[tokio::main] async fn main() { - let (sender, mut receiver) = mpsc::channel(100); - let (builder_sender, mut builder_receiver) = mpsc::channel(100); + let (sender, receiver) = mpsc::channel(100); + let (builder_sender, builder_receiver) = mpsc::channel(100); let token_list: TokenList = Arc::new(Mutex::new(HashMap::new())); let worker = AppState { @@ -261,16 +279,19 @@ async fn main() { token_list: token_list.clone(), }; - let receiver_runner = tokio::spawn( - setup_request_receiver(receiver, token_list.clone(), builder_sender.clone())); - let builder_runner = tokio::spawn( - setup_builder(builder_receiver, sender.clone())); + tokio::spawn(setup_request_receiver( + receiver, + token_list.clone(), + builder_sender.clone(), + )); + tokio::spawn(setup_builder(builder_receiver, sender.clone())); - let rocket_runner = rocket::build() + rocket::build() .manage(worker) .mount("/", routes![build]) .mount("/", routes![status]) .mount("/", routes![download]) .launch() - .await; + .await + .ok(); }