Rust 實現 NTP 客戶端 - 2

在《Rust 實現 NTP 客戶端 - 1》部分中，我們創建了從 NTP 服務器獲取時間的基本網絡調用，這次我們將加固代碼，考慮各種異常情況，並添加單元測試。

讓我們從單元測試開始：

#[cfg(test)]
mod main_tests {
    use super::*;
    use byteorder::{BigEndian, ReadBytesExt};
    use std::io::{Cursor, Seek, SeekFrom};

    #[test]
    fn unable_to_bind() {
        let result = ntp_main("0.0.0.0:80", "pool.ntp.org:123");
        assert!(result.is_err());
    }
}

在代碼中，ntp_main() 是新的函數。直接測試 main 函數是很尷尬的，因爲我們需要提供命令行參數等。將邏輯代碼移到一個函數中，在 main() 中只保留調用邏輯函數的代碼，是非常易於測試的：

fn main() {
    let time = ntp_main("0.0.0.0:0", "pool.ntp.org:123").unwrap();
    println!("Time is {time}");
}

fn ntp_main(bind_address: &str, ntp_server: &str) -> Result<u64, std::io::Error> {
    let socket = UdpSocket::bind(bind_address)?;
    socket.set_write_timeout(Some(Duration::from_millis(500)))?;
    socket.set_read_timeout(Some(Duration::from_millis(500)))?;

    let mut transmit: Vec<u8> = vec![0; 48];
    transmit[0] = 0x1b;

    let mut retries = 3;
    while retries > 0 {
        retries -= 1;
        let _bytes_transmitted = match socket.send_to(&transmit, ntp_server) {
            Ok(bytes) => bytes,
            Err(_) => continue,
        };
        let mut buf = [0; 48];
        let _bytes_received = match socket.recv(&mut buf) {
            Ok(bytes) => bytes,
            Err(_) => continue,
        };
        let ntp_time = process_ntp_packet(&buf);
        if ntp_time > 2208988800 {
            let unix_time: u64 = ntp_time - 2208988800;
            return Ok(unix_time)
        }else {
            return Ok(0)
        }
    }

    Err(std::io::Error::new(
        std::io::ErrorKind::TimedOut,
        "Timed out getting response from server",
    ))
}

fn process_ntp_packet(buffer: &[u8]) -> u64 {
    let mut reader = Cursor::new(buffer);
    reader.seek(SeekFrom::Start(40)).unwrap();

    let transmit_timestamp_seconds = reader.read_u32::<BigEndian>().unwrap();

    u64::from(transmit_timestamp_seconds)
}

與上次相比有很多變化，首先 ntp_main() 函數返回了一個 Result<u64, std::io::Error>，它允許我們在發生錯誤時通知調用函數。

第 8 行和第 9 行對網絡讀寫操作設置了合理的超時時間。第 15 行是一個 while 循環，允許函數在發生錯誤，退出之前嘗試 3 次以獲得響應。

在第 35 行，我們使用 std::io::Error::new() 函數返回 Err；雖然，我們不知道是什麼導致了錯誤，但超時是一個很好的猜測。

現在讓我們添加一些更多的測試：

#[test]
fn response_timeout() {
    let result = ntp_main("0.0.0.0:0", "google.com:21");
    assert_eq!(result.unwrap_err().kind(), std::io::ErrorKind::TimedOut);
}

這一個嘗試連接 google.com:21 (FTP)，我們預計會出現 TimeOut 錯誤，我們使用 assert_eq! 宏。

我們也可以做更復雜的測試：

#[test]
fn correct_response() {
    let test_socket = UdpSocket::bind("0.0.0.0:0").unwrap();
    let test_port = test_socket.local_addr().unwrap().port();

    let join_handle = std::thread::spawn(move || {
        let mut buf = [0; 48];
        let (_, sender) = test_socket.recv_from(&mut buf).unwrap();
        let mut cur = Cursor::new(vec![0; 48]);
        cur.seek(SeekFrom::Start(40)).unwrap();
        cur.write_u32::<BigEndian>(2208988800).unwrap();
        test_socket.send_to(cur.get_ref(), sender).unwrap();
    });

    ntp_main("0.0.0.0:0",format!("127.0.0.1:                      {}", test_port).as_str()).unwrap();

    join_handle.join().unwrap();
}

我們將綁定到端口 0，該端口分配給我們一個隨機未使用的端口，我們在第 4 行提取該端口。

然後在第 6 行，我們創建了一個線程，這意味着 test_socket 被移動到新線程，當新線程退出時，對象將被釋放。

線程偵聽套接字上的一個包，然後用一個新包響應給發送方。

我們測試包解碼器得到正確的字節：

#[test]
fn packet_decoder() {
    let mut cur = Cursor::new(vec![0; 48]);
    cur.seek(SeekFrom::Start(40)).unwrap();
    cur.write_u32::<BigEndian>(2208988800).unwrap();
    assert_eq!(process_ntp_packet(cur.get_ref()), 2208988800);
    assert_ne!(process_ntp_packet(cur.get_ref()), 10);
}

還有一個針對各種失敗模式的測試，它使用與 correct_response() 測試是相同的機制，但針對不同的測試用例重複了三次：

#[test]
fn incorrect_response() {
    let test_socket = UdpSocket::bind("0.0.0.0:0").unwrap();
    let test_port = test_socket.local_addr().unwrap().port();

    let join_handle = std::thread::spawn(move || {
        let mut buf = [0; 48];
        // First test - blank response
        let (_, sender) = test_socket.recv_from(&mut buf).unwrap();
        let mut tx = [0; 48];
        test_socket.send_to(&mut tx, sender).unwrap();
        // Second test - too short
        let (_, sender) = test_socket.recv_from(&mut buf).unwrap();
        let mut tx = [0; 8];
        test_socket.send_to(&mut tx, sender).unwrap();
        // Third test - too long
        let (_, sender) = test_socket.recv_from(&mut buf).unwrap();
        let mut tx = [0; 200];
        test_socket.send_to(&mut tx, sender).unwrap();
    });

    // Test 1 - blank response
    ntp_main("0.0.0.0:0", format!("127.0.0.1:                        {}", test_port).as_str()).unwrap();

    // Test 2 - too short
    ntp_main("0.0.0.0:0", format!("127.0.0.1:                        {}", test_port).as_str()).unwrap();

    // Test 3 - too long
    ntp_main("0.0.0.0:0", format!("127.0.0.1:                        {}", test_port).as_str()).unwrap();

    join_handle.join().unwrap();
}

運行 cargo test ：

master:ntp_part1 Justin$ cargo test
    Finished test [unoptimized + debuginfo] target(s) in 0.07s
     Running unittests src/main.rs (target/debug/deps/ntp_part1-59f5a83e3e1bfa9a)
running 5 tests
test main_tests::packet_decoder ... ok
test main_tests::correct_response ... ok
test main_tests::incorrect_response ... ok
test main_tests::unable_to_bind ... ok
test main_tests::response_timeout ... ok
test result: ok. 5 passed; 0 failed; 0 ignored; 0 measured; 0 filtered out; finished in 1.52s

本文翻譯自：

https://medium.com/towardsdev/rust-foo-ntp-client-part-2-e82832bfaf03

coding 到燈火闌珊 專注於技術分享，包括 Rust、Golang、分佈式架構、雲原生等。

本文由 Readfog 進行 AMP 轉碼，版權歸原作者所有。
來源：https://mp.weixin.qq.com/s/rGit8D7FQCCG0wIzmcucDQ