1 files changed, 156 insertions, 0 deletions

diff --git a/src/tcp.rs b/src/tcp.rs
new file mode 100644
index 0000000..5684c41
--- /dev/null
+++ b/src/tcp.rs
@@ -0,0 +1,156 @@
+//! TCP backed channel and listener implementations.
+//!
+//! ```no_run
+//! #[urpc::service]
+//! trait Hello {
+//!     type Error = ();
+//!
+//!     async fn hello(name: String) -> String;
+//! }
+//!
+//! struct HelloServer;
+//!
+//! impl Hello for HelloServer {
+//!     async fn hello(&self, _ctx: urpc::Context, name: String) -> Result<String, ()> {
+//!         Ok(format!("Hello, {name}!"))
+//!     }
+//! }
+//!
+//! #[tokio::main]
+//! async fn main() -> Result<(), Box<dyn std::error::Error>>{
+//!     let listener = urpc::tcp::bind("0.0.0.0:3000").await?;
+//!
+//!     // spawn the server
+//!     tokio::spawn(async move {
+//!         urpc::Server::default()
+//!             .with_listener(listener)
+//!             .with_service(HelloServer.into_service())
+//!             .serve()
+//!             .await
+//!     });
+//!
+//!     // create a client
+//!     let channel = urpc::ClientChannel::new(urpc::tcp::connect("127.0.0.1:3000").await?);
+//!     let client = HelloClient::new(channel);
+//!     let greeting = client.hello("World".into()).await.unwrap();
+//!     assert_eq!(greeting, "Hello, World!");
+//!     Ok(())
+//! }
+//! ```
+use std::pin::Pin;
+
+use futures::{Sink, Stream};
+use tokio::{
+    net::{TcpListener, TcpStream, ToSocketAddrs},
+    sync::mpsc::Receiver,
+    task::AbortHandle,
+};
+use tokio_util::codec::Framed;
+
+use crate::{
+    Channel, Listener,
+    protocol::{RpcMessage, RpcMessageCodec},
+};
+
+pub struct TcpChannel(Framed<TcpStream, RpcMessageCodec>);
+
+impl TcpChannel {
+    fn new(stream: TcpStream) -> Self {
+        Self(Framed::new(stream, RpcMessageCodec::default()))
+    }
+
+    pub async fn connect(addrs: impl ToSocketAddrs) -> std::io::Result<Self> {
+        let stream = TcpStream::connect(addrs).await?;
+        Ok(Self::new(stream))
+    }
+}
+
+impl Sink<RpcMessage> for TcpChannel {
+    type Error = std::io::Error;
+
+    fn poll_ready(
+        self: Pin<&mut Self>,
+        cx: &mut std::task::Context<'_>,
+    ) -> std::task::Poll<std::result::Result<(), Self::Error>> {
+        Sink::poll_ready(Pin::new(&mut self.get_mut().0), cx)
+    }
+
+    fn start_send(self: Pin<&mut Self>, item: RpcMessage) -> std::result::Result<(), Self::Error> {
+        Sink::start_send(Pin::new(&mut self.get_mut().0), item)
+    }
+
+    fn poll_flush(
+        self: Pin<&mut Self>,
+        cx: &mut std::task::Context<'_>,
+    ) -> std::task::Poll<std::result::Result<(), Self::Error>> {
+        Sink::poll_flush(Pin::new(&mut self.get_mut().0), cx)
+    }
+
+    fn poll_close(
+        self: Pin<&mut Self>,
+        cx: &mut std::task::Context<'_>,
+    ) -> std::task::Poll<std::result::Result<(), Self::Error>> {
+        Sink::poll_close(Pin::new(&mut self.get_mut().0), cx)
+    }
+}
+
+impl Stream for TcpChannel {
+    type Item = std::io::Result<RpcMessage>;
+
+    fn poll_next(
+        self: Pin<&mut Self>,
+        cx: &mut std::task::Context<'_>,
+    ) -> std::task::Poll<Option<Self::Item>> {
+        Stream::poll_next(Pin::new(&mut self.get_mut().0), cx)
+    }
+}
+
+impl Channel for TcpChannel {}
+
+pub struct TcpChannelListener {
+    receiver: Receiver<TcpChannel>,
+    abort: AbortHandle,
+}
+
+impl Drop for TcpChannelListener {
+    fn drop(&mut self) {
+        self.abort.abort();
+    }
+}
+
+impl TcpChannelListener {
+    pub async fn bind(addrs: impl ToSocketAddrs) -> std::io::Result<Self> {
+        let listener = TcpListener::bind(addrs).await?;
+        let (sender, receiver) = tokio::sync::mpsc::channel(8);
+        let abort = tokio::spawn(async move {
+            while let Ok((stream, _addr)) = listener.accept().await {
+                if sender.send(TcpChannel::new(stream)).await.is_err() {
+                    break;
+                }
+            }
+        })
+        .abort_handle();
+        Ok(Self { receiver, abort })
+    }
+}
+
+impl Stream for TcpChannelListener {
+    type Item = std::io::Result<TcpChannel>;
+
+    fn poll_next(
+        self: Pin<&mut Self>,
+        cx: &mut std::task::Context<'_>,
+    ) -> std::task::Poll<Option<Self::Item>> {
+        self.get_mut().receiver.poll_recv(cx).map(|v| v.map(Ok))
+    }
+}
+
+impl Listener<TcpChannel> for TcpChannelListener {}
+
+pub async fn bind(addrs: impl ToSocketAddrs) -> std::io::Result<TcpChannelListener> {
+    TcpChannelListener::bind(addrs).await
+}
+
+pub async fn connect(addrs: impl ToSocketAddrs) -> std::io::Result<TcpChannel> {
+    TcpChannel::connect(addrs).await
+}