117 lines
3.8 KiB
Rust
117 lines
3.8 KiB
Rust
use anyhow::Result;
|
|
use async_trait::async_trait;
|
|
use futures_util::{SinkExt, StreamExt};
|
|
use tokio_tungstenite::tungstenite::Message;
|
|
|
|
use crate::transport::WsStream;
|
|
|
|
// ============================================================================
|
|
// Transport trait abstraction
|
|
// ============================================================================
|
|
|
|
/// Outbound half of a VPN transport connection.
|
|
#[async_trait]
|
|
pub trait TransportSink: Send + 'static {
|
|
/// Send a framed binary message on the reliable channel.
|
|
async fn send_reliable(&mut self, data: Vec<u8>) -> Result<()>;
|
|
|
|
/// Send a datagram (unreliable, best-effort).
|
|
/// Falls back to reliable if the transport does not support datagrams.
|
|
async fn send_datagram(&mut self, data: Vec<u8>) -> Result<()>;
|
|
|
|
/// Gracefully close the transport.
|
|
async fn close(&mut self) -> Result<()>;
|
|
}
|
|
|
|
/// Inbound half of a VPN transport connection.
|
|
#[async_trait]
|
|
pub trait TransportStream: Send + 'static {
|
|
/// Receive the next reliable binary message. Returns `None` on close.
|
|
async fn recv_reliable(&mut self) -> Result<Option<Vec<u8>>>;
|
|
|
|
/// Receive the next datagram. Returns `None` if datagrams are unsupported
|
|
/// or the connection is closed.
|
|
async fn recv_datagram(&mut self) -> Result<Option<Vec<u8>>>;
|
|
|
|
/// Whether this transport supports unreliable datagrams.
|
|
fn supports_datagrams(&self) -> bool;
|
|
}
|
|
|
|
// ============================================================================
|
|
// WebSocket implementation
|
|
// ============================================================================
|
|
|
|
/// WebSocket transport sink (wraps the write half of a split WsStream).
|
|
pub struct WsTransportSink {
|
|
inner: futures_util::stream::SplitSink<WsStream, Message>,
|
|
}
|
|
|
|
impl WsTransportSink {
|
|
pub fn new(inner: futures_util::stream::SplitSink<WsStream, Message>) -> Self {
|
|
Self { inner }
|
|
}
|
|
}
|
|
|
|
#[async_trait]
|
|
impl TransportSink for WsTransportSink {
|
|
async fn send_reliable(&mut self, data: Vec<u8>) -> Result<()> {
|
|
self.inner.send(Message::Binary(data.into())).await?;
|
|
Ok(())
|
|
}
|
|
|
|
async fn send_datagram(&mut self, data: Vec<u8>) -> Result<()> {
|
|
// WebSocket has no datagram support — fall back to reliable.
|
|
self.send_reliable(data).await
|
|
}
|
|
|
|
async fn close(&mut self) -> Result<()> {
|
|
self.inner.close().await?;
|
|
Ok(())
|
|
}
|
|
}
|
|
|
|
/// WebSocket transport stream (wraps the read half of a split WsStream).
|
|
pub struct WsTransportStream {
|
|
inner: futures_util::stream::SplitStream<WsStream>,
|
|
}
|
|
|
|
impl WsTransportStream {
|
|
pub fn new(inner: futures_util::stream::SplitStream<WsStream>) -> Self {
|
|
Self { inner }
|
|
}
|
|
}
|
|
|
|
#[async_trait]
|
|
impl TransportStream for WsTransportStream {
|
|
async fn recv_reliable(&mut self) -> Result<Option<Vec<u8>>> {
|
|
loop {
|
|
match self.inner.next().await {
|
|
Some(Ok(Message::Binary(data))) => return Ok(Some(data.to_vec())),
|
|
Some(Ok(Message::Close(_))) | None => return Ok(None),
|
|
Some(Ok(Message::Ping(_))) => {
|
|
// Ping handling is done at the tungstenite layer automatically
|
|
// when the sink side is alive. Just skip here.
|
|
continue;
|
|
}
|
|
Some(Ok(_)) => continue,
|
|
Some(Err(e)) => return Err(anyhow::anyhow!("WebSocket error: {}", e)),
|
|
}
|
|
}
|
|
}
|
|
|
|
async fn recv_datagram(&mut self) -> Result<Option<Vec<u8>>> {
|
|
// WebSocket does not support datagrams.
|
|
Ok(None)
|
|
}
|
|
|
|
fn supports_datagrams(&self) -> bool {
|
|
false
|
|
}
|
|
}
|
|
|
|
/// Split a WebSocket stream into transport sink and stream halves.
|
|
pub fn split_ws(ws: WsStream) -> (WsTransportSink, WsTransportStream) {
|
|
let (sink, stream) = ws.split();
|
|
(WsTransportSink::new(sink), WsTransportStream::new(stream))
|
|
}
|