borer-core 0.5.9

use std::{
    io,
    pin::Pin,
    task::{Context, Poll},
};

use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};

use crate::store::store;

/// Stream wrapper that records per-connection traffic in the shared store.
#[derive(Debug)]
pub struct StatsStream<S> {
    inner: S,
    conn_id: String,
}

impl<S> StatsStream<S> {
    /// Create a tracked stream and register it in the shared traffic store.
    pub fn new(
        inner: S,
        user: impl Into<String>,
        hash: impl Into<String>,
        peer_addr: impl Into<String>,
        req_addr: impl Into<String>,
        padding: bool,
    ) -> Self {
        let conn_id = gen_conn_id();
        let user: String = user.into();
        let hash: String = hash.into();
        let peer_addr: String = peer_addr.into();
        let req_addr: String = req_addr.into();

        store().insert_conn(&user, &hash, &conn_id, &peer_addr, &req_addr, padding);
        StatsStream { inner, conn_id }
    }
}

impl<S> Drop for StatsStream<S> {
    fn drop(&mut self) {
        store().delete_conn(&self.conn_id);
    }
}

impl<S> AsyncRead for StatsStream<S>
where
    S: AsyncRead + Unpin,
{
    fn poll_read(
        self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        buf: &mut ReadBuf,
    ) -> Poll<io::Result<()>> {
        let conn_id = self.conn_id.clone();
        let me = self.get_mut();
        let filled_before = buf.filled().len();
        let poll = Pin::new(&mut me.inner).poll_read(cx, buf);
        if let Poll::Ready(Ok(())) = poll {
            let filled_after = buf.filled().len();
            let bytes_read = filled_after.saturating_sub(filled_before);
            store().add_up(&conn_id, bytes_read);
            debug!("stats => read buf: {bytes_read}");
        }
        poll
    }
}

impl<S> AsyncWrite for StatsStream<S>
where
    S: AsyncWrite + Unpin,
{
    fn poll_write(
        mut self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        buf: &[u8],
    ) -> Poll<io::Result<usize>> {
        let conn_id = self.conn_id.clone();
        debug!("stats => write buf: {}", buf.len());
        match Pin::new(&mut self.inner).poll_write(cx, buf) {
            Poll::Ready(Ok(bytes_written)) => {
                store().add_down(&conn_id, bytes_written);
                Poll::Ready(Ok(bytes_written))
            }
            other => other,
        }
    }

    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
        Pin::new(&mut self.inner).poll_flush(cx)
    }

    fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
        Pin::new(&mut self.inner).poll_shutdown(cx)
    }
}

pub(crate) fn gen_conn_id() -> String {
    let uuid = uuid::Uuid::new_v4().as_u128();
    format!("{:032x}", uuid)
}

#[cfg(test)]
mod tests {
    use std::time::{SystemTime, UNIX_EPOCH};

    use std::{
        io,
        pin::Pin,
        task::{Context, Poll},
    };

    use tokio::io::{AsyncReadExt, AsyncWrite, AsyncWriteExt};

    use super::{StatsStream, gen_conn_id};
    use crate::store::store;

    #[test]
    fn gen_conn_id_is_32_char_hex() {
        let id = gen_conn_id();

        assert_eq!(id.len(), 32);
        assert!(id.bytes().all(|b| b.is_ascii_hexdigit()));
        assert!(id.bytes().all(|b| !b.is_ascii_uppercase()));
    }

    #[tokio::test]
    async fn stats_stream_tracks_read_write_and_drop() {
        let store = store();
        let suffix = SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .unwrap()
            .as_nanos();
        let user = format!("alice-{suffix}");
        let hash = format!("hash-{suffix}");
        let (client, mut peer) = tokio::io::duplex(64);
        let mut stream = StatsStream::new(
            client,
            user.clone(),
            hash.clone(),
            "127.0.0.1:1".to_string(),
            "example.com:443".to_string(),
            false,
        );

        let conn = only_conn(&store, &user);
        assert_eq!(conn.user, user);
        assert_eq!(conn.hash, hash);

        stream.write_all(b"hello").await.unwrap();
        let mut downstream = [0u8; 5];
        peer.read_exact(&mut downstream).await.unwrap();
        assert_eq!(&downstream, b"hello");

        peer.write_all(b"world").await.unwrap();
        let mut upstream = [0u8; 5];
        stream.read_exact(&mut upstream).await.unwrap();
        assert_eq!(&upstream, b"world");

        let conn = only_conn(&store, &conn.user);
        assert_eq!(conn.traffic.down, 5);
        assert_eq!(conn.traffic.up, 5);
        let traffic = store.get_traffic_by_user(&conn.user).unwrap();
        assert_eq!(traffic.down, 5);
        assert_eq!(traffic.up, 5);

        drop(stream);
        assert_eq!(store.get_conns_by_user(&conn.user).unwrap().len(), 0);
    }

    #[tokio::test]
    async fn stats_stream_counts_only_bytes_written_by_inner_stream() {
        let store = store();
        let suffix = SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .unwrap()
            .as_nanos();
        let user = format!("writer-{suffix}");
        let hash = format!("hash-{suffix}");
        let mut stream = StatsStream::new(
            PartialWriter { max_write: 2 },
            user.clone(),
            hash,
            "127.0.0.1:1".to_string(),
            "example.com:443".to_string(),
            false,
        );

        let bytes_written = stream.write(b"hello").await.unwrap();

        assert_eq!(bytes_written, 2);
        let conn = only_conn(&store, &user);
        assert_eq!(conn.traffic.down, 2);
    }

    fn only_conn(
        store: &std::sync::Arc<crate::store::Store>,
        user: &str,
    ) -> crate::store::Connection {
        let mut conns = store.get_conns_by_user(user).unwrap();
        assert_eq!(conns.len(), 1);
        conns.pop().unwrap()
    }

    struct PartialWriter {
        max_write: usize,
    }

    impl AsyncWrite for PartialWriter {
        fn poll_write(
            self: Pin<&mut Self>,
            _cx: &mut Context<'_>,
            buf: &[u8],
        ) -> Poll<io::Result<usize>> {
            Poll::Ready(Ok(buf.len().min(self.max_write)))
        }

        fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<io::Result<()>> {
            Poll::Ready(Ok(()))
        }

        fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<io::Result<()>> {
            Poll::Ready(Ok(()))
        }
    }
}