nganhkhoa/ssm-builder
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

format code

Browse Source 
and add strip command
This commit is contained in:
nganhkhoa 2022-05-05 23:54:41 +07:00
parent 459754bd8c
commit 1f5a82ff91
1 changed files with 100 additions and 79 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

179
src/main.rs
View File

@ -2,35 +2,32 @@
extern crate rocket; extern crate rocket;
use std::collections::HashMap; use std::collections::HashMap;
use std::fs;
use std::error::Error; use std::error::Error;
use std::path::Path; use std::path::Path;
use std::sync::{Arc, Mutex}; use std::sync::{Arc, Mutex};
use std::time::Duration;
use std::thread;
use rocket::tokio as tokio; use rocket::serde;
use rocket::serde as serde; use rocket::tokio;
use rocket::State;
use rocket::http::Status; use rocket::http::Status;
use rocket::State;
use rocket_download_response::DownloadResponse; use rocket_download_response::DownloadResponse;
use serde::{Serialize, Deserialize, json::Json}; use serde::{json::Json, Deserialize, Serialize};
use tokio::sync::mpsc; use tokio::fs;
use tokio::fs as async_fs;
use tokio::process::Command; use tokio::process::Command;
use tokio::sync::mpsc;
use uuid::Uuid; use uuid::Uuid;
#[derive(Clone, Debug)] #[derive(Clone, Debug)]
struct RequestBuildArgs { struct RequestBuildArgs {
token: String, token: String,
vendor_name: String, vendor_name: String,
vendor_metamask_account: String, vendor_metamask_account: String,
aead_key: String, aead_key: String,
ssl_private_key: String ssl_private_key: String,
} }
#[derive(Debug)] #[derive(Debug)]
@ -42,9 +39,6 @@ enum Request {
} }
async fn perform_build(args: &RequestBuildArgs) -> Result<(), Box<dyn Error>> { async fn perform_build(args: &RequestBuildArgs) -> Result<(), Box<dyn Error>> {
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") Command::new("python3")
.current_dir("/home/r00t/work/ollvm/ssm-for-test/ssm/dynamic/") .current_dir("/home/r00t/work/ollvm/ssm-for-test/ssm/dynamic/")
.env("VENDOR_NAME", &args.vendor_name) .env("VENDOR_NAME", &args.vendor_name)
@ -58,36 +52,61 @@ async fn perform_build(args: &RequestBuildArgs) -> Result<(), Box<dyn Error>> {
Command::new("/home/r00t/work/ollvm/llvm-install/bin/go") Command::new("/home/r00t/work/ollvm/llvm-install/bin/go")
.current_dir("/home/r00t/work/ollvm/ssm-for-test/ssm/") .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([ .args([
"build", "-a", "build",
"-o", &format!("ssm-{}", args.token), "-a",
"-gccgoflags=\"-static-libgo -Wl,--version-script=ssm.version\"", "-o",
"-mllvm", "-obfuscation=gvo", &format!("ssm-{}", args.token),
"-mllvm", "-obfuscation=sub", "-gccgoflags",
"-mllvm", "-obfuscation=flatten", "-static-libgo -Wl,--version-script=ssm.version",
"-mllvm", "-obfuscation=idr-branch", "-mllvm",
"-obfuscation=gvo",
"-mllvm",
"-obfuscation=sub",
"-mllvm",
"-obfuscation=flatten",
"-mllvm",
"-obfuscation=idr-branch",
]) ])
.spawn()? .spawn()?
.wait() .wait()
.await?; .await?;
async_fs::create_dir("out").await.ok(); Command::new("strip")
async_fs::rename( .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!("/home/r00t/work/ollvm/ssm-for-test/ssm/ssm-{}", args.token),
format!("./out/ssm-{}", args.token)) format!("./out/ssm-{}", args.token),
.await )
.ok(); .await
.ok();
println!("build complete out/ssm-{}", args.token); println!("build complete out/ssm-{}", args.token);
Ok(()) Ok(())
} }
async fn process_request(request: Request, token_list: TokenList, build_sender: mpsc::Sender<RequestBuildArgs>) { async fn process_request(
request: Request,
token_list: TokenList,
build_sender: mpsc::Sender<RequestBuildArgs>,
) {
match request { match request {
Request::InitToken(token) => { Request::InitToken(token) => {
let mut token_list = token_list.lock().unwrap(); let mut token_list = token_list.lock().unwrap();
token_list.insert(token.clone(), TokenStatus::Initialized); token_list.insert(token.clone(), TokenStatus::Initialized);
}, }
Request::SetToken(token, new_status) => { Request::SetToken(token, new_status) => {
println!("set token {:?} {:?}", token, new_status); println!("set token {:?} {:?}", token, new_status);
let mut token_list = token_list.lock().unwrap(); 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) => { Request::Build(args) => {
build_sender.send(args).await.ok(); build_sender.send(args).await.ok();
}, }
Request::Delete(token) => { Request::Delete(token) => {
fs::remove_file(format!("./out/ssm-{}", token)).await.ok(); fs::remove_file(format!("./out/ssm-{}", token)).await.ok();
let mut token_list = token_list.lock().unwrap(); let mut token_list = token_list.lock().unwrap();
@ -118,21 +137,19 @@ enum TokenStatus {
} }
type TokenList = Arc<Mutex<HashMap<String, TokenStatus>>>; type TokenList = Arc<Mutex<HashMap<String, TokenStatus>>>;
// type BuildQueue = Arc<Mutex<Vec<RequestBuildArgs>>>;
struct AppState { struct AppState {
pub sender: mpsc::Sender<Request>, pub sender: mpsc::Sender<Request>,
pub token_list: TokenList, pub token_list: TokenList,
} }
#[derive(Debug)] #[derive(Debug, Deserialize)]
#[derive(Deserialize)]
#[serde(crate = "rocket::serde")] #[serde(crate = "rocket::serde")]
struct BuildArgs { struct BuildArgs {
vendor_name: String, vendor_name: String,
vendor_metamask_account: String, vendor_metamask_account: String,
aead_key: String, aead_key: String,
ssl_private_key: String ssl_private_key: String,
} }
#[derive(Serialize)] #[derive(Serialize)]
@ -141,7 +158,7 @@ struct BuildResponse {
token: String, token: String,
} }
#[post("/build", data="<args>")] #[post("/build", data = "<args>")]
async fn build(state: &State<AppState>, args: Json<BuildArgs>) -> Json<BuildResponse> { async fn build(state: &State<AppState>, args: Json<BuildArgs>) -> Json<BuildResponse> {
let args = &*args; let args = &*args;
println!("build with args={:?}", args); println!("build with args={:?}", args);
@ -149,46 +166,36 @@ async fn build(state: &State<AppState>, args: Json<BuildArgs>) -> Json<BuildResp
let sender = &state.sender; let sender = &state.sender;
let token = Uuid::new_v4().to_string(); let token = Uuid::new_v4().to_string();
sender.send(Request::InitToken(token.clone())) sender.send(Request::InitToken(token.clone())).await.ok();
.await sender
.ok(); .send(Request::Build(RequestBuildArgs {
sender.send(Request::Build(
RequestBuildArgs {
token: token.clone(), token: token.clone(),
vendor_name: args.vendor_name.clone(), vendor_name: args.vendor_name.clone(),
vendor_metamask_account: args.vendor_metamask_account.clone(), vendor_metamask_account: args.vendor_metamask_account.clone(),
aead_key: args.aead_key.clone(), aead_key: args.aead_key.clone(),
ssl_private_key: args.ssl_private_key.clone(), ssl_private_key: args.ssl_private_key.clone(),
} }))
))
.await .await
.ok(); .ok();
Json(BuildResponse { Json(BuildResponse { token: token })
token: token,
})
} }
#[derive(Debug)] #[derive(Debug, Deserialize)]
#[derive(Deserialize)]
#[serde(crate = "rocket::serde")] #[serde(crate = "rocket::serde")]
struct StatusArgs { struct StatusArgs {
token: String token: String,
} }
#[post("/status", data="<args>")] #[post("/status", data = "<args>")]
async fn status(state: &State<AppState>, args: Json<StatusArgs>) -> String { async fn status(state: &State<AppState>, args: Json<StatusArgs>) -> String {
let args = &*args; let args = &*args;
println!("args: {:?}", args); println!("args: {:?}", args);
let token_list = state.token_list.lock().unwrap(); let token_list = state.token_list.lock().unwrap();
let result = match token_list.get(&args.token) { let result = match token_list.get(&args.token) {
Some(state) => { Some(state) => format!("{:?}", state).to_string(),
format!("{:?}", state).to_string() _ => "token non exist".to_string(),
},
_ => {
"token non exist".to_string()
}
}; };
println!("{}", result); println!("{}", result);
result result
@ -198,51 +205,62 @@ async fn status(state: &State<AppState>, args: Json<StatusArgs>) -> String {
async fn download(state: &State<AppState>, token: String) -> Result<DownloadResponse, Status> { async fn download(state: &State<AppState>, token: String) -> Result<DownloadResponse, Status> {
println!("download with token={}", token); 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 token_list = state.token_list.lock().unwrap();
let file_ready = match token_list.get(&token) { let file_ready = match token_list.get(&token) {
Some(TokenStatus::Finished) => true, Some(TokenStatus::Finished) => true,
_ => false _ => false,
}; };
if !file_ready { if !file_ready {
return Err(Status::NotFound); 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 path = Path::join(Path::new("out"), format!("ssm-{}", token));
let content = async_fs::read(path) let content = fs::read(path).await.map_err(|_| Status::NotFound)?;
.await
.map_err(|_| Status::NotFound)?;
state.sender.send(Request::Delete(token)) state.sender.send(Request::Delete(token)).await.ok();
.await
.ok();
Ok(DownloadResponse::from_vec(content, Some("ssm"), None)) Ok(DownloadResponse::from_vec(content, Some("ssm"), None))
} }
async fn setup_request_receiver(mut receiver: mpsc::Receiver<Request>, token_list: TokenList, builder_sender: mpsc::Sender<RequestBuildArgs>) { 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 { while let Some(request) = receiver.recv().await {
println!("requesting {:?}", request); println!("requesting {:?}", request);
process_request(request, token_list.clone(), builder_sender.clone()).await; process_request(request, token_list.clone(), builder_sender.clone()).await;
} }
} }
async fn setup_builder(mut receiver: mpsc::Receiver<RequestBuildArgs>, sender: mpsc::Sender<Request>) { async fn setup_builder(
mut receiver: mpsc::Receiver<RequestBuildArgs>,
sender: mpsc::Sender<Request>,
) {
while let Some(args) = receiver.recv().await { while let Some(args) = receiver.recv().await {
println!("requesting build {:?}", args); println!("requesting build {:?}", args);
sender.send(Request::SetToken(args.token.clone(), TokenStatus::Building)) sender
.send(Request::SetToken(args.token.clone(), TokenStatus::Building))
.await .await
.ok(); .ok();
let build_result = perform_build(&args).await.ok(); let build_result = perform_build(&args).await.ok();
match build_result { match build_result {
Some(_) => { Some(_) => {
sender.send(Request::SetToken(args.token.clone(), TokenStatus::Finished)) sender
.send(Request::SetToken(args.token.clone(), TokenStatus::Finished))
.await .await
.ok(); .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 .await
.ok(); .ok();
} }
@ -252,8 +270,8 @@ async fn setup_builder(mut receiver: mpsc::Receiver<RequestBuildArgs>, sender: m
#[tokio::main] #[tokio::main]
async fn main() { async fn main() {
let (sender, mut receiver) = mpsc::channel(100); let (sender, receiver) = mpsc::channel(100);
let (builder_sender, mut builder_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: TokenList = Arc::new(Mutex::new(HashMap::new()));
let worker = AppState { let worker = AppState {
@ -261,16 +279,19 @@ async fn main() {
token_list: token_list.clone(), token_list: token_list.clone(),
}; };
let receiver_runner = tokio::spawn( tokio::spawn(setup_request_receiver(
setup_request_receiver(receiver, token_list.clone(), builder_sender.clone())); receiver,
let builder_runner = tokio::spawn( token_list.clone(),
setup_builder(builder_receiver, sender.clone())); builder_sender.clone(),
));
tokio::spawn(setup_builder(builder_receiver, sender.clone()));
let rocket_runner = rocket::build() rocket::build()
.manage(worker) .manage(worker)
.mount("/", routes![build]) .mount("/", routes![build])
.mount("/", routes![status]) .mount("/", routes![status])
.mount("/", routes![download]) .mount("/", routes![download])
.launch() .launch()
.await; .await
.ok();
} }