89 lines
3.0 KiB
Rust
89 lines
3.0 KiB
Rust
use std::time::{Duration, Instant};
|
|
|
|
use tokio::{
|
|
io::{AsyncReadExt, AsyncWriteExt},
|
|
net::{TcpListener, TcpStream},
|
|
process::Command,
|
|
};
|
|
|
|
#[tokio::main(flavor = "current_thread")]
|
|
async fn main() {
|
|
eprintln!(
|
|
"USAGE:
|
|
- Set env var ADDR=addr:port
|
|
- Set env var INTERVAL=seconds
|
|
- Set env vars ADDR1-ADDRn for n proxies
|
|
- Create scripts ./proxy1 - ./proxy(n-1) to check if proxy target is online (exit status 0)
|
|
- If all scripts fail, proxy n is used"
|
|
);
|
|
let interval: u64 = std::env::var("INTERVAL")
|
|
.expect("expected env var INTERVAL to be set")
|
|
.trim()
|
|
.parse()
|
|
.expect(
|
|
"expected env var INTERVAL to be set to an integer value greater than or equal to zero",
|
|
);
|
|
let proxies = (1..)
|
|
.map_while(|i| std::env::var(format!("ADDR{i}")).ok())
|
|
.collect::<Vec<_>>();
|
|
if proxies.is_empty() {
|
|
panic!("expected env vars ADDR1-ADDRn");
|
|
}
|
|
|
|
let listener =
|
|
TcpListener::bind(std::env::var("ADDR").expect("expected env var ADDR to be set"))
|
|
.await
|
|
.unwrap();
|
|
|
|
let mut next_check = Instant::now();
|
|
let mut selected_proxy = 0;
|
|
loop {
|
|
let (con, _) = listener.accept().await.unwrap();
|
|
let now = Instant::now();
|
|
if now >= next_check {
|
|
next_check = now + Duration::from_secs(interval);
|
|
let pproxy = selected_proxy;
|
|
selected_proxy = proxies.len();
|
|
for i in 1..proxies.len() {
|
|
if Command::new(format!("./proxy{i}"))
|
|
.status()
|
|
.await
|
|
.is_ok_and(|status| status.success())
|
|
{
|
|
selected_proxy = i;
|
|
}
|
|
}
|
|
if selected_proxy != pproxy {
|
|
eprintln!(
|
|
"Using proxy {selected_proxy}: {}",
|
|
proxies[selected_proxy - 1]
|
|
);
|
|
}
|
|
}
|
|
if selected_proxy > 0 {
|
|
if let Some(proxy) = proxies.get(selected_proxy - 1).cloned() {
|
|
if let Ok(prox) = TcpStream::connect(&proxy).await {
|
|
let (mut con_read, mut con_write) = con.into_split();
|
|
let (mut prox_read, mut prox_write) = prox.into_split();
|
|
tokio::task::spawn(async move {
|
|
let mut buf = [0u8; 256];
|
|
while let Ok(n) = con_read.read(&mut buf).await {
|
|
if n > 0 && prox_write.write_all(&buf[0..n]).await.is_err() {
|
|
break;
|
|
}
|
|
}
|
|
});
|
|
tokio::task::spawn(async move {
|
|
let mut buf = [0u8; 256];
|
|
while let Ok(n) = prox_read.read(&mut buf).await {
|
|
if n > 0 && con_write.write_all(&buf[0..n]).await.is_err() {
|
|
break;
|
|
}
|
|
}
|
|
});
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|