pub mod client;
pub mod message;
pub mod server;
pub mod socket;
pub const VERSION: [u8; 3] = [0, 1, 0];
#[cfg(test)]
mod tests {
use super::*;
use crate::message::{BtpFile, BtpMessage, StatusCode};
use crate::server::BtpListener;
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_message(body: &str) -> BtpMessage {
let len = body.len() as u32;
BtpMessage::from(
len.to_be_bytes(),
[StatusCode::Message as u8],
VERSION,
body.to_string(),
)
}
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()
}
#[cfg(not(feature = "tls"))]
mod tls_tests {
use super::*;
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 = make_test_message("ACK");
server_socket.write_message(response).await?;
Ok::<_, Error>(())
});
let client_conf = make_config(server_addr);
let mut client_socket = BtpSocket::connect(server_addr, client_conf).await?;
let request = make_test_message("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 = make_test_message("FILE_RECEIVED");
server_socket.write_message(ack).await?;
Ok::<_, Error>(())
});
let client_conf = make_config(server_addr);
let mut client_socket = BtpSocket::connect(server_addr, 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(server_addr, client_conf).await?;
let msg = make_test_message("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 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 = BtpListener::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 = make_test_message("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(server_addr, connector, client_conf).await?;
let request = make_test_message("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 = BtpListener::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 = make_test_message("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(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(())
}
}
}