use crate::errors::*;
use crate::p2p::proto::PeerAddr;
use std::net::SocketAddr;
use std::time::Duration;
use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
use tokio::net::TcpStream;
use tokio::time;
use tokio_socks::tcp::Socks5Stream;
pub const CONNECT_TIMEOUT: Duration = Duration::from_secs(15);
pub const HANDSHAKE_TIMEOUT: Duration = Duration::from_secs(20);
pub const PROXY_TIMEOUT: Duration = Duration::from_secs(30);
pub async fn connect(addr: &PeerAddr, proxy: Option<SocketAddr>) -> Result<TcpStream> {
let PeerAddr::Inet(addr) = addr else {
bail!("Connecting to onions is not yet implemented")
};
let target = proxy.unwrap_or(*addr);
info!("Creating tcp connection to {target:?}");
let sock = TcpStream::connect(target);
let mut sock = time::timeout(CONNECT_TIMEOUT, sock)
.await
.with_context(|| anyhow!("Connecting to {target:?} timed out"))?
.with_context(|| anyhow!("Failed to connect to {target:?}"))?;
if let Some(proxy) = proxy {
debug!("Requesting socks5 connection to {addr:?}");
let connect = Socks5Stream::connect_with_socket(sock, addr);
sock = time::timeout(PROXY_TIMEOUT, connect)
.await
.with_context(|| anyhow!("Connecting to {addr:?} (with socks5 {proxy:?}) timed out"))?
.with_context(|| anyhow!("Failed to connect to {addr:?} (with socks5 {proxy:?})"))?
.into_inner()
}
debug!("Connection has been established");
Ok(sock)
}
pub async fn handshake<R: AsyncRead + Unpin, W: AsyncWrite + Unpin>(
mut rx: R,
mut tx: W,
) -> Result<()> {
debug!("Sending protocol probe to remote peer");
time::timeout(HANDSHAKE_TIMEOUT, tx.write_all(b"//\n"))
.await
.context("Sending handshake protocol probe timed out")?
.context("Failed to send protocol probe")?;
let mut buf = [0u8; 3];
time::timeout(HANDSHAKE_TIMEOUT, rx.read_exact(&mut buf))
.await
.context("Sending handshake protocol probe timed out")?
.context("Failed to receive handshake response")?;
if buf == *b":0\n" {
debug!("Remote peer has sent expected response");
Ok(())
} else {
bail!("Invalid handshake response: {buf:?}")
}
}