Add support for multiple executables

This commit is contained in:
2024-12-21 23:55:33 -06:00
parent 042e2cf0d6
commit ccba5a6149

View File

@@ -1,3 +1,7 @@
use std::collections::HashMap;
use std::iter::Map;
use std::path;
use salvo::http::HeaderValue; use salvo::http::HeaderValue;
use salvo::logging::Logger; use salvo::logging::Logger;
@@ -52,7 +56,7 @@ async fn hello() -> &'static str {
#[handler] #[handler]
async fn download(depot: &mut Depot, req: &mut Request, res: &mut Response) { async fn download(depot: &mut Depot, req: &mut Request, res: &mut Response) {
let state = depot.obtain::<State>().unwrap(); let state = depot.obtain::<State>().unwrap();
let data = state.with_key(b"test"); // Clone the data let data = state.exe_with_key("windows", b"test"); // Clone the data
if let Err(e) = res.write_body(data) { if let Err(e) = res.write_body(data) {
eprintln!("Error writing body: {}", e); eprintln!("Error writing body: {}", e);
@@ -60,34 +64,62 @@ async fn download(depot: &mut Depot, req: &mut Request, res: &mut Response) {
res.headers.insert( res.headers.insert(
"Content-Disposition", "Content-Disposition",
HeaderValue::from_static("attachment; filename=demo"), HeaderValue::from_str(format!("attachment; filename=\"{}\"", "demo.exe").as_str()).unwrap(),
); );
res.headers.insert( res.headers.insert(
"Content-Type", "Content-Type",
HeaderValue::from_static("application/octet-stream"), HeaderValue::from_static("application/octet-stream"),
); );
} }
#[allow(dead_code)]
#[derive(Default, Clone, Debug)] #[derive(Default, Clone, Debug)]
struct State { struct State {
executables: HashMap<String, Executable>,
}
#[derive(Default, Clone, Debug)]
struct Executable {
data: Vec<u8>, data: Vec<u8>,
name: String,
key_start: usize, key_start: usize,
key_end: usize, key_end: usize,
} }
impl State { impl State {
fn with_key(&self, new_key: &[u8]) -> Vec<u8> { fn add_executable(&mut self, exe_type: String, exe_path: String) {
let mut data = self.data.clone(); let data = std::fs::read(&exe_path).expect("Unable to read file");
let pattern = "a".repeat(1024);
let key_start = search(&data, pattern.as_bytes(), 0).unwrap();
let key_end = key_start + pattern.len();
let filename = path::Path::new(&exe_path)
.file_name()
.unwrap()
.to_string_lossy()
.into_owned();
let exe = Executable {
data,
name: filename,
key_start: key_start,
key_end: key_end,
};
self.executables.insert(exe_type, exe);
}
fn exe_with_key(&self, executable_type: &str, new_key: &[u8]) -> Vec<u8> {
let exe = self.executables.get(executable_type).unwrap();
let mut data = exe.data.clone();
// Copy the key into the data // Copy the key into the data
for i in 0..new_key.len() { for i in 0..new_key.len() {
data[self.key_start + i] = new_key[i]; data[exe.key_start + i] = new_key[i];
} }
// If the new key is shorter than the old key, we just write over the remaining data // If the new key is shorter than the old key, we just write over the remaining data
if new_key.len() < self.key_end - self.key_start { if new_key.len() < exe.key_end - exe.key_start {
for i in self.key_start + new_key.len()..self.key_end { for i in exe.key_start + new_key.len()..exe.key_end {
data[i] = b' '; data[i] = b' ';
} }
} }
@@ -102,19 +134,19 @@ async fn main() {
let addr = format!("0.0.0.0:{}", port); let addr = format!("0.0.0.0:{}", port);
tracing_subscriber::fmt().init(); tracing_subscriber::fmt().init();
let path = "./demo/target/release/demo"; let mut state = State {
let data = std::fs::read(&path).expect("Unable to read file"); executables: HashMap::new(),
let pattern = "a".repeat(1024);
let key_start = search(&data, pattern.as_bytes(), 0).unwrap();
let key_end = key_start + pattern.len();
let state = State {
data,
key_start: key_start,
key_end: key_end,
}; };
state.add_executable(
"windows".to_string(),
"./demo/target/x86_64-pc-windows-gnu/release/demo.exe".to_string(),
);
state.add_executable(
"linux".to_string(),
"./demo/target/release/demo".to_string(),
);
let router = Router::new() let router = Router::new()
.hoop(affix_state::inject(state)) .hoop(affix_state::inject(state))
.get(hello) .get(hello)