split multiple files

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<PathBuf>,
+}
+
+impl Builder {
+    pub fn new(llvm_install: Box<PathBuf>) -> Self {
+        Builder { llvm_install }
+    }
+
+    async fn build(&self, args: &RequestBuildArgs) -> Result<(), Box<dyn Error>> {
+        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<RequestBuildArgs>,
+    sender: mpsc::Sender<Request>,
+    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();
+            }
+        };
+    }
+}

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)]
+fn get_llvm_install_path() -> Result<Box<PathBuf>, Box<dyn Error>> {
-struct RequestBuildArgs {
+    let p =
-    token: String,
+        std::env::var("LLVM_INSTALL").map_err(|_| "No LLVM_INSTALL environment variable set")?;
-    vendor_name: String,
+    let p = Box::new(PathBuf::from(p));
-    vendor_metamask_account: String,
+    if !Path::exists(&*p) {
-    aead_key: String,
+        return Err("LLVM_INSTALL is set however the location is non-existant".into());
-    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<dyn Error>> {
-    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<RequestBuildArgs>,
-) {
-    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<Mutex<HashMap<String, TokenStatus>>>;
-
-struct AppState {
-    pub sender: mpsc::Sender<Request>,
-    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 = "<args>")]
-async fn build(state: &State<AppState>, args: Json<BuildArgs>) -> Json<BuildResponse> {
-    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 = "<args>")]
-async fn status(state: &State<AppState>, args: Json<StatusArgs>) -> 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?<token>")]
-async fn download(state: &State<AppState>, token: String) -> Result<DownloadResponse, Status> {
-    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);
-        }
     }
 
-    // load the content into memory, so that we can delete file
+    // check if go compiler is working
-    let path = Path::join(Path::new("out"), format!("ssm-{}", token));
+    use std::process::{Command, Stdio};
-    let content = fs::read(path).await.map_err(|_| Status::NotFound)?;
+    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();
+    if !go.success() {
-    Ok(DownloadResponse::from_vec(content, Some("ssm"), None))
+        return Err("go compiler can't give version".into());
-}
-
-async fn setup_request_receiver(
-    mut receiver: mpsc::Receiver<Request>,
-    token_list: TokenList,
-    builder_sender: mpsc::Sender<RequestBuildArgs>,
-) {
-    while let Some(request) = receiver.recv().await {
-        println!("requesting {:?}", request);
-        process_request(request, token_list.clone(), builder_sender.clone()).await;
     }
-}
 
-async fn setup_builder(
+    Ok(p)
-    mut receiver: mpsc::Receiver<RequestBuildArgs>,
-    sender: mpsc::Sender<Request>,
-) {
-    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();
-            }
-        };
-    }
 }
 
 #[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 {
+    // main request handler
-        sender: sender.clone(),
+    tokio::spawn(request::setup_request_receiver(
-        token_list: token_list.clone(),
-    };
-
-    tokio::spawn(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)
+        .manage(types::AppState {
-        .mount("/", routes![build])
+            sender: sender.clone(),
-        .mount("/", routes![status])
+            token_list: token_list.clone(),
-        .mount("/", routes![download])
+        })
+        .mount(
+            "/",
+            routes![
+                routes::route_build,
+                routes::route_status,
+                routes::route_download
+            ],
+        )
         .launch()
         .await
         .ok();

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<RequestBuildArgs>,
+) {
+    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<Request>,
+    token_list: TokenList,
+    builder_sender: mpsc::Sender<RequestBuildArgs>,
+) {
+    while let Some(request) = receiver.recv().await {
+        process_request(request, token_list.clone(), builder_sender.clone()).await;
+    }
+}

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 = "<args>")]
+pub async fn route(state: &State<AppState>, args: Json<BuildArgs>) -> Json<BuildResponse> {
+    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 })
+}

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?<token>")]
+pub async fn route(state: &State<AppState>, token: String) -> Result<DownloadResponse, Status> {
+    {
+        // 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))
+}

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;

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<String>,
+    error: Option<String>,
+}
+
+#[post("/status", data = "<args>")]
+pub async fn route(state: &State<AppState>, args: Json<StatusArgs>) -> Json<StatusResponse> {
+    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()),
+        }),
+    }
+}

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<String> 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<Mutex<HashMap<String, TokenStatus>>>;
+
+pub struct AppState {
+    pub sender: mpsc::Sender<Request>,
+    pub token_list: TokenList,
+}