pfs-tftp-claude/crates/pfs-tftp/tests/integration.rs

//! Integration tests for TFTP client and server.
//!
//! These tests verify that the client and server can communicate correctly
//! by running an actual server in a background thread and making client requests.

use std::{
    fs,
    io::Write,
    net::SocketAddr,
    path::PathBuf,
    sync::{
        Arc,
        atomic::{AtomicBool, Ordering},
    },
    thread,
    time::Duration,
};

use pfs_tftp::{Client, Mode, ServerBuilder};

/// Create a temporary directory for test files.
fn create_test_dir() -> PathBuf {
    use std::sync::atomic::AtomicU64;
    static COUNTER: AtomicU64 = AtomicU64::new(0);
    let id = COUNTER.fetch_add(1, Ordering::Relaxed);
    let dir = std::env::temp_dir().join(format!(
        "tftp_test_{}_{}_{}",
        std::process::id(),
        id,
        std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .map(|d| d.as_nanos())
            .unwrap_or(0)
    ));
    if dir.exists() {
        let _ = fs::remove_dir_all(&dir);
    }
    fs::create_dir_all(&dir).expect("create test dir");
    dir
}

/// Clean up the test directory.
fn cleanup_test_dir(dir: &PathBuf) {
    let _ = fs::remove_dir_all(dir);
}

/// Find an available port for the test server.
fn find_available_port() -> u16 {
    let socket = std::net::UdpSocket::bind("127.0.0.1:0").expect("bind");
    socket.local_addr().expect("local_addr").port()
}

/// Test context that manages server and client setup.
struct TestContext {
    server_dir: PathBuf,
    client_dir: PathBuf,
    server_addr: SocketAddr,
    running: Arc<AtomicBool>,
}

impl TestContext {
    fn new(writable: bool) -> Self {
        let server_dir = create_test_dir();
        let client_dir = create_test_dir();
        let port = find_available_port();
        let server_addr: SocketAddr = ([127, 0, 0, 1], port).into();

        let running = Arc::new(AtomicBool::new(true));
        let running_clone = running.clone();
        let server_dir_clone = server_dir.clone();

        // Start server in background thread
        thread::spawn(move || {
            let server = ServerBuilder::new(&server_dir_clone)
                .allow_write(writable)
                .allow_overwrite(writable)
                .timeout(Duration::from_secs(1))
                .retries(2)
                .bind(server_addr)
                .expect("bind server");

            while running_clone.load(Ordering::Relaxed) {
                // Use handle_one with a short timeout to allow checking the running flag
                let _ = server.handle_one();
            }
        });

        // Give server time to start
        thread::sleep(Duration::from_millis(50));

        Self {
            server_dir,
            client_dir,
            server_addr,
            running,
        }
    }

    fn client(&self) -> Client {
        Client::new(self.server_addr)
            .expect("create client")
            .with_timeout(Duration::from_secs(2))
            .with_retries(3)
    }

    fn create_server_file(&self, name: &str, content: &[u8]) -> PathBuf {
        let path = self.server_dir.join(name);
        let mut file = fs::File::create(&path).expect("create file");
        file.write_all(content).expect("write file");
        path
    }

    fn read_server_file(&self, name: &str) -> Vec<u8> {
        let path = self.server_dir.join(name);
        fs::read(path).expect("read file")
    }

    fn server_file_exists(&self, name: &str) -> bool {
        self.server_dir.join(name).exists()
    }
}

impl Drop for TestContext {
    fn drop(&mut self) {
        self.running.store(false, Ordering::Relaxed);
        // Small delay to let server thread exit
        thread::sleep(Duration::from_millis(100));
        cleanup_test_dir(&self.server_dir);
        cleanup_test_dir(&self.client_dir);
    }
}

#[test]
fn test_get_small_file() {
    let ctx = TestContext::new(false);
    let content = b"Hello, TFTP!";
    ctx.create_server_file("hello.txt", content);

    let client = ctx.client();
    let received = client.get("hello.txt", Mode::Octet).expect("get");

    assert_eq!(received, content);
}

