btp 0.1.1

//! src/lib.rs

//! # Welcome
//! Btp is stands for 'Blog Transfer Protocol' which developed while developing this project.
//! The protocol has a basic structure, a 'header' which consists of 10 bytes, and the rest of the message, which can be reach 4GiB.
//!
//! This protocol handles a basic transmission task between server and client.
//! To add this project to repository
//! ```bash
//! cargo add btp
//! ```
//! or with TLS,
//! ```bash
//! cargo add btp --features tls
//! ```
//!
//! After you added the library to project, It will be accessible under the crate name btp.
//!
//! # Examples
//! Let me teach the structure with examples here;
//!
//!
//! ## NO TLS
//!
//! Examples below will not use any safe transmission with TLS encryption,
//! but bare and crystal clear transmission.
//!
//! ### Basic client example without TSL;
//!
//! ```no_run
//! use btp::message::BtpMessage;
//! use btp::socket::{BtpConfig, BtpSocket};
//!
//! #[tokio::main]
//! async fn main() {
//!     let address = "127.0.0.1:8080".parse().unwrap();
//!     // First we created a basic client config, includes our server address,
//!     // and the VERSION of the btp protocol used.
//!     let client_conf = BtpConfig::from_addr(address);
//!     // Connected to given address using conf, and awaited since it's an async connection.
//!     let mut client_socket = BtpSocket::connect(client_conf).await.unwrap();
//!
//!     // Created a message as [&str] and converted it into request.
//!     let body: &str = "Hello, BTP!";
//!     let msg = BtpMessage::from_str(body);
//!
//!     // End sended to client (I guess)
//!     client_socket.write_message(msg).await.unwrap();
//! }
//! ```
//!
//! ### Basic server example without TSL;
//!
//! ```no_run
//! use btp::message::BtpMessage;
//! use btp::socket::{BtpConfig, BtpSocket};
//!
//! #[tokio::main]
//! async fn main() {
//!     let address = "127.0.0.1:8080".parse().unwrap();
//!     // First we created a basic client config, includes our server address,
//!     // and the VERSION of the btp protocol used.
//!     let client_conf = BtpConfig::from_addr(address);
//!     // Connected to given address using conf, and awaited since it's an async connection.
//!     let mut client_socket = BtpSocket::connect(client_conf).await.unwrap();
//!
//!     // Created a message as [&str] and converted it into request.
//!     let body: &str = "Hello, BTP!";
//!     let msg = BtpMessage::from_str(body);
//!
//!     // End sended to client (I guess)
//!     client_socket.write_message(msg).await.unwrap();
//! }
//! ```

pub mod client;
pub mod message;
pub mod server;
pub mod socket;

/// Actual current version of the protocol. Used in [`socket::BtpConfig`]
pub const VERSION: [u8; 3] = [0, 1, 1];

#[cfg(test)]
mod tests {
    use super::*;
    use crate::message::{BtpFile, BtpMessage, StatusCode};
    use crate::socket::{BtpConfig, BtpSocket};
    use std::io::Write;
    use std::net::SocketAddr;
    use tempfile::tempdir;
    use tokio::net::TcpListener;

    fn make_config(addr: SocketAddr) -> BtpConfig {
        BtpConfig { ver: VERSION, addr }
    }

    fn make_test_file(content: &str, name: &str) -> BtpFile {
        let mut temp_file = tempfile::NamedTempFile::new().unwrap();
        write!(temp_file, "{}", content).unwrap();
        let file = temp_file.reopen().unwrap();
        BtpFile::from([StatusCode::File as u8], VERSION, file, name.to_string()).unwrap()
    }

