borer-core 0.5.9

use anyhow::Context as _;
use tokio::{
    io::{AsyncRead, AsyncWrite, AsyncWriteExt as _, copy_bidirectional},
    net::TcpStream,
};

/// Handles connections that should receive a masquerade response instead of proxying.
pub struct MasqueradeConnection<T> {
    inner: T,
}

impl<T> MasqueradeConnection<T>
where
    T: AsyncRead + AsyncWrite + Unpin,
{
    /// Create a masquerade connection wrapper around the underlying stream.
    pub fn new(ts: T) -> Self {
        Self { inner: ts }
    }

    /// Send the configured masquerade response or a default 404 response.
    pub async fn handle(mut self) -> anyhow::Result<()> {
        let stream = &mut self.inner;
        let local_addr: Option<String> = None;

        match local_addr {
            Some(ref local_addr) => {
                info!("requested proxy local {}", local_addr);
                let mut local_ts = TcpStream::connect(local_addr)
                    .await
                    .context(format!("proxy can't connect addr {}", local_addr))?;
                debug!("proxy connect success {:?}", local_addr);

                copy_bidirectional(stream, &mut local_ts).await.ok();
            }
            None => {
                info!("response not found");
                Self::resp_html(stream).await
            }
        }

        Ok(())
    }
    async fn resp_html(stream: &mut T) {
        stream
            .write_all(
                &b"HTTP/1.0 404 Not Found\r\n\
                Content-Type: text/plain; charset=utf-8\r\n\
                Content-length: 13\r\n\r\n\
                404 Not Found"[..],
            )
            .await
            .unwrap();
    }
}

#[cfg(test)]
mod tests {
    use tokio::io::AsyncReadExt;

    use super::MasqueradeConnection;

    #[tokio::test]
    async fn handle_without_local_target_returns_404_response() {
        let (client, mut peer) = tokio::io::duplex(256);

        let task = tokio::spawn(async move { MasqueradeConnection::new(client).handle().await });

        let mut buf = Vec::new();
        peer.read_to_end(&mut buf).await.unwrap();
        task.await.unwrap().unwrap();

        let response = String::from_utf8(buf).unwrap();
        assert!(response.starts_with("HTTP/1.0 404 Not Found\r\n"));
        assert!(response.contains("Content-Type: text/plain; charset=utf-8\r\n"));
        assert!(response.ends_with("404 Not Found"));
    }
}