From 938b29d418a9df2f93ee273a394f34adc99ea25d Mon Sep 17 00:00:00 2001
From: hozan23 <hozan23@proton.me>
Date: Sat, 18 Nov 2023 13:36:19 +0300
Subject: net: improve Conn API

---
 net/src/transports/tcp.rs  | 17 ++++++++++-------
 net/src/transports/udp.rs  | 12 +++++-------
 net/src/transports/unix.rs | 17 ++++++++++-------
 3 files changed, 25 insertions(+), 21 deletions(-)

(limited to 'net/src/transports')

diff --git a/net/src/transports/tcp.rs b/net/src/transports/tcp.rs
index 5ff7b28..9b2b928 100644
--- a/net/src/transports/tcp.rs
+++ b/net/src/transports/tcp.rs
@@ -10,7 +10,7 @@ use crate::{
     connection::Connection,
     endpoint::{Addr, Endpoint, Port},
     listener::Listener,
-    Result,
+    Error, Result,
 };
 
 /// TCP network connection implementations of the `Connection` trait.
@@ -42,14 +42,17 @@ impl Connection for TcpConn {
         Ok(Endpoint::new_tcp_addr(&self.inner.local_addr()?))
     }
 
-    async fn recv(&self, buf: &mut [u8]) -> Result<usize> {
-        self.read.lock().await.read_exact(buf).await?;
-        Ok(buf.len())
+    async fn read(&self, buf: &mut [u8]) -> Result<usize> {
+        self.read.lock().await.read(buf).await.map_err(Error::from)
     }
 
-    async fn send(&self, buf: &[u8]) -> Result<usize> {
-        self.write.lock().await.write_all(buf).await?;
-        Ok(buf.len())
+    async fn write(&self, buf: &[u8]) -> Result<usize> {
+        self.write
+            .lock()
+            .await
+            .write(buf)
+            .await
+            .map_err(Error::from)
     }
 }
 
diff --git a/net/src/transports/udp.rs b/net/src/transports/udp.rs
index 27fb9ae..2050226 100644
--- a/net/src/transports/udp.rs
+++ b/net/src/transports/udp.rs
@@ -6,7 +6,7 @@ use smol::net::UdpSocket;
 use crate::{
     connection::Connection,
     endpoint::{Addr, Endpoint, Port},
-    Result,
+    Error, Result,
 };
 
 /// UDP network connection implementations of the `Connection` trait.
@@ -47,14 +47,12 @@ impl Connection for UdpConn {
         Ok(Endpoint::new_udp_addr(&self.inner.local_addr()?))
     }
 
-    async fn recv(&self, buf: &mut [u8]) -> Result<usize> {
-        let size = self.inner.recv(buf).await?;
-        Ok(size)
+    async fn read(&self, buf: &mut [u8]) -> Result<usize> {
+        self.inner.recv(buf).await.map_err(Error::from)
     }
 
-    async fn send(&self, buf: &[u8]) -> Result<usize> {
-        let size = self.inner.send(buf).await?;
-        Ok(size)
+    async fn write(&self, buf: &[u8]) -> Result<usize> {
+        self.inner.send(buf).await.map_err(Error::from)
     }
 }
 
diff --git a/net/src/transports/unix.rs b/net/src/transports/unix.rs
index c89832e..1a32311 100644
--- a/net/src/transports/unix.rs
+++ b/net/src/transports/unix.rs
@@ -6,7 +6,7 @@ use smol::{
     net::unix::{UnixListener, UnixStream},
 };
 
-use crate::{connection::Connection, endpoint::Endpoint, listener::Listener, Result};
+use crate::{connection::Connection, endpoint::Endpoint, listener::Listener, Error, Result};
 
 /// Unix domain socket implementations of the `Connection` trait.
 pub struct UnixConn {
@@ -37,14 +37,17 @@ impl Connection for UnixConn {
         Ok(Endpoint::new_unix_addr(&self.inner.local_addr()?))
     }
 
-    async fn recv(&self, buf: &mut [u8]) -> Result<usize> {
-        self.read.lock().await.read_exact(buf).await?;
-        Ok(buf.len())
+    async fn read(&self, buf: &mut [u8]) -> Result<usize> {
+        self.read.lock().await.read(buf).await.map_err(Error::from)
     }
 
-    async fn send(&self, buf: &[u8]) -> Result<usize> {
-        self.write.lock().await.write_all(buf).await?;
-        Ok(buf.len())
+    async fn write(&self, buf: &[u8]) -> Result<usize> {
+        self.write
+            .lock()
+            .await
+            .write(buf)
+            .await
+            .map_err(Error::from)
     }
 }
 
-- 
cgit v1.2.3