use tokio::net::{TcpStream, UdpSocket};
use std::io::{self, Result};
use tokio::io::{AsyncReadExt, AsyncWriteExt};
pub struct Remote {
tcp_stream: Option<TcpStream>,
udp_socket: Option<UdpSocket>,
}
impl Remote {
pub async fn connect_tcp(address: &str) -> Result<Self> {
let tcp_stream = TcpStream::connect(address).await?;
Ok(Remote {
tcp_stream: Some(tcp_stream),
udp_socket: None,
})
}
pub async fn connect_udp(address: &str) -> Result<Self> {
let udp_socket = UdpSocket::bind("0.0.0.0:0").await?;
udp_socket.connect(address).await?;
Ok(Remote {
tcp_stream: None,
udp_socket: Some(udp_socket),
})
}
pub async fn send(&mut self, data: &[u8]) -> Result<()> {
if let Some(ref mut stream) = self.tcp_stream {
stream.write_all(data).await?;
} else if let Some(ref mut socket) = self.udp_socket {
socket.send(data).await?;
}
Ok(())
}
pub async fn receive(&mut self, buffer: &mut [u8]) -> Result<usize> {
if let Some(ref mut stream) = self.tcp_stream {
let n = stream.read(buffer).await?;
Ok(n)
} else if let Some(ref mut socket) = self.udp_socket {
let (n, _) = socket.recv_from(buffer).await?;
Ok(n)
} else {
Err(io::Error::new(io::ErrorKind::Other, "No connection established"))
}
}
}