    mod tests {
        use super::*;
        use crate::server::BtpListener;
        type Error = Box<dyn std::error::Error + Send + Sync>;
        #[tokio::test]
        async fn test_message_exchange_no_tls() -> Result<(), Error> {
            let listener = TcpListener::bind("127.0.0.1:0").await?;
            let server_addr = listener.local_addr()?;
            let btp_conf = make_config(server_addr);
            let server_task = tokio::spawn(async move {
                let btp_listener = BtpListener::from(listener, btp_conf).await;
                let mut server_socket = btp_listener.accept().await?;
                let (_file_opt, msg_opt) = server_socket.read().await?;
                let msg = msg_opt.ok_or("No message received")?;
                assert_eq!(msg.body, "Hello, BTP!");
                let response = BtpMessage::from_str("ACK");
                server_socket.write_message(response).await?;

                Ok::<_, Error>(())
            });
            let client_conf = make_config(server_addr);
            let mut client_socket = BtpSocket::connect(client_conf).await?;

            let request = BtpMessage::from_str("Hello, BTP!");
            client_socket.write_message(request).await?;

            let (_file_opt, msg_opt) = client_socket.read().await?;
            let response = msg_opt.ok_or("No response received")?;
            assert_eq!(response.body, "ACK");

            server_task.await??;

            Ok(())
        }

        #[tokio::test]
        async fn test_file_transfer_no_tls() -> Result<(), Error> {
            let tmp_dir = tempdir()?;
            let file_path = tmp_dir.path().join("test.txt");
            let original_content = "Hello, this is a test file.";
            tokio::fs::write(&file_path, original_content).await?;

            let listener = TcpListener::bind("127.0.0.1:0").await?;
            let server_addr = listener.local_addr()?;
            let btp_conf = make_config(server_addr);

            let server_task = tokio::spawn(async move {
                let btp_listener = BtpListener::from(listener, btp_conf).await;
                let mut server_socket = btp_listener.accept().await?;

                let (file_opt, _msg_opt) = server_socket.read().await?;
                let file = file_opt.ok_or("No file received")?;
                assert_eq!(file.file_content, original_content);

                let ack = BtpMessage::from_str("FILE_RECEIVED");
                server_socket.write_message(ack).await?;

                Ok::<_, Error>(())
            });

            let client_conf = make_config(server_addr);
            let mut client_socket = BtpSocket::connect(client_conf).await?;

            let file = make_test_file(original_content, "test.txt");
            client_socket.write_file(file).await?;

            let (_file_opt, msg_opt) = client_socket.read().await?;
            let ack = msg_opt.ok_or("No acknowledgment received")?;
            assert_eq!(ack.body, "FILE_RECEIVED");

            server_task.await??;

            Ok(())
        }

        #[tokio::test]
        async fn test_full_session_no_tls() -> Result<(), Error> {
            let listener = TcpListener::bind("127.0.0.1:0").await?;
            let server_addr = listener.local_addr()?;
            let btp_conf = make_config(server_addr);

            let server_task = tokio::spawn(async move {
                let btp_listener = BtpListener::from(listener, btp_conf).await;
                let server_socket = btp_listener.accept().await?;

                server_socket.attach_handler(
                    |msg: BtpMessage| {
                        println!("Server received: {}", msg.body);
                    },
                    |_file: BtpFile| {
                        println!("Server received a file");
                    },
                )?;

                tokio::time::sleep(tokio::time::Duration::from_millis(100)).await;
                Ok::<_, Error>(())
            });

            let client_conf = make_config(server_addr);
            let mut client_socket = BtpSocket::connect(client_conf).await?;

            let msg = BtpMessage::from_str("Hello, echo!");
            client_socket.write_message(msg).await?;

            tokio::time::sleep(tokio::time::Duration::from_millis(50)).await;

            server_task.await??;

            Ok(())
        }
    }

    #[cfg(feature = "tls")]
    mod tls_tests {

        type Error = Box<dyn std::error::Error + Send + Sync>;
        use crate::server::BtpListenerTls;

        use super::*;
        use rustls::ServerConfig;
        use rustls::pki_types::{PrivateKeyDer, PrivatePkcs8KeyDer};
        use std::sync::Arc;
        use tokio_rustls::TlsConnector;

