Added process spawning & improved code
This commit is contained in:
parent
608bef56c1
commit
c3bbc6e541
44
conf.rs
Normal file
44
conf.rs
Normal file
|
@ -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<ProxyConf>
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
pub struct ProxyConf {
|
||||
pub host: String,
|
||||
pub target: String,
|
||||
pub socket: Option<bool>,
|
||||
pub spawn: Option<SpawnConf>
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
pub struct SpawnConf {
|
||||
pub command: String,
|
||||
pub args: Option<Vec<String>>,
|
||||
pub envs: Option<Vec<(String, String)>>
|
||||
}
|
||||
|
||||
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);}
|
||||
}
|
||||
}
|
|
@ -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"
|
||||
|
|
|
@ -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<String, usize> = generate_host_table();
|
||||
}
|
||||
|
||||
|
||||
#[derive(Debug, Deserialize, Serialize)]
|
||||
pub struct Root {
|
||||
pub proxy: Vec<ProxyItem>
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize, Serialize)]
|
||||
pub struct ProxyItem {
|
||||
pub host: String,
|
||||
pub target: String,
|
||||
pub socket: Option<bool>
|
||||
}
|
||||
|
||||
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<String, usize> {
|
||||
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)
|
||||
}
|
58
data.rs
Normal file
58
data.rs
Normal file
|
@ -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<String, usize> = generate_host_map();
|
||||
pub static ref SERVICES: Arc<Mutex<Vec<ServiceData>>> = Arc::new(Mutex::new(vec![]));
|
||||
}
|
||||
|
||||
pub struct ServiceData {
|
||||
pub child: Option<Child>,
|
||||
pub running: bool
|
||||
}
|
||||
|
||||
impl ServiceData {
|
||||
pub fn new(child: Option<Child>) -> 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<String, usize> {
|
||||
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)
|
||||
}
|
92
main.rs
92
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<Body>) -> Result<Response<Body>, hyper::Error> {
|
||||
use crate::conf::CONFIG;
|
||||
|
||||
async fn run(req: Request<Body>) -> Result<Response<Body>, 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<Body>) -> Result<Response<Body>, 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);
|
||||
}
|
||||
|
|
Loading…
Reference in a new issue