poc working

This commit is contained in:
nganhkhoa 2022-05-05 23:34:14 +07:00
commit f67267b151
4 changed files with 2164 additions and 0 deletions

1
.gitignore vendored Normal file
View File

@ -0,0 +1 @@
/target

1871
Cargo.lock generated Normal file

File diff suppressed because it is too large Load Diff

18
Cargo.toml Normal file
View File

@ -0,0 +1,18 @@
[package]
name = "ssm-builder"
version = "0.1.0"
edition = "2021"
[dependencies]
rocket = {version = "0.5.0-rc.1", features = ["json"]}
rocket-download-response = "0.5.1"
async-std = "1.11.0"
[dependencies.tokio]
features = ["fs", "process"]
[dependencies.uuid]
version = "1.0.0"
features = [
"v4",
]

274
src/main.rs Normal file
View File

@ -0,0 +1,274 @@
#[macro_use]
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::State;
use rocket::http::Status;
use rocket_download_response::DownloadResponse;
use serde::{Serialize, Deserialize, json::Json};
use tokio::sync::mpsc;
use tokio::fs as async_fs;
use tokio::process::Command;
use uuid::Uuid;
#[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<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")
.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?;
async_fs::create_dir("out").await.ok();
async_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) => {
println!("delete file with token={}", 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>>>;
// type BuildQueue = Arc<Mutex<Vec<RequestBuildArgs>>>;
struct AppState {
pub sender: mpsc::Sender<Request>,
pub token_list: TokenList,
}
#[derive(Debug)]
#[derive(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)]
#[derive(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);
}
}
let path = Path::join(Path::new("out"), format!("ssm-{}", token));
let content = async_fs::read(path)
.await
.map_err(|_| Status::NotFound)?;
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<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(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 (sender, mut receiver) = mpsc::channel(100);
let (builder_sender, mut builder_receiver) = mpsc::channel(100);
let token_list: TokenList = Arc::new(Mutex::new(HashMap::new()));
let worker = AppState {
sender: sender.clone(),
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()));
let rocket_runner = rocket::build()
.manage(worker)
.mount("/", routes![build])
.mount("/", routes![status])
.mount("/", routes![download])
.launch()
.await;
}