Something went wrong. Try again.
Monorepo for Tangled
Something went wrong. Try again.
2.7 kB · 92 lines
Rust
at pull2
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293use anyhow::{Context, Result};use tokio::net::{TcpListener, TcpStream};use tokio::task::{JoinError, JoinHandle, JoinSet};use tokio_vsock::{VsockAddr, VsockStream};use tracing::{info, warn};
// this implements a vsock <-> tcp proxy for communicating with spindlepub struct VsockTcpProxy { url: String, handle: JoinHandle<()>,}
impl VsockTcpProxy { pub async fn start( name: &'static str, bind_addr: &str, host_cid: u32, host_port: u32, ) -> Result<Self> { if host_port == 0 { anyhow::bail!("port 0 cant be requested"); }
let listener = TcpListener::bind(bind_addr) .await .with_context(|| format!("bind {name} listener {bind_addr}"))?; let local_addr = listener .local_addr() .with_context(|| format!("{name} local address"))?; let url = format!("http://{local_addr}");
let handle = tokio::spawn(async move { accept_loop(name, listener, host_cid, host_port).await; });
info!(%url, host_cid, host_port, "{name} ready"); Ok(Self { url, handle }) }
pub fn url(&self) -> &str { &self.url }}
impl Drop for VsockTcpProxy { fn drop(&mut self) { self.handle.abort(); }}
async fn accept_loop(name: &'static str, listener: TcpListener, host_cid: u32, host_port: u32) { let mut tasks = JoinSet::new(); loop { tokio::select! { accepted = listener.accept() => match accepted { Ok((conn, _addr)) => { tasks.spawn(async move { if let Err(error) = proxy_conn(name, conn, host_cid, host_port).await { warn!(%error, "{name} connection failed"); } }); } Err(error) => warn!(%error, "{name} accept failed"), }, Some(result) = tasks.join_next(), if !tasks.is_empty() => { log_proxy_task_result(result); } } }}
fn log_proxy_task_result(result: Result<(), JoinError>) { if let Err(error) = result { warn!(%error, "proxy task failed"); }}
async fn proxy_conn( name: &'static str, mut tcp: TcpStream, host_cid: u32, host_port: u32,) -> Result<()> { let mut host = VsockStream::connect(VsockAddr::new(host_cid, host_port)) .await .with_context(|| format!("dial host {name} cid={host_cid} port={host_port}"))?;
tokio::io::copy_bidirectional(&mut tcp, &mut host) .await .context("proxy connection copy")?; Ok(())}