From 9759fe6e28789dd3bb0f526d0d42cf924f62d3af Mon Sep 17 00:00:00 2001
From: Mike Dilger <mike@mikedilger.com>
Date: Sat, 22 Jun 2024 13:04:49 +1200
Subject: [PATCH] Restore (and improve) byte counting with CountingStream

---
 src/bin/chorus.rs      |  7 ++++--
 src/counting_stream.rs | 53 ++++++++++++++++++++++++++++++++++++++++++
 src/lib.rs             |  6 +++--
 3 files changed, 62 insertions(+), 4 deletions(-)
 create mode 100644 src/counting_stream.rs

diff --git a/src/bin/chorus.rs b/src/bin/chorus.rs
index 52e2040..207a8fb 100644
--- a/src/bin/chorus.rs
+++ b/src/bin/chorus.rs
@@ -1,4 +1,5 @@
 use chorus::config::{Config, FriendlyConfig};
+use chorus::counting_stream::CountingStream;
 use chorus::error::Error;
 use chorus::globals::GLOBALS;
 use chorus::ip::HashedPeer;
@@ -92,6 +93,8 @@ async fn main() -> Result<(), Error> {
                     (tcp_stream, hashed_peer)
                 };
 
+                let counting_stream = CountingStream(tcp_stream);
+
                 // Possibly IP block
                 if GLOBALS.config.read().enable_ip_blocking {
                     let ip_data = chorus::get_ip_data(GLOBALS.store.get().unwrap(), hashed_peer.ip())?;
@@ -108,7 +111,7 @@ async fn main() -> Result<(), Error> {
                 tokio::spawn(async move {
                     let stream: Box<dyn FullStream> = match maybe_tls_acceptor_clone {
                         Some(tls_acceptor) => {
-                            match tls_acceptor.accept(tcp_stream).await {
+                            match tls_acceptor.accept(counting_stream).await {
                                 Ok(stream) => Box::new(stream),
                                 Err(e) => {
                                     log::error!(
@@ -119,7 +122,7 @@ async fn main() -> Result<(), Error> {
                                 }
                             }
                         },
-                        None => Box::new(tcp_stream)
+                        None => Box::new(counting_stream)
                     };
                     if let Err(e) = chorus::serve(stream, hashed_peer).await {
                         log::error!(
diff --git a/src/counting_stream.rs b/src/counting_stream.rs
new file mode 100644
index 0000000..ef7debd
--- /dev/null
+++ b/src/counting_stream.rs
@@ -0,0 +1,53 @@
+use crate::globals::GLOBALS;
+use std::pin::Pin;
+use std::sync::atomic::Ordering;
+use std::task::{Context, Poll};
+use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};
+
+#[derive(Debug)]
+pub struct CountingStream<S>(pub S);
+
+impl<S: AsyncRead + Unpin> AsyncRead for CountingStream<S> {
+    fn poll_read(
+        self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+        buf: &mut ReadBuf<'_>,
+    ) -> Poll<std::io::Result<()>> {
+        // Count bytes for statistics
+        let pre = buf.filled().len();
+        let result = Pin::new(&mut self.get_mut().0).poll_read(cx, buf);
+        let post = buf.filled().len();
+        let count = post - pre;
+        if count > 0 {
+            let _ = GLOBALS
+                .bytes_inbound
+                .fetch_add(count as u64, Ordering::SeqCst);
+        }
+        result
+    }
+}
+
+impl<S: AsyncWrite + Unpin> AsyncWrite for CountingStream<S> {
+    fn poll_write(
+        self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+        buf: &[u8],
+    ) -> Poll<Result<usize, std::io::Error>> {
+        // Count bytes for statistics
+        let _ = GLOBALS
+            .bytes_outbound
+            .fetch_add(buf.len() as u64, Ordering::SeqCst);
+        Pin::new(&mut self.get_mut().0).poll_write(cx, buf)
+    }
+
+    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), std::io::Error>> {
+        Pin::new(&mut self.get_mut().0).poll_flush(cx)
+    }
+
+    fn poll_shutdown(
+        self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+    ) -> Poll<Result<(), std::io::Error>> {
+        Pin::new(&mut self.get_mut().0).poll_shutdown(cx)
+    }
+}
diff --git a/src/lib.rs b/src/lib.rs
index ce21da3..41ec5f7 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -1,4 +1,5 @@
 pub mod config;
+pub mod counting_stream;
 pub mod error;
 pub mod globals;
 pub mod ip;
@@ -8,6 +9,7 @@ pub mod tls;
 pub mod web;
 
 use crate::config::{Config, FriendlyConfig};
+use crate::counting_stream::CountingStream;
 use crate::error::{ChorusError, Error};
 use crate::globals::GLOBALS;
 use crate::ip::{HashedIp, HashedPeer, IpData, SessionExit};
@@ -42,8 +44,8 @@ use tungstenite::protocol::WebSocketConfig;
 use tungstenite::Message;
 
 pub trait FullStream: AsyncRead + AsyncWrite + Unpin + Send {}
-impl FullStream for TcpStream {}
-impl FullStream for TlsStream<TcpStream> {}
+impl FullStream for CountingStream<TcpStream> {}
+impl FullStream for TlsStream<CountingStream<TcpStream>> {}
 
 /// Serve a single network connection
 pub async fn serve(stream: Box<dyn FullStream>, peer: HashedPeer) -> Result<(), Error> {