#[test]
fn test_get_exact_block_size() {
    let ctx = TestContext::new(false);
    // Exactly 512 bytes = one full block + empty final block
    let content = vec![0xAB; 512];
    ctx.create_server_file("exact.bin", &content);

    let client = ctx.client();
    let received = client.get("exact.bin", Mode::Octet).expect("get");

    assert_eq!(received, content);
}

#[test]
fn test_get_large_file() {
    let ctx = TestContext::new(false);
    // Multi-block file: 2.5 blocks
    let content: Vec<u8> = (0..1280).map(|i| (i % 256) as u8).collect();
    ctx.create_server_file("large.bin", &content);

    let client = ctx.client();
    let received = client.get("large.bin", Mode::Octet).expect("get");

    assert_eq!(received.len(), content.len());
    assert_eq!(received, content);
}

#[test]
fn test_get_empty_file() {
    let ctx = TestContext::new(false);
    ctx.create_server_file("empty.txt", b"");

    let client = ctx.client();
    let received = client.get("empty.txt", Mode::Octet).expect("get");

    assert!(received.is_empty());
}

#[test]
fn test_get_file_not_found() {
    let ctx = TestContext::new(false);

    let client = ctx.client();
    let result = client.get("nonexistent.txt", Mode::Octet);

    assert!(result.is_err());
    let err = result.unwrap_err();
    assert!(err.to_string().contains("not found") || err.to_string().contains("Remote error"));
}

#[test]
fn test_put_small_file() {
    let ctx = TestContext::new(true);
    let content = b"Hello from client!";

    let client = ctx.client();
    client
        .put("uploaded.txt", Mode::Octet, content)
        .expect("put");

    // Give server time to flush the file
    thread::sleep(Duration::from_millis(50));

    assert!(ctx.server_file_exists("uploaded.txt"));
    assert_eq!(ctx.read_server_file("uploaded.txt"), content);
}

#[test]
fn test_put_large_file() {
    let ctx = TestContext::new(true);
    // Multi-block file
    let content: Vec<u8> = (0..2000).map(|i| (i % 256) as u8).collect();

    let client = ctx.client();
    client
        .put("big_upload.bin", Mode::Octet, &content)
        .expect("put");

    // Give server time to flush the file
    thread::sleep(Duration::from_millis(50));

    assert!(ctx.server_file_exists("big_upload.bin"));
    assert_eq!(ctx.read_server_file("big_upload.bin"), content);
}

#[test]
fn test_put_write_denied() {
    let ctx = TestContext::new(false); // Write not allowed

    let client = ctx.client();
    let result = client.put("test.txt", Mode::Octet, b"data");

    assert!(result.is_err());
}

#[test]
fn test_roundtrip() {
    let ctx = TestContext::new(true);
    let original = b"This is test data for roundtrip!";

    // Upload - use a fresh client
    let upload_client = ctx.client();
    upload_client
        .put("roundtrip.txt", Mode::Octet, original)
        .expect("put");

    // Give server time to flush the file
    thread::sleep(Duration::from_millis(50));

    // Download - use a fresh client to avoid socket state issues
    let download_client = ctx.client();
    let downloaded = download_client
        .get("roundtrip.txt", Mode::Octet)
        .expect("get");

    assert_eq!(downloaded, original);
}

#[test]
fn test_path_traversal_blocked() {
    let ctx = TestContext::new(false);

    let client = ctx.client();
    let result = client.get("../../../etc/passwd", Mode::Octet);

    assert!(result.is_err());
}

#[test]
fn test_netascii_mode() {
    let ctx = TestContext::new(false);
    let content = b"Line 1\nLine 2\n";
    ctx.create_server_file("text.txt", content);

    let client = ctx.client();
    // NetASCII mode should work for simple ASCII text
    let received = client.get("text.txt", Mode::NetAscii).expect("get");

    // Content should be transferred (exact handling depends on implementation)
    assert!(!received.is_empty());
}