use futures_util::future::try_join;
use std::fs;
use std::{convert::Infallible, net::SocketAddr};
use hyper::service::{make_service_fn, service_fn};
use hyper::{Body, Client, Method, Request, Response, Server};
use anyhow::anyhow;
use h2sr::Pattern;
use serde::Deserialize;
use std::borrow::Cow;
use std::path::Path;
use tokio::net::tcp::{ReadHalf, WriteHalf};
use tokio::net::{lookup_host, TcpStream};
type HttpClient = Client<hyper::client::HttpConnector>;
#[derive(Clone, Copy)]
pub struct Socks5Proxy(SocketAddr);
impl Socks5Proxy {
pub async fn connect(&self, auth: &str, req: Request<Body>) {
match tokio_socks::tcp::Socks5Stream::connect(self.0, &*auth).await {
Ok(mut stream) => tunnel(auth, req, stream.split()).await,
Err(e) => eprintln!("{}: socks5 connect error: {:?}", auth, e),
}
}
}
pub struct Env {
pub listen: SocketAddr,
pub to: Socks5Proxy,
pub proxy: Pattern,
pub block: Pattern,
}
#[derive(Deserialize)]
pub struct Config {
pub listen: SocketAddr,
pub to: SocketAddr,
pub proxy: Vec<String>,
pub block: Vec<String>,
}
async fn load_env() -> Env {
let matches = clap::App::new("Http socks router")
.author("Jason D.H. Lee <jason5lee@hotmail.com>")
.arg(
clap::Arg::with_name("config")
.short("c")
.long("config")
.value_name("FILE")
.help("The config file, default \".h2sr.toml\" in user directory.")
.takes_value(true),
)
.get_matches();
let config: Cow<'_, Path> = matches
.value_of("config")
.map(|path| Path::new(path).into())
.unwrap_or_else(|| {
let user_dir = directories::UserDirs::new().expect("unable to get user directory");
let mut path = user_dir.home_dir().to_owned();
path.push(".h2sr.toml");
path.into()
});
let config = fs::read(config).expect("unable to read config file");
let config: Config = toml::from_slice(&config).expect("config file error");
let proxy =
Pattern::from_strs(config.proxy.iter().map(|s| &*s as &str)).expect("proxy config error");
let block =
Pattern::from_strs(config.block.iter().map(|s| &*s as &str)).expect("block config error");
Env {
listen: config.listen,
to: Socks5Proxy(config.to),
proxy,
block,
}
}
#[tokio::main]
async fn main() {
let env: &'static _ = Box::leak(Box::new(load_env().await));
let client = HttpClient::new();
let make_service = make_service_fn(move |_| {
let client = client.clone();
async move { Ok::<_, Infallible>(service_fn(move |req| proxy(env, client.clone(), req))) }
});
let server = Server::bind(&env.listen).serve(make_service);
println!("Listening on http://{}", env.listen);
if let Err(e) = server.await {
eprintln!("server error: {}", e);
}
}
async fn proxy(
env: &'static Env,
client: HttpClient,
req: Request<Body>,
) -> anyhow::Result<Response<Body>> {
if Method::CONNECT == req.method() {
let uri = req.uri();
let auth = uri
.authority()
.ok_or_else(|| anyhow!("CONNECT host is not socket addr: {:?}", uri))?;
let host = auth.host().as_bytes();
let auth = auth.to_string();
let to = &env.to;
if env.block.contain_host(host) {
println!("BLOCK {}", auth);
} else if env.proxy.contain_host(host) {
println!("PROXY {}", auth);
tokio::task::spawn(async move { to.connect(&auth, req).await });
} else {
tokio::task::spawn(async move {
match lookup_host(&auth).await {
Err(e) => eprintln!("lookup host {} error: {:?}.", auth, e),
Ok(mut host) => match host.next() {
None => eprintln!("lookup host {} error: no ip found.", auth),
Some(ip) => {
if env.proxy.contain_ip(&ip.ip()) {
println!("PROXY {}", auth);
to.connect(&auth, req).await;
} else {
println!("DIRECT {}", auth);
match TcpStream::connect(&*auth).await {
Ok(mut stream) => tunnel(&auth, req, stream.split()).await,
Err(e) => eprintln!("{}: direct connect error: {:?}", auth, e),
}
}
}
},
}
});
};
Ok(Response::new(Body::empty()))
} else {
Ok(client.request(req).await?)
}
}
async fn tunnel(
auth: &str,
req: Request<Body>,
(mut server_rd, mut server_wr): (ReadHalf<'_>, WriteHalf<'_>),
) {
match hyper::upgrade::on(req).await {
Ok(upgraded) => {
let amounts = {
let (mut client_rd, mut client_wr) = tokio::io::split(upgraded);
let client_to_server = tokio::io::copy(&mut client_rd, &mut server_wr);
let server_to_client = tokio::io::copy(&mut server_rd, &mut client_wr);
try_join(client_to_server, server_to_client).await
};
match amounts {
Ok((from_client, from_server)) => {
println!(
"{}: client wrote {} bytes and received {} bytes",
auth, from_client, from_server
);
}
Err(e) => {
eprintln!("{}: tunnel error: {}", auth, e);
}
};
}
Err(e) => eprintln!("{}: upgrade error: {}", auth, e),
}
}