rust-rpxy/src/main.rs
2022-07-05 19:30:28 +09:00

115 lines
3 KiB
Rust

#[cfg(not(target_env = "msvc"))]
#[global_allocator]
static GLOBAL: tikv_jemallocator::Jemalloc = tikv_jemallocator::Jemalloc;
mod backend;
mod config;
mod constants;
mod error;
mod globals;
mod log;
mod proxy;
use crate::{
backend::{Backend, Backends},
config::parse_opts,
constants::*,
error::*,
globals::*,
log::*,
proxy::Proxy,
};
use futures::future::select_all;
use hyper::Client;
// use hyper_trust_dns::TrustDnsResolver;
use rustc_hash::FxHashMap as HashMap;
use std::{io::Write, sync::Arc};
use tokio::time::Duration;
fn main() {
// env::set_var("RUST_LOG", "info");
env_logger::Builder::from_env(env_logger::Env::default().default_filter_or("info"))
.format(|buf, record| {
let ts = buf.timestamp();
writeln!(
buf,
"{} [{}] {}",
ts,
record.level(),
// record.target(),
record.args(),
// record.file().unwrap_or("unknown"),
// record.line().unwrap_or(0),
)
})
.init();
info!("Start http (reverse) proxy");
let mut runtime_builder = tokio::runtime::Builder::new_multi_thread();
runtime_builder.enable_all();
runtime_builder.thread_name("rpxy");
let runtime = runtime_builder.build().unwrap();
runtime.block_on(async {
let mut globals = Globals {
listen_sockets: Vec::new(),
http_port: None,
https_port: None,
http3: false,
timeout: Duration::from_secs(TIMEOUT_SEC),
max_clients: MAX_CLIENTS,
clients_count: Default::default(),
max_concurrent_streams: MAX_CONCURRENT_STREAMS,
keepalive: true,
runtime_handle: runtime.handle().clone(),
};
let mut backends = Backends {
default_app: None,
apps: HashMap::<String, Backend>::default(),
};
let _ = parse_opts(&mut globals, &mut backends).expect("Invalid configuration");
entrypoint(Arc::new(globals), Arc::new(backends))
.await
.unwrap()
});
warn!("Exit the program");
}
// entrypoint creates and spawns tasks of proxy services
async fn entrypoint(globals: Arc<Globals>, backends: Arc<Backends>) -> Result<()> {
// let connector = TrustDnsResolver::default().into_rustls_webpki_https_connector();
let connector = hyper_rustls::HttpsConnectorBuilder::new()
.with_webpki_roots()
.https_or_http()
.enable_http1()
.enable_http2()
.build();
let forwarder = Arc::new(Client::builder().build::<_, hyper::Body>(connector));
let addresses = globals.listen_sockets.clone();
let futures = select_all(addresses.into_iter().map(|addr| {
let mut tls_enabled = false;
if let Some(https_port) = globals.https_port {
tls_enabled = https_port == (addr.port() as u16)
}
let proxy = Proxy {
globals: globals.clone(),
listening_on: addr,
tls_enabled,
backends: backends.clone(),
forwarder: forwarder.clone(),
};
globals.runtime_handle.spawn(proxy.start())
}));
// wait for all future
if let (Ok(_), _, _) = futures.await {
error!("Some proxy services are down");
};
Ok(())
}