diff --git a/src/builder.rs b/src/builder.rs new file mode 100644 index 0000000..d072cc0 --- /dev/null +++ b/src/builder.rs @@ -0,0 +1,106 @@ +use std::error::Error; +use std::path::PathBuf; + +use rocket::tokio; + +use tokio::fs; +use tokio::process::Command; +use tokio::sync::mpsc; + +use crate::types::{Request, RequestBuildArgs, TokenStatus}; + +pub struct Builder { + llvm_install: Box, +} + +impl Builder { + pub fn new(llvm_install: Box) -> Self { + Builder { llvm_install } + } + + async fn build(&self, args: &RequestBuildArgs) -> Result<(), Box> { + Command::new("python3") + .current_dir("./ssm/dynamic/") + .env("VENDOR_NAME", &args.vendor_name) + .env("VENDOR_METAMASK_ACCOUNT", &args.vendor_metamask_account) + .env("RAW_AEAD_KEY", &args.aead_key) + .env("RAW_SSL_PRIVATE_KEY", &args.ssl_private_key) + .args(["gen.py"]) + .spawn()? + .wait() + .await?; + + Command::new(self.llvm_install.join("bin/go")) + .current_dir("./ssm/") + .env("LD_LIBRARY_PATH", self.llvm_install.join("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", + ]) + .spawn()? + .wait() + .await?; + + Command::new("strip") + .current_dir("./ssm/") + .args(["-s", &format!("ssm-{}", args.token)]) + .spawn()? + .wait() + .await?; + + fs::create_dir("out").await.ok(); + fs::rename( + format!("./ssm/ssm-{}", args.token), + format!("./out/ssm-{}", args.token), + ) + .await + .ok(); + println!("build complete out/ssm-{}", args.token); + Ok(()) + } +} + +pub async fn setup_builder( + mut receiver: mpsc::Receiver, + sender: mpsc::Sender, + builder: Builder, +) { + while let Some(args) = receiver.recv().await { + println!("requesting build {:?}", args); + sender + .send(Request::SetToken(args.token.clone(), TokenStatus::Building)) + .await + .ok(); + + let build_result = builder.build(&args).await.ok(); + match build_result { + Some(_) => { + 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()), + )) + .await + .ok(); + } + }; + } +} diff --git a/src/main.rs b/src/main.rs index dcb09d6..3234a7c 100644 --- a/src/main.rs +++ b/src/main.rs @@ -3,294 +3,76 @@ extern crate rocket; use std::collections::HashMap; use std::error::Error; -use std::path::Path; +use std::path::{Path, PathBuf}; use std::sync::{Arc, Mutex}; -use rocket::serde; use rocket::tokio; - -use rocket::http::Status; -use rocket::State; -use rocket_download_response::DownloadResponse; - -use serde::{json::Json, Deserialize, Serialize}; - -use tokio::fs; -use tokio::process::Command; use tokio::sync::mpsc; -use uuid::Uuid; +mod builder; +mod request; +mod routes; +mod types; -#[derive(Clone, Debug)] -struct RequestBuildArgs { - token: String, - vendor_name: String, - vendor_metamask_account: String, - aead_key: String, - ssl_private_key: String, -} - -#[derive(Debug)] -enum Request { - InitToken(String), - SetToken(String, TokenStatus), - Build(RequestBuildArgs), - Delete(String), -} - -async fn perform_build(args: &RequestBuildArgs) -> Result<(), Box> { - Command::new("python3") - .current_dir("/home/r00t/work/ollvm/ssm-for-test/ssm/dynamic/") - .env("VENDOR_NAME", &args.vendor_name) - .env("VENDOR_METAMASK_ACCOUNT", &args.vendor_metamask_account) - .env("RAW_AEAD_KEY", &args.aead_key) - .env("RAW_SSL_PRIVATE_KEY", &args.ssl_private_key) - .args(["gen.py"]) - .spawn()? - .wait() - .await?; - - 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", - ) - .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", - ]) - .spawn()? - .wait() - .await?; - - 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(); - println!("build complete out/ssm-{}", args.token); - Ok(()) -} - -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(); - if let Some(state) = token_list.get_mut(&token) { - *state = new_status; - } - } - 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(); - token_list.remove(&token); - } - } -} - -#[derive(Debug)] -enum TokenStatus { - // new token is assigned and waiting in queue - Initialized, - Building, - Finished, - // if the build can't be complete - // Error(msg) - Error(String), -} - -type TokenList = Arc>>; - -struct AppState { - pub sender: mpsc::Sender, - pub token_list: TokenList, -} - -#[derive(Debug, Deserialize)] -#[serde(crate = "rocket::serde")] -struct BuildArgs { - vendor_name: String, - vendor_metamask_account: String, - aead_key: String, - ssl_private_key: String, -} - -#[derive(Serialize)] -#[serde(crate = "rocket::serde")] -struct BuildResponse { - token: String, -} - -#[post("/build", data = "")] -async fn build(state: &State, args: Json) -> Json { - let args = &*args; - println!("build with args={:?}", args); - - let sender = &state.sender; - let token = Uuid::new_v4().to_string(); - - sender.send(Request::InitToken(token.clone())).await.ok(); - sender - .send(Request::Build(RequestBuildArgs { - token: token.clone(), - vendor_name: args.vendor_name.clone(), - vendor_metamask_account: args.vendor_metamask_account.clone(), - aead_key: args.aead_key.clone(), - ssl_private_key: args.ssl_private_key.clone(), - })) - .await - .ok(); - - Json(BuildResponse { token: token }) -} - -#[derive(Debug, Deserialize)] -#[serde(crate = "rocket::serde")] -struct StatusArgs { - token: String, -} - -#[post("/status", data = "")] -async fn status(state: &State, 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(), - }; - println!("{}", result); - result -} - -#[get("/download?")] -async fn download(state: &State, token: String) -> Result { - println!("download with token={}", token); - - { - // 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, - }; - if !file_ready { - return Err(Status::NotFound); - } +fn get_llvm_install_path() -> Result, Box> { + let p = + std::env::var("LLVM_INSTALL").map_err(|_| "No LLVM_INSTALL environment variable set")?; + let p = Box::new(PathBuf::from(p)); + if !Path::exists(&*p) { + return Err("LLVM_INSTALL is set however the location is non-existant".into()); } - // load the content into memory, so that we can delete file - let path = Path::join(Path::new("out"), format!("ssm-{}", token)); - let content = fs::read(path).await.map_err(|_| Status::NotFound)?; + // check if go compiler is working + use std::process::{Command, Stdio}; + let go = Command::new(p.join("bin/go")) + .env("LD_LIBRARY_PATH", p.join("lib64")) + .args(["version"]) + .stdout(Stdio::null()) + .status() + .map_err(|_| "Testing go compiler error, non-existant file LLVM_INSTALL/bin/go")?; - 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, -) { - while let Some(request) = receiver.recv().await { - println!("requesting {:?}", request); - process_request(request, token_list.clone(), builder_sender.clone()).await; + if !go.success() { + return Err("go compiler can't give version".into()); } -} -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)) - .await - .ok(); - - let build_result = perform_build(&args).await.ok(); - match build_result { - Some(_) => { - 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()), - )) - .await - .ok(); - } - }; - } + Ok(p) } #[tokio::main] async fn main() { + let llvm_install = get_llvm_install_path().expect("Cannot get llvm installation path due to"); + 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 token_list: types::TokenList = Arc::new(Mutex::new(HashMap::new())); - let worker = AppState { - sender: sender.clone(), - token_list: token_list.clone(), - }; - - tokio::spawn(setup_request_receiver( + // main request handler + tokio::spawn(request::setup_request_receiver( receiver, token_list.clone(), builder_sender.clone(), )); - tokio::spawn(setup_builder(builder_receiver, sender.clone())); + + // builder handler, process build request, called by request handler + tokio::spawn(builder::setup_builder( + builder_receiver, + sender.clone(), + builder::Builder::new(llvm_install), + )); rocket::build() - .manage(worker) - .mount("/", routes![build]) - .mount("/", routes![status]) - .mount("/", routes![download]) + .manage(types::AppState { + sender: sender.clone(), + token_list: token_list.clone(), + }) + .mount( + "/", + routes![ + routes::route_build, + routes::route_status, + routes::route_download + ], + ) .launch() .await .ok(); diff --git a/src/request.rs b/src/request.rs new file mode 100644 index 0000000..2755e6e --- /dev/null +++ b/src/request.rs @@ -0,0 +1,42 @@ +use tokio::fs; +use tokio::sync::mpsc; + +use crate::types::{Request, RequestBuildArgs, TokenList, TokenStatus}; + +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(); + if let Some(state) = token_list.get_mut(&token) { + *state = new_status; + } + } + 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(); + token_list.remove(&token); + } + } +} + +pub async fn setup_request_receiver( + mut receiver: mpsc::Receiver, + token_list: TokenList, + builder_sender: mpsc::Sender, +) { + while let Some(request) = receiver.recv().await { + process_request(request, token_list.clone(), builder_sender.clone()).await; + } +} diff --git a/src/routes/build.rs b/src/routes/build.rs new file mode 100644 index 0000000..1af17c9 --- /dev/null +++ b/src/routes/build.rs @@ -0,0 +1,47 @@ +use rocket::serde; + +use rocket::State; + +use serde::{json::Json, Deserialize, Serialize}; + +use uuid::Uuid; + +use crate::types::{AppState, Request, RequestBuildArgs}; + +#[derive(Debug, Deserialize)] +#[serde(crate = "rocket::serde")] +pub struct BuildArgs { + vendor_name: String, + vendor_metamask_account: String, + aead_key: String, + ssl_private_key: String, +} + +#[derive(Serialize)] +#[serde(crate = "rocket::serde")] +pub struct BuildResponse { + token: String, +} + +#[post("/build", data = "")] +pub async fn route(state: &State, args: Json) -> Json { + let args = &*args; + println!("build with args={:?}", args); + + let sender = &state.sender; + let token = Uuid::new_v4().to_string(); + + sender.send(Request::InitToken(token.clone())).await.ok(); + sender + .send(Request::Build(RequestBuildArgs { + token: token.clone(), + vendor_name: args.vendor_name.clone(), + vendor_metamask_account: args.vendor_metamask_account.clone(), + aead_key: args.aead_key.clone(), + ssl_private_key: args.ssl_private_key.clone(), + })) + .await + .ok(); + + Json(BuildResponse { token }) +} diff --git a/src/routes/download.rs b/src/routes/download.rs new file mode 100644 index 0000000..4110fa8 --- /dev/null +++ b/src/routes/download.rs @@ -0,0 +1,33 @@ +use std::path::Path; + +use rocket::tokio; + +use rocket::http::Status; +use rocket::State; +use rocket_download_response::DownloadResponse; + +use tokio::fs; + +use crate::types::{AppState, Request, TokenStatus}; + +#[get("/download?")] +pub async fn route(state: &State, token: String) -> Result { + { + // 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, + }; + 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 = fs::read(path).await.map_err(|_| Status::NotFound)?; + + state.sender.send(Request::Delete(token)).await.ok(); + Ok(DownloadResponse::from_vec(content, Some("ssm"), None)) +} diff --git a/src/routes/mod.rs b/src/routes/mod.rs new file mode 100644 index 0000000..5a09571 --- /dev/null +++ b/src/routes/mod.rs @@ -0,0 +1,7 @@ +mod build; +mod download; +mod status; + +pub use build::route as route_build; +pub use download::route as route_download; +pub use status::route as route_status; diff --git a/src/routes/status.rs b/src/routes/status.rs new file mode 100644 index 0000000..548be5a --- /dev/null +++ b/src/routes/status.rs @@ -0,0 +1,39 @@ +use rocket::serde; + +use rocket::State; + +use serde::{json::Json, Deserialize, Serialize}; + +use crate::types::AppState; + +#[derive(Debug, Deserialize)] +#[serde(crate = "rocket::serde")] +pub struct StatusArgs { + token: String, +} + +#[derive(Serialize)] +#[serde(crate = "rocket::serde")] +pub struct StatusResponse { + token: String, + status: Option, + error: Option, +} + +#[post("/status", data = "")] +pub async fn route(state: &State, args: Json) -> Json { + let args = &*args; + let token_list = state.token_list.lock().unwrap(); + match token_list.get(&args.token) { + Some(state) => Json(StatusResponse { + token: args.token.clone(), + status: Some(state.into()), + error: None, + }), + _ => Json(StatusResponse { + token: args.token.clone(), + status: None, + error: Some("invalid token".into()), + }), + } +} diff --git a/src/types.rs b/src/types.rs new file mode 100644 index 0000000..0d2827d --- /dev/null +++ b/src/types.rs @@ -0,0 +1,52 @@ +use std::collections::HashMap; +use std::sync::{Arc, Mutex}; + +use rocket::tokio; + +use tokio::sync::mpsc; + +#[derive(Debug)] +pub enum Request { + InitToken(String), + SetToken(String, TokenStatus), + Build(RequestBuildArgs), + Delete(String), +} + +#[derive(Clone, Debug)] +pub struct RequestBuildArgs { + pub token: String, + pub vendor_name: String, + pub vendor_metamask_account: String, + pub aead_key: String, + pub ssl_private_key: String, +} + +#[derive(Debug)] +pub enum TokenStatus { + // new token is assigned and waiting in queue + Initialized, + Building, + Finished, + // if the build can't be complete + // Error(msg) + Error(String), +} + +impl Into for &TokenStatus { + fn into(self) -> String { + match self { + TokenStatus::Initialized => "Initialized".into(), + TokenStatus::Building => "Building".into(), + TokenStatus::Finished => "Finished".into(), + TokenStatus::Error(err) => format!("Error: {}", &err), + } + } +} + +pub type TokenList = Arc>>; + +pub struct AppState { + pub sender: mpsc::Sender, + pub token_list: TokenList, +}