From 196a887f695f5fb9064df22e93abd714c0855537 Mon Sep 17 00:00:00 2001
From: "Dr. Maxim Orlovsky" <orlovsky@cyphernet.org>
Date: Fri, 27 Jan 2023 12:40:35 +0100
Subject: [PATCH] AsConnection trait

---
 src/connection.rs |  7 ++++++-
 src/lib.rs        |  2 +-
 src/split.rs      | 25 +++++++++++++++++++++++--
 3 files changed, 30 insertions(+), 4 deletions(-)

diff --git a/src/connection.rs b/src/connection.rs
index e4ce53c..c74a39d 100644
--- a/src/connection.rs
+++ b/src/connection.rs
@@ -33,8 +33,13 @@ impl<T> Address for T where T: Addr + Send + Clone + Eq + Hash + Debug + Display
 
 pub trait NetStream: Send + io::Read + io::Write {}
 
+pub trait AsConnection {
+    type Connection: NetConnection;
+    fn as_connection(&self) -> &Self::Connection;
+}
+
 /// Network stream is an abstraction of TCP stream object.
-pub trait NetConnection: Send + NetStream + AsRawFd + Debug {
+pub trait NetConnection: NetStream + AsRawFd + Debug {
     type Addr: Address;
 
     fn connect_blocking(addr: Self::Addr) -> io::Result<Self>
diff --git a/src/lib.rs b/src/lib.rs
index dd655c3..9e85006 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -40,7 +40,7 @@ pub mod resource;
 
 pub const READ_BUFFER_SIZE: usize = u16::MAX as usize;
 
-pub use connection::{Address, NetConnection, NetStream};
+pub use connection::{Address, AsConnection, NetConnection, NetStream};
 pub use frame::{Frame, Marshaller};
 pub use listener::NetListener;
 #[cfg(feature = "io-reactor")]
diff --git a/src/split.rs b/src/split.rs
index 62d93ac..b70d1cc 100644
--- a/src/split.rs
+++ b/src/split.rs
@@ -22,6 +22,7 @@
 use std::io;
 use std::net::TcpStream;
 
+use crate::connection::AsConnection;
 use crate::{NetConnection, NetSession, NetStateMachine};
 
 #[derive(Debug, Display)]
@@ -34,8 +35,8 @@ pub struct SplitIoError<T: SplitIo> {
 impl<T: SplitIo + std::fmt::Debug> std::error::Error for SplitIoError<T> {}
 
 pub trait SplitIo: Sized {
-    type Read: io::Read + Sized;
-    type Write: io::Write + Sized;
+    type Read: AsConnection + io::Read + Sized;
+    type Write: AsConnection + io::Write + Sized;
 
     /// # Panics
     ///
@@ -49,6 +50,11 @@ pub struct NetReader<S: NetSession> {
     pub(crate) reader: <S as SplitIo>::Read,
 }
 
+impl<S: NetSession> AsConnection for NetReader<S> {
+    type Connection = <<S as SplitIo>::Read as AsConnection>::Connection;
+    fn as_connection(&self) -> &Self::Connection { self.reader.as_connection() }
+}
+
 impl<S: NetSession> io::Read for NetReader<S> {
     fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> { self.reader.read(buf) }
 }
@@ -59,6 +65,11 @@ pub struct NetWriter<M: NetStateMachine, S: NetSession> {
     pub(crate) writer: <S as SplitIo>::Write,
 }
 
+impl<M: NetStateMachine, S: NetSession> AsConnection for NetWriter<M, S> {
+    type Connection = <<S as SplitIo>::Write as AsConnection>::Connection;
+    fn as_connection(&self) -> &Self::Connection { self.writer.as_connection() }
+}
+
 impl<M: NetStateMachine, S: NetSession> io::Write for NetWriter<M, S> {
     fn write(&mut self, buf: &[u8]) -> io::Result<usize> { self.writer.write(buf) }
     fn flush(&mut self) -> io::Result<()> { self.writer.flush() }
@@ -73,6 +84,11 @@ impl<C: NetConnection> io::Read for TcpReader<C> {
     fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> { self.connection.read(buf) }
 }
 
+impl<C: NetConnection> AsConnection for TcpReader<C> {
+    type Connection = C;
+    fn as_connection(&self) -> &Self::Connection { &self.connection }
+}
+
 pub struct TcpWriter<C: NetConnection> {
     unique_id: u64,
     connection: C,
@@ -84,6 +100,11 @@ impl<C: NetConnection> io::Write for TcpWriter<C> {
     fn flush(&mut self) -> io::Result<()> { self.connection.flush() }
 }
 
+impl<C: NetConnection> AsConnection for TcpWriter<C> {
+    type Connection = C;
+    fn as_connection(&self) -> &Self::Connection { &self.connection }
+}
+
 impl SplitIo for TcpStream {
     type Read = TcpReader<Self>;
     type Write = TcpWriter<Self>;