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

split multiple files

Browse Source
This commit is contained in:
nganhkhoa 2022-05-06 12:40:03 +07:00
parent 1f5a82ff91
commit 832555cbad
8 changed files with 372 additions and 264 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

106
src/builder.rs Normal file
View File

@ -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();
}
};
}
}

312
src/main.rs
View File

@ -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,
fn get_llvm_install_path() -> Result<Box<PathBuf>, Box<dyn Error>> {
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());
}
#[derive(Debug)]
enum Request {
InitToken(String),
SetToken(String, TokenStatus),
Build(RequestBuildArgs),
Delete(String),
// 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")?;
if !go.success() {
return Err("go compiler can't give version".into());
}
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
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))
}
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();
}
};
}
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();

42
src/request.rs Normal file
View File

@ -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;
}
}

47
src/routes/build.rs Normal file
View File

@ -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 })
}

33
src/routes/download.rs Normal file
View File

@ -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))
}

7
src/routes/mod.rs Normal file
View File

@ -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;

39
src/routes/status.rs Normal file
View File

@ -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()),
}),
}
}

52
src/types.rs Normal file
View File

@ -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,
}