        fn make_test_tls_config() -> (Arc<ServerConfig>, Vec<u8>) {
            let certified = rcgen::generate_simple_self_signed(vec!["127.0.0.1".into()]).unwrap();

            let cert_der = certified.cert.der().to_vec();

            let server_cert = rustls::pki_types::CertificateDer::from(certified.cert.der().clone());

            let key_der = PrivateKeyDer::Pkcs8(PrivatePkcs8KeyDer::from(
                certified.signing_key.serialize_der(),
            ));

            let config = ServerConfig::builder()
                .with_no_client_auth()
                .with_single_cert(vec![server_cert], key_der)
                .unwrap();

            (Arc::new(config), cert_der)
        }

        #[tokio::test]
        async fn test_message_exchange_with_tls() -> Result<(), Error> {
            let listener = TcpListener::bind("127.0.0.1:0").await?;
            let server_addr = listener.local_addr()?;
            let btp_conf = make_config(server_addr);
            let (tls_conf, cert_der) = make_test_tls_config();

            let mut root_store = rustls::RootCertStore::empty();
            root_store.add(cert_der.into()).unwrap();

            let client_config = rustls::ClientConfig::builder()
                .with_root_certificates(root_store)
                .with_no_client_auth();

            let server_task = tokio::spawn(async move {
                let btp_listener = BtpListenerTls::from(listener, btp_conf, tls_conf).await;
                let mut server_socket = btp_listener.accept().await?;

                let (_file_opt, msg_opt) = server_socket.read().await?;
                let msg = msg_opt.ok_or("No message received")?;
                assert_eq!(msg.body, "Hello, TLS BTP!");

                let response = BtpMessage::from_str("TLS ACK");
                server_socket.write_message(response).await?;

                Ok::<_, Error>(())
            });

            let connector = TlsConnector::from(Arc::new(client_config));
            let client_conf = make_config(server_addr);
            let mut client_socket =
                BtpSocket::connect_tls(server_addr, connector, client_conf).await?;

            let request = BtpMessage::from_str("Hello, TLS BTP!");
            client_socket.write_message(request).await?;

            let (_file_opt, msg_opt) = client_socket.read().await?;
            let response = msg_opt.ok_or("No response received")?;
            assert_eq!(response.body, "TLS ACK");

            server_task.await??;

            Ok(())
        }

        #[tokio::test]
        async fn test_file_transfer_with_tls() -> Result<(), Error> {
            let tmp_dir = tempdir()?;
            let file_path = tmp_dir.path().join("tls_test.txt");
            let original_content = "TLS file transfer test content.";
            tokio::fs::write(&file_path, original_content).await?;

            let listener = TcpListener::bind("127.0.0.1:0").await?;
            let server_addr = listener.local_addr()?;
            let btp_conf = make_config(server_addr);
            let (tls_conf, cert_der) = make_test_tls_config();

            let mut root_store = rustls::RootCertStore::empty();
            root_store.add(cert_der.into()).unwrap();

            let client_config = rustls::ClientConfig::builder()
                .with_root_certificates(root_store)
                .with_no_client_auth();

            let server_task = tokio::spawn(async move {
                let btp_listener = BtpListenerTls::from(listener, btp_conf, tls_conf).await;
                let mut server_socket = btp_listener.accept().await?;

                let (file_opt, _msg_opt) = server_socket.read().await?;
                let file = file_opt.ok_or("No file received")?;
                assert_eq!(file.file_content, original_content);

                let ack = BtpMessage::from_str("TLS FILE RECEIVED");
                server_socket.write_message(ack).await?;

                Ok::<_, Error>(())
            });

            let connector = TlsConnector::from(Arc::new(client_config));
            let client_conf = make_config(server_addr);
            let mut client_socket =
                BtpSocket::connect_tls(server_addr, connector, client_conf).await?;

            let file = make_test_file(original_content, "tls_test.txt");
            client_socket.write_file(file).await?;

            let (_file_opt, msg_opt) = client_socket.read().await?;
            let ack = msg_opt.ok_or("No acknowledgment received")?;
            assert_eq!(ack.body, "TLS FILE RECEIVED");

            server_task.await??;

            Ok(())
        }
    }
}