From c3bbc6e541efa57edee8352ecad6432262562498 Mon Sep 17 00:00:00 2001 From: Filip Znachor Date: Fri, 21 Apr 2023 00:04:58 +0200 Subject: [PATCH] Added process spawning & improved code --- conf.rs | 44 +++++++++++++++++++++++ config.toml | 4 +++ configuration.rs | 64 --------------------------------- data.rs | 58 ++++++++++++++++++++++++++++++ main.rs | 92 ++++++++++++++++++++++++++++++++++++++++++------ 5 files changed, 187 insertions(+), 75 deletions(-) create mode 100644 conf.rs delete mode 100644 configuration.rs create mode 100644 data.rs diff --git a/conf.rs b/conf.rs new file mode 100644 index 0000000..e671928 --- /dev/null +++ b/conf.rs @@ -0,0 +1,44 @@ +use std::{fs::File, process::exit}; +use std::io::prelude::*; +use toml::{de::from_str}; +use serde::Deserialize; +use lazy_static::lazy_static; + +lazy_static! { + pub static ref CONFIG: RootConf = load_config(); +} + +#[derive(Debug, Deserialize)] +pub struct RootConf { + pub proxy: Vec +} + +#[derive(Debug, Deserialize)] +pub struct ProxyConf { + pub host: String, + pub target: String, + pub socket: Option, + pub spawn: Option +} + +#[derive(Debug, Deserialize)] +pub struct SpawnConf { + pub command: String, + pub args: Option>, + pub envs: Option> +} + +fn load_config() -> RootConf { + let file = File::open("config.toml"); + if file.is_err() { + println!("[!] Unable to read config file"); exit(-1); + } + let mut contents = String::new(); + if file.unwrap().read_to_string(&mut contents).is_err() { + println!("[!] Unable to read config file"); exit(-1); + } + match from_str(&contents) { + Ok(conf) => conf, + Err(_) => {println!("[!] Unable to parse config"); exit(0);} + } +} \ No newline at end of file diff --git a/config.toml b/config.toml index 657ce9e..5a4992f 100644 --- a/config.toml +++ b/config.toml @@ -2,6 +2,10 @@ host = "website.local.gd" socket = true target = "./www.sock" +[proxy.spawn] +command = "/usr/bin/node" +args = [ "./webserver/index.mjs" ] +envs = [ ["PORT", "www.sock"] ] [[proxy]] host = "git.local.gd" diff --git a/configuration.rs b/configuration.rs deleted file mode 100644 index fab3f47..0000000 --- a/configuration.rs +++ /dev/null @@ -1,64 +0,0 @@ -use std::{fs::File, process::exit}; -use std::io::prelude::*; -use toml::{de::from_str}; -use serde::{Deserialize, Serialize}; -use lazy_static::lazy_static; - -use std::collections::HashMap; -use hyper::http::HeaderValue; - -lazy_static! { - pub static ref CONFIG: Root = load_config(); - pub static ref HOST_MAP: HashMap = generate_host_table(); -} - - -#[derive(Debug, Deserialize, Serialize)] -pub struct Root { - pub proxy: Vec -} - -#[derive(Debug, Deserialize, Serialize)] -pub struct ProxyItem { - pub host: String, - pub target: String, - pub socket: Option -} - -fn load_config() -> Root { - let file = File::open("config.toml"); - if file.is_err() { - println!("[!] Unable to read config file"); exit(-1); - } - let mut contents = String::new(); - if file.unwrap().read_to_string(&mut contents).is_err() { - println!("[!] Unable to read config file"); exit(-1); - } - match from_str(&contents) { - Ok(conf) => {conf}, - Err(_) => {println!("[!] Unable to parse config"); exit(0);} - } -} - -pub fn get_host(host: Option<&HeaderValue>) -> Option<&ProxyItem> { - if host.is_some() { - let host_parts: Vec<&str> = host.unwrap().to_str().unwrap().split(":").collect(); - let domain = host_parts.get(0); - if domain.is_some() { - let host_index = HOST_MAP.get(&domain.unwrap().to_string()); - if host_index.is_some() { - let res = CONFIG.proxy.get(host_index.unwrap().clone()); - return res; - } - } - } - return None; -} - -pub fn generate_host_table() -> HashMap { - let mut hosts: Vec<(String, usize)> = vec![]; - for (i, proxy) in CONFIG.proxy.iter().enumerate() { - hosts.push((proxy.host.to_string(), i)); - } - HashMap::from_iter(hosts) -} \ No newline at end of file diff --git a/data.rs b/data.rs new file mode 100644 index 0000000..37ad11b --- /dev/null +++ b/data.rs @@ -0,0 +1,58 @@ +use std::{process::Child, sync::{Arc, Mutex}}; +use lazy_static::lazy_static; + +use std::collections::HashMap; +use hyper::http::HeaderValue; + +use crate::{conf::{CONFIG, ProxyConf}}; + +lazy_static! { + pub static ref HOST_MAP: HashMap = generate_host_map(); + pub static ref SERVICES: Arc>> = Arc::new(Mutex::new(vec![])); +} + +pub struct ServiceData { + pub child: Option, + pub running: bool +} + +impl ServiceData { + pub fn new(child: Option) -> ServiceData { + ServiceData { + child, + running: false + } + } + pub fn set_running(&mut self, running: bool) { + self.running = running; + } +} + +pub fn get_proxy(host_index: Option<&usize>) -> Option<&ProxyConf> { + match host_index { + Some(i) => CONFIG.proxy.get(i.clone()), + None => None + } +} + +pub fn get_proxy_index(host: Option<&HeaderValue>) -> Option<&usize> { + match host { + Some(host) => { + let host_parts: Vec<&str> = host.to_str().unwrap().split(":").collect(); + let domain = host_parts.get(0); + match domain { + Some(domain) => HOST_MAP.get(&domain.to_string()), + None => None + } + }, + None => None + } +} + +pub fn generate_host_map() -> HashMap { + let mut hosts: Vec<(String, usize)> = vec![]; + for (i, proxy) in CONFIG.proxy.iter().enumerate() { + hosts.push((proxy.host.to_string(), i)); + } + HashMap::from_iter(hosts) +} \ No newline at end of file diff --git a/main.rs b/main.rs index 92e9d76..53a6d59 100644 --- a/main.rs +++ b/main.rs @@ -1,23 +1,27 @@ -use std::{alloc::System, str::FromStr}; +mod conf; +mod data; -#[global_allocator] -static A: System = System; - -mod configuration; - -use std::net::SocketAddr; +use std::{net::SocketAddr, str::FromStr, process::Command, path::Path, time::Duration}; +use conf::{ProxyConf, SpawnConf}; +use data::{HOST_MAP, SERVICES, ServiceData}; use hyperlocal::{UnixClientExt}; +use tokio::{fs, time::sleep}; use tower::make::Shared; use hyper::{service::service_fn, Body, Client, Request, Response, Server}; -async fn log(req: Request) -> Result, hyper::Error> { +use crate::conf::CONFIG; + +async fn run(req: Request) -> Result, hyper::Error> { let host = req.headers().get("host"); - let p = configuration::get_host(host); - match p { + let host_index = data::get_proxy_index(host); + let proxy = data::get_proxy(host_index); + match proxy { Some(p) => { + check_service(host_index.unwrap().clone(),p).await; + // Create new Request let mut request_builder = Request::builder().method(req.method()); let path = req.uri().path_and_query().unwrap().as_str(); @@ -54,13 +58,79 @@ async fn log(req: Request) -> Result, hyper::Error> { } +fn set_service_running(index: usize) { + SERVICES.lock().unwrap().get_mut(index).unwrap().set_running(true); +} + +fn is_service_running(index: usize) -> bool { + SERVICES.lock().unwrap().get(index).unwrap().running +} + +async fn check_service(index: usize, proxy: &ProxyConf) { + + match &proxy.spawn { + Some(spawn) => { + spawn_service(index, spawn); + if !is_service_running(index) { + wait_for_service(proxy).await; + set_service_running(index); + } + }, + None => {} + } + +} + +fn spawn_service(index: usize, spawn: &SpawnConf) -> bool { + match SERVICES.lock() { + Ok(mut array) => { + if array.get(index).is_none() { + let command = spawn.command.clone(); + let args = spawn.args.clone().unwrap_or(vec![]); + let envs = spawn.envs.clone().unwrap_or(vec![]); + let spawned_child = Command::new(command).args(args).envs(envs).spawn(); + match spawned_child { + Ok(child) => { + array.insert(index, ServiceData::new(Some(child))); + return true; + }, + Err(_) => println!("Error while spawning process!") + } + } + }, + Err(_) => {} + } + return false; +} + +async fn wait_for_service(proxy: &ProxyConf) { + let path = Path::new(&proxy.target); + while !path.exists() { + sleep(Duration::from_millis(100)).await; + } +} + #[tokio::main] async fn main() { - let make_service = Shared::new(service_fn(log)); + + for proxy in CONFIG.proxy.iter() { + if proxy.socket.unwrap_or(false) { + let path = Path::new(&proxy.target); + if path.exists() { + fs::remove_file(path).await.unwrap(); + } + } + } + + let make_service = Shared::new(service_fn(run)); let addr = SocketAddr::from(([127, 0, 0, 1], 3000)); let server = Server::bind(&addr).serve(make_service); + let host_count = HOST_MAP.len(); + let service_count = CONFIG.proxy.len(); + println!("odproxy is running with {} hosts and {} services", host_count, service_count); + if let Err(e) = server.await { println!("error: {}", e); }