use anyhow::{anyhow, Result};
use backoff::{backoff::Backoff, Notify};
use socket2::{SockRef, TcpKeepalive};
use std::{future::Future, net::SocketAddr, time::Duration};
use tokio::{
net::{lookup_host, TcpStream, ToSocketAddrs, UdpSocket},
sync::broadcast,
};
use tracing::trace;
pub fn try_set_tcp_keepalive(
conn: &TcpStream,
keepalive_duration: Duration,
keepalive_interval: Duration,
) -> Result<()> {
let s = SockRef::from(conn);
let keepalive = TcpKeepalive::new()
.with_time(keepalive_duration)
.with_interval(keepalive_interval);
trace!(
"Set TCP keepalive {:?} {:?}",
keepalive_duration,
keepalive_interval
);
Ok(s.set_tcp_keepalive(&keepalive)?)
}
#[allow(dead_code)]
pub fn feature_not_compile(feature: &str) -> ! {
panic!(
"The feature '{}' is not compiled in this binary. Please re-compile rathole",
feature
)
}
pub async fn udp_connect<A: ToSocketAddrs>(addr: A) -> Result<UdpSocket> {
let addr = lookup_host(addr)
.await?
.next()
.ok_or(anyhow!("Failed to lookup the host"))?;
let bind_addr = match addr {
SocketAddr::V4(_) => "0.0.0.0:0",
SocketAddr::V6(_) => ":::0",
};
let s = UdpSocket::bind(bind_addr).await?;
s.connect(addr).await?;
Ok(s)
}
pub async fn retry_notify_with_deadline<I, E, Fn, Fut, B, N>(
backoff: B,
operation: Fn,
notify: N,
deadline: &mut broadcast::Receiver<bool>,
) -> Result<I>
where
E: std::error::Error + Send + Sync + 'static,
B: Backoff,
Fn: FnMut() -> Fut,
Fut: Future<Output = std::result::Result<I, backoff::Error<E>>>,
N: Notify<E>,
{
tokio::select! {
v = backoff::future::retry_notify(backoff, operation, notify) => {
v.map_err(anyhow::Error::new)
}
_ = deadline.recv() => {
Err(anyhow!("shutdown"))
}
}
}