btp 0.1.2

//! 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
//!
//! Example below is not using any safe transmission with TLS encryption,
//! but bare and crystal clear transmission without any security feature,
//! this type of communication may cause some
//!
//! ```
//! # {
//! type Error = Box<dyn std::error::Error + Send + Sync>;
//! use btp::server::BtpListener;
//! use btp::message::BtpMessage;
//! use btp::socket::{BtpConfig, BtpSocket};
//! use tokio::net::TcpListener;
//! // Add the items from libraries.
//!
//! #[tokio::main]
//! async fn main() -> Result<(), Error> {
//!     // Create a listener from a tcp listener for port 0.
//!     let listener = TcpListener::bind("127.0.0.1:0").await?;
//!     let server_addr = listener.local_addr()?;
//!     let btp_conf = BtpConfig::from_addr(server_addr);
//!
//!     // Spawn another async thread to handle incoming btp connection.
//!     let server_task = tokio::spawn(async move {
//!         // Assign that listener to a BtpListener.
//!         let btp_listener = BtpListener::from(listener, btp_conf).await;
//!         // Wait for someone to connect it.
//!         let mut server_socket = btp_listener.accept().await?;
//!
//!         // When connection established, wait for the incoming message.
//!         // And save it safely.
//!         let (_file_opt, msg_opt) = server_socket.read().await?;
//!         let msg = msg_opt.ok_or("No message received")?;
//!         // Message should be same with the one you sent.
//!         assert_eq!(msg.body, "Hello, BTP!");
//!
//!         // Response back with the confirmation message.
//!         let response = BtpMessage::from_str("ACK");
//!         server_socket.write_message(response).await?;
//!         Ok::<_, Error>(())
//!     });
//!
//!     // Create a config for client and connect to btp port you've listened.
//!     let client_conf = BtpConfig::from_addr(server_addr);
//!     let mut client_socket = BtpSocket::connect(client_conf).await?;
//!
//!     // Prepare the request and send it.
//!     let request = BtpMessage::from_str("Hello, BTP!");
//!     client_socket.write_message(request).await?;
//!
//!     // Wait for the response and handle it safely, then verify it.
//!     let (_file_opt, msg_opt) = client_socket.read().await?;
//!     let response = msg_opt.ok_or("No response received")?;
//!     assert_eq!(response.body, "ACK");
//!
//!     // Wait for server to finish jobs.
//!     server_task.await??;
//!     Ok(())
//! }
//!
//! # }
//! ```
//! ---
//! # With TLS
//!
//! ```no_run
//! type Error = Box<dyn std::error::Error + Send + Sync>;
//! use btp::message::BtpMessage;
//! use btp::server::BtpListenerTls;
//! use btp::socket::{BtpConfig, BtpSocket};
//! use rustls::ServerConfig;
//! use rustls::pki_types::{PrivateKeyDer, PrivatePkcs8KeyDer};
//! use std::sync::Arc;
//! use tokio::net::TcpListener;
//! use tokio_rustls::TlsConnector;
//!
//! fn make_test_tls_config() -> (Arc<ServerConfig>, Vec<u8>) { ... }
//!
//! #[tokio::main]
//! async fn main() -> Result<(), Error> {
//!     // Created a listener, then a TLS configurator and cert_der for conf.
//!     let listener = TcpListener::bind("127.0.0.1:0").await?;
//!     let server_addr = listener.local_addr()?;
//!     let btp_conf = BtpConfig::from_addr(server_addr);
//!     let (tls_conf, cert_der) = make_test_tls_config();
//!
//!     // Create a root certification store for client with cert_der you've created.
//!     let mut root_store = rustls::RootCertStore::empty();
//!     root_store.add(cert_der.into()).unwrap();
//!
//!     // Create the TLS config for client with that root store.
//!     let client_config = rustls::ClientConfig::builder()
//!         .with_root_certificates(root_store)
//!         .with_no_client_auth();
//!
//!     let server_task = tokio::spawn(async move {
//!         // Create a thread to assign listener to BtpListenerTls.
//!         let btp_listener = BtpListenerTls::from(listener, btp_conf, tls_conf).await;
//!
//!         // Wait for the client to connect.
//!         let mut server_socket = btp_listener.accept().await?;
//!
//!         // Wait for the message after client connects.
//!         let (_file_opt, msg_opt) = server_socket.read().await?;
//!
//!         // Handle it safely and validate it.
//!         let msg = msg_opt.ok_or("No message received")?;
//!         assert_eq!(msg.body, "Hello, TLS BTP!");
//!
//!         // Response to client back.
//!         let response = BtpMessage::from_str("TLS ACK");
//!         server_socket.write_message(response).await?;
//!
//!         Ok::<_, Error>(())
//!     });
//!
//!     // Create the tls_connector for client, then connect to server.
//!     let connector = TlsConnector::from(Arc::new(client_config));
//!     let client_conf = BtpConfig::from_addr(server_addr);
//!
//!     // Wait for connecting to server.
//!     let mut client_socket = BtpSocket::connect_tls(connector, client_conf).await?;
//!
//!     // Send message.
//!     let request = BtpMessage::from_str("Hello, TLS BTP!");
//!     client_socket.write_message(request).await?;
//!
//!     // Take the response back, handle it correctly and verify.
//!     let (_file_opt, msg_opt) = client_socket.read().await?;
//!     let response = msg_opt.ok_or("No response received")?;
//!     assert_eq!(response.body, "TLS ACK");
//!
//!     // Wait for server to complete its jobs.
//!     server_task.await??;
//!
//!     Ok(())
//! }
//!
//! ```
//!

/// For client connections.
#[cfg(feature = "client")]
pub mod client;

/// For parsing messages and files and it headers safely, without any issue.
pub mod message;

/// For server connections, brings [`BtpListener`].
///
/// [`BtpListener`]: server::BtpListener
#[cfg(feature = "server")]
pub mod server;

/// For transmitting messages and files safely.
pub mod socket;

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

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

    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 std::net::SocketAddr;

        use tokio::net::TcpStream;

        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 = BtpConfig::from_addr(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 = BtpConfig::from_addr(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 = BtpConfig::from_addr(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 = BtpConfig::from_addr(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 = BtpConfig::from_addr(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, _: BtpSocket<TcpStream, SocketAddr>| {
                        println!("Server received: {}", msg.body);
                    },
                    |_file: BtpFile, _: BtpSocket<TcpStream, SocketAddr>| {
                        println!("Server received a file");
                    },
                )?;

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

            let client_conf = BtpConfig::from_addr(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 = BtpConfig::from_addr(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 = BtpConfig::from_addr(server_addr);
            let mut client_socket = BtpSocket::connect_tls(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 = BtpConfig::from_addr(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 = BtpConfig::from_addr(server_addr);
            let mut client_socket = BtpSocket::connect_tls(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(())
        }
    }
}