use super::{Transport, TransportError};
use async_trait::async_trait;
use std::pin::Pin;
use std::task::{Context, Poll};
use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};
use tokio::net::TcpStream;
pub struct TcpTransport {
stream: TcpStream,
}
impl TcpTransport {
pub fn from_stream(stream: TcpStream) -> Self {
Self { stream }
}
pub fn get_ref(&self) -> &TcpStream {
&self.stream
}
pub fn get_mut(&mut self) -> &mut TcpStream {
&mut self.stream
}
pub fn into_inner(self) -> TcpStream {
self.stream
}
}
#[async_trait]
impl Transport for TcpTransport {
async fn connect(addr: &str) -> Result<Self, TransportError> {
let stream = TcpStream::connect(addr).await.map_err(|e| {
TransportError::ConnectionFailed(format!("TCP connection failed: {}", e))
})?;
Ok(Self { stream })
}
async fn close(&mut self) -> Result<(), TransportError> {
Ok(())
}
fn peer_addr(&self) -> Result<String, TransportError> {
self.stream
.peer_addr()
.map(|addr| addr.to_string())
.map_err(TransportError::Io)
}
fn local_addr(&self) -> Result<String, TransportError> {
self.stream
.local_addr()
.map(|addr| addr.to_string())
.map_err(TransportError::Io)
}
fn set_nodelay(&self, nodelay: bool) -> Result<(), TransportError> {
self.stream.set_nodelay(nodelay).map_err(TransportError::Io)
}
}
impl AsyncRead for TcpTransport {
fn poll_read(
mut self: Pin<&mut Self>,
ctx: &mut Context<'_>,
buf: &mut ReadBuf<'_>,
) -> Poll<std::io::Result<()>> {
Pin::new(&mut self.stream).poll_read(ctx, buf)
}
}
impl AsyncWrite for TcpTransport {
fn poll_write(
mut self: Pin<&mut Self>,
ctx: &mut Context<'_>,
buf: &[u8],
) -> Poll<std::io::Result<usize>> {
Pin::new(&mut self.stream).poll_write(ctx, buf)
}
fn poll_flush(mut self: Pin<&mut Self>, ctx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
Pin::new(&mut self.stream).poll_flush(ctx)
}
fn poll_shutdown(mut self: Pin<&mut Self>, ctx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
Pin::new(&mut self.stream).poll_shutdown(ctx)
}
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
#[ignore] async fn test_tcp_transport_connect() {
let result = TcpTransport::connect("broker.emqx.io:1883").await;
assert!(result.is_ok(), "Should connect to public broker");
if let Ok(mut transport) = result {
let peer = transport.peer_addr();
assert!(peer.is_ok(), "Should get peer address");
let local = transport.local_addr();
assert!(local.is_ok(), "Should get local address");
let result = transport.set_nodelay(true);
assert!(result.is_ok(), "Should set TCP_NODELAY");
let result = transport.close().await;
assert!(result.is_ok(), "Should close connection");
}
}
#[tokio::test]
async fn test_tcp_transport_invalid_address() {
let result = TcpTransport::connect("invalid-address-that-does-not-exist:1883").await;
assert!(result.is_err(), "Should fail with invalid address");
}
#[tokio::test]
async fn test_tcp_transport_connection_refused() {
let result = TcpTransport::connect("127.0.0.1:65535").await;
assert!(result.is_err(), "Should fail when connection is refused");
}
}