Add debug to blocking

This commit is contained in:
Aram 🍐 2021-08-05 01:38:27 -04:00
parent 3b3f9278de
commit cc409f853b
2 changed files with 30 additions and 3 deletions

View file

@ -24,9 +24,15 @@ impl Client {
let tcp_stream = TcpStream::connect_timeout(&host.addr, config.timeout)?; let tcp_stream = TcpStream::connect_timeout(&host.addr, config.timeout)?;
let mut client = Client { let mut client = Client {
config: config.clone(), config: config.clone(),
stream: ConnectionStream::Tcp(tcp_stream).buffered(), stream: ConnectionStream::Tcp(tcp_stream),
}; };
if config.debug {
client.stream = client.stream.debug();
}
client.stream = client.stream.buffered();
client = client.enable_ssl()?; client = client.enable_ssl()?;
Ok(client) Ok(client)

View file

@ -15,6 +15,11 @@ pub enum ConnectionStream {
/// It can then be un-wrapped with `.unbuffered()`. /// It can then be un-wrapped with `.unbuffered()`.
Buffered(Box<BufReader<ConnectionStream>>), Buffered(Box<BufReader<ConnectionStream>>),
/// A stream wrapped with a debug logging mechanism.
///
/// Use `.debug()` to wrap any stream with this wrapper.
Debug(Box<ConnectionStream>),
/// A client stream wrapped with SSL using `rustls`. /// A client stream wrapped with SSL using `rustls`.
#[cfg(feature = "ssl")] #[cfg(feature = "ssl")]
SslClient(Box<rustls::StreamOwned<rustls::ClientSession, ConnectionStream>>), SslClient(Box<rustls::StreamOwned<rustls::ClientSession, ConnectionStream>>),
@ -129,7 +134,7 @@ impl ConnectionStream {
} }
/// Wraps the current stream with a `BufReader`. /// Wraps the current stream with a `BufReader`.
pub fn buffered(self) -> ConnectionStream { pub fn buffered(self) -> Self {
Self::Buffered(Box::new(BufReader::new(self))) Self::Buffered(Box::new(BufReader::new(self)))
} }
@ -137,13 +142,18 @@ impl ConnectionStream {
/// If the current stream is not buffered, it returns itself (no-op). /// If the current stream is not buffered, it returns itself (no-op).
/// ///
/// Note that, if the stream is buffered, any un-consumed data will be discarded. /// Note that, if the stream is buffered, any un-consumed data will be discarded.
pub fn unbuffered(self) -> ConnectionStream { pub fn unbuffered(self) -> Self {
if let Self::Buffered(buf) = self { if let Self::Buffered(buf) = self {
buf.into_inner() buf.into_inner()
} else { } else {
self self
} }
} }
/// Wraps the current stream with a debug logging mechanism.
pub fn debug(self) -> Self {
Self::Debug(Box::new(self))
}
} }
impl Read for ConnectionStream { impl Read for ConnectionStream {
@ -151,6 +161,11 @@ impl Read for ConnectionStream {
match self { match self {
Self::Tcp(stream) => stream.read(buf), Self::Tcp(stream) => stream.read(buf),
Self::Buffered(reader) => reader.read(buf), Self::Buffered(reader) => reader.read(buf),
Self::Debug(stream) => {
let str = String::from_utf8_lossy(buf).trim().to_string();
println!("DEBUG::READ <- '{}'", str); // TODO: Replace with logger?
stream.read(buf)
}
#[cfg(feature = "ssl")] #[cfg(feature = "ssl")]
Self::SslClient(stream) => stream.read(buf), Self::SslClient(stream) => stream.read(buf),
#[cfg(feature = "ssl")] #[cfg(feature = "ssl")]
@ -184,6 +199,11 @@ impl Write for ConnectionStream {
match self { match self {
Self::Tcp(stream) => stream.write(buf), Self::Tcp(stream) => stream.write(buf),
Self::Buffered(reader) => reader.get_mut().write(buf), Self::Buffered(reader) => reader.get_mut().write(buf),
Self::Debug(stream) => {
let str = String::from_utf8_lossy(buf).trim().to_string();
println!("DEBUG::WRITE -> '{}'", str); // TODO: Replace with logger?
stream.write(buf)
}
#[cfg(feature = "ssl")] #[cfg(feature = "ssl")]
Self::SslClient(stream) => stream.write(buf), Self::SslClient(stream) => stream.write(buf),
#[cfg(feature = "ssl")] #[cfg(feature = "ssl")]
@ -201,6 +221,7 @@ impl Write for ConnectionStream {
match self { match self {
Self::Tcp(stream) => stream.flush(), Self::Tcp(stream) => stream.flush(),
Self::Buffered(reader) => reader.get_mut().flush(), Self::Buffered(reader) => reader.get_mut().flush(),
Self::Debug(stream) => stream.flush(),
#[cfg(feature = "ssl")] #[cfg(feature = "ssl")]
Self::SslClient(stream) => stream.flush(), Self::SslClient(stream) => stream.flush(),
#[cfg(feature = "ssl")] #[cfg(feature = "ssl")]