dns-server 0.2.6

use std::net::{IpAddr, Ipv6Addr, SocketAddr, UdpSocket};
use std::process::Command;
use std::time::Duration;

use permit::Permit;

use dns_server::{Builder, DnsRecord};

// TODO: Test rate limiting.

#[test]
fn example() {
    let top_permit = Permit::new();
    let permit = top_permit.new_sub();
    let (builder, addr) = Builder::new_random_port().unwrap();
    let records = vec![
        DnsRecord::new_a("name1", "1.2.3.4").unwrap(),
        DnsRecord::new_aaaa("name2", "2606:2800:220:1:248:1893:25c8:1946").unwrap(),
        DnsRecord::new_a("name3", "10.0.0.1").unwrap(),
        DnsRecord::new_a("name3", "10.0.0.2").unwrap(),
        DnsRecord::new_a("name4", "1.1.1.1").unwrap(),
        DnsRecord::new_aaaa("name4", "2:2:2:2:2:2:2:2").unwrap(),
        DnsRecord::new_cname("cname5.example.com", "target1.example.com").unwrap(),
        DnsRecord::new_ns("example.com", "ns1.example.com").unwrap(),
        DnsRecord::new_ns("example.com", "ns2.example.com").unwrap(),
        DnsRecord::new_txt("txt.example.com", "abc").unwrap(),
        DnsRecord::new_txt_multi("txt2.example.com", &["a", "b"]).unwrap(),
    ];
    let join_handle = std::thread::spawn(move || {
        builder.with_permit(permit).serve_static(&records).unwrap();
    });
    for (kind, name, expected_results) in [
        ("a", "notfound1", vec![""]),
        // Wrong type.
        ("cname", "name1", vec![""]),
        ("a", "cname5.example.com", vec![""]),
        ("aaaa", "cname5.example.com", vec![""]),
        ("ns", "cname5.example.com", vec![""]),
        // Exact match.
        ("a", "name1", vec!["1.2.3.4\n"]),
        (
            "aaaa",
            "name2",
            vec!["2606:2800:220:1:248:1893:25c8:1946\n"],
        ),
        (
            "a",
            "name3",
            vec!["10.0.0.1\n10.0.0.2\n", "10.0.0.2\n10.0.0.1\n"],
        ),
        (
            "cname",
            "cname5.example.com",
            vec!["target1.example.com.\n"],
        ),
        (
            "ns",
            "example.com",
            vec![
                "ns1.example.com.\nns2.example.com.\n",
                "ns2.example.com.\nns1.example.com.\n",
            ],
        ),
        ("txt", "txt.example.com", vec!["\"abc\"\n"]),
        ("txt", "txt2.example.com", vec!["\"a\" \"b\"\n"]),
        // Any.
        ("any", "name1", vec!["1.2.3.4\n"]),
        ("any", "name2", vec!["2606:2800:220:1:248:1893:25c8:1946\n"]),
        (
            "any",
            "name3",
            vec!["10.0.0.1\n10.0.0.2\n", "10.0.0.2\n10.0.0.1\n"],
        ),
        (
            "any",
            "name4",
            vec!["1.1.1.1\n2:2:2:2:2:2:2:2\n", "2:2:2:2:2:2:2:2\n1.1.1.1\n"],
        ),
        ("any", "cname5.example.com", vec!["target1.example.com.\n"]),
        (
            "any",
            "example.com",
            vec![
                "ns1.example.com.\nns2.example.com.\n",
                "ns2.example.com.\nns1.example.com.\n",
            ],
        ),
    ] {
        let output = Command::new("dig")
            .arg("@localhost")
            .arg("-p")
            .arg(addr.port().to_string())
            .arg("+notcp")
            .arg("+time=1")
            .arg("+tries=1")
            .arg("+short")
            .arg(kind)
            .arg(name)
            .output()
            .unwrap();
        let stdout = String::from_utf8(output.stdout).unwrap();
        assert!(
            output.status.success(),
            "dig command failed for kind={kind:?} name={name:?} stdout={stdout:?}"
        );
        assert!(
            expected_results.contains(&stdout.as_str()),
            "unexpected response for kind={kind:?} name={name:?} expected_results={expected_results:?}, got: {stdout:?}"
        );
    }
    top_permit.revoke();
    join_handle.join().unwrap();
}

#[test]
fn result_randomization() {
    let top_permit = Permit::new();
    let permit = top_permit.new_sub();
    let (builder, addr) = Builder::new_random_port().unwrap();
    let records = vec![
        DnsRecord::new_a("name1", "10.0.0.1").unwrap(),
        DnsRecord::new_a("name1", "10.0.0.2").unwrap(),
    ];
    let join_handle = std::thread::spawn(move || {
        builder.with_permit(permit).serve_static(&records).unwrap();
    });
    let mut order1 = 0usize;
    let mut order2 = 0usize;
    for _ in 0..20 {
        let output = Command::new("dig")
            .arg("@127.0.0.1")
            .arg("-p")
            .arg(addr.port().to_string())
            .arg("+time=1")
            .arg("+tries=1")
            .arg("+short")
            .arg("a")
            .arg("name1")
            .output()
            .unwrap();
        let stdout = String::from_utf8(output.stdout).unwrap();
        assert!(output.status.success());
        if stdout.as_str() == "10.0.0.1\n10.0.0.2\n" {
            order1 += 1;
        } else if stdout.as_str() == "10.0.0.2\n10.0.0.1\n" {
            order2 += 1;
        } else {
            panic!("unexpected dig output: {stdout:?}");
        }
    }
    assert_ne!(order1, 0);
    assert_ne!(order2, 0);
    top_permit.revoke();
    join_handle.join().unwrap();
}

#[test]
#[allow(clippy::unusual_byte_groupings)]
#[allow(clippy::too_many_lines)]
fn hard_coded() {
    let top_permit = Permit::new();
    let permit = top_permit.new_sub();
    let (builder, addr) = Builder::new_random_port().unwrap();
    let records = vec![
        DnsRecord::new_a("aaa.example.com", "10.0.0.1").unwrap(),
        DnsRecord::new_aaaa("aaa.example.com", "2606:2800:220:1:248:1893:25c8:1946").unwrap(),
    ];
    let join_handle = std::thread::spawn(move || {
        builder.with_permit(permit).serve_static(&records).unwrap();
    });
    let client_sock =
        UdpSocket::bind(SocketAddr::new(IpAddr::V6(Ipv6Addr::UNSPECIFIED), 0)).unwrap();
    client_sock
        .set_write_timeout(Some(Duration::from_secs(1)))
        .unwrap();
    client_sock
        .set_read_timeout(Some(Duration::from_secs(1)))
        .unwrap();
    client_sock.connect(addr).unwrap();
    let mut buf = [0_u8; 512];
    // request type=A
    client_sock
        .send(&[
            // ID
            0x9A,
            0x9A,
            // is_response=0b0, opcode=0b0000 QUERY, authoritative=0b0, truncated=0b0, recursion_desired=0b0
            0b0_0000_0_0_1,
            // recursion_available=0b0, reserved=0b000, response_code=0b0000 NOERROR
            0b0_000_0000,
            // question_count=1
            0x00,
            0x01,
            // answer_count=0
            0x00,
            0x00,
            // name_server_count=0
            0x00,
            0x00,
            // additional_count=0
            0x00,
            0x00,
            // question 0
            // name=aaa.example.com
            0x03,
            97,
            97,
            97,
            0x07,
            101,
            120,
            97,
            109,
            112,
            108,
            101,
            0x03,
            99,
            111,
            109,
            0x00,
            // type=1 A
            0x00,
            0x01,
            // class=1 IN
            0x00,
            0x01,
        ])
        .unwrap();
    let response_len = client_sock.recv(&mut buf).unwrap();
    let response = &buf[0..response_len];
    assert_eq!(
        vec![
            // ID
            0x9A,
            0x9A,
            // is_response=0b1, opcode=0b0000 QUERY, authoritative=0b1, truncated=0b0, recursion_desired=0b1
            0b1_0000_1_0_1,
            // recursion_available=0b0, reserved=0b000, response_code=0b0000 NOERROR
            0b0_000_0000,
            // question_count=1
            0x00,
            0x01,
            // answer_count=1
            0x00,
            0x01,
            // name_server_count=0
            0x00,
            0x00,
            // additional_count=0
            0x00,
            0x00,
            // question 0
            // name=aaa.example.com
            0x03,
            97,
            97,
            97,
            0x07,
            101,
            120,
            97,
            109,
            112,
            108,
            101,
            0x03,
            99,
            111,
            109,
            0x00,
            // type=1 A
            0x00,
            0x01,
            // class=1 IN
            0x00,
            0x01,
            // answer 0
            // name=aaa.example.com
            0x03,
            97,
            97,
            97,
            0x07,
            101,
            120,
            97,
            109,
            112,
            108,
            101,
            0x03,
            99,
            111,
            109,
            0x00,
            // type=1 A
            0x00,
            0x01,
            // class=1 IN
            0x00,
            0x01,
            // ttl_seconds=300
            0x00,
            0x00,
            0x01,
            0x2C,
            // rdlength=4
            0x00,
            0x04,
            // ipv4_addr=10.0.0.1
            10,
            0,
            0,
            1,
        ],
        response
    );
    // request type=AAAA
    client_sock
        .send(&[
            // ID
            0x9A,
            0x9A,
            // is_response=0b0, opcode=0b0000 QUERY, authoritative=0b0, truncated=0b0, recursion_desired=0b0
            0b0_0000_0_0_1,
            // recursion_available=0b0, reserved=0b000, response_code=0b0000 NOERROR
            0b0_000_0000,
            // question_count=1
            0x00,
            0x01,
            // answer_count=0
            0x00,
            0x00,
            // name_server_count=0
            0x00,
            0x00,
            // additional_count=0
            0x00,
            0x00,
            // question 0
            // name=aaa.example.com
            0x03,
            97,
            97,
            97,
            0x07,
            101,
            120,
            97,
            109,
            112,
            108,
            101,
            0x03,
            99,
            111,
            109,
            0x00,
            // type=28 AAAA
            0x00,
            0x1C,
            // class=1 IN
            0x00,
            0x01,
        ])
        .unwrap();
    let response_len = client_sock.recv(&mut buf).unwrap();
    let response = &buf[0..response_len];
    assert_eq!(
        vec![
            // ID
            0x9A,
            0x9A,
            // is_response=0b1, opcode=0b0000 QUERY, authoritative=0b1, truncated=0b0, recursion_desired=0b1
            0b1_0000_1_0_1,
            // recursion_available=0b0, reserved=0b000, response_code=0b0000 NOERROR
            0b0_000_0000,
            // question_count=1
            0x00,
            0x01,
            // answer_count=1
            0x00,
            0x01,
            // name_server_count=0
            0x00,
            0x00,
            // additional_count=0
            0x00,
            0x00,
            // question 0
            // name=aaa.example.com
            0x03,
            97,
            97,
            97,
            0x07,
            101,
            120,
            97,
            109,
            112,
            108,
            101,
            0x03,
            99,
            111,
            109,
            0x00,
            // type=28 AAAA
            0x00,
            0x1C,
            // class=1 IN
            0x00,
            0x01,
            // answer 0
            // name=aaa.example.com
            0x03,
            97,
            97,
            97,
            0x07,
            101,
            120,
            97,
            109,
            112,
            108,
            101,
            0x03,
            99,
            111,
            109,
            0x00,
            // type=28 AAAA
            0x00,
            0x1C,
            // class=1 IN
            0x00,
            0x01,
            // ttl_seconds=300
            0x00,
            0x00,
            0x01,
            0x2C,
            // rdlength=16
            0x00,
            0x10,
            // ipv6_addr=2606:2800:220:1:248:1893:25c8:1946
            0x26,
            0x06,
            0x28,
            0x00,
            0x02,
            0x20,
            0x00,
            0x01,
            0x02,
            0x48,
            0x18,
            0x93,
            0x25,
            0xc8,
            0x19,
            0x46
        ],
        response
    );
    // request type=ANY
    client_sock
        .send(&[
            // ID
            0x9A,
            0x9A,
            // is_response=0b0, opcode=0b0000 QUERY, authoritative=0b0, truncated=0b0, recursion_desired=0b0
            0b0_0000_0_0_1,
            // recursion_available=0b0, reserved=0b000, response_code=0b0000 NOERROR
            0b0_000_0000,
            // question_count=1
            0x00,
            0x01,
            // answer_count=0
            0x00,
            0x00,
            // name_server_count=0
            0x00,
            0x00,
            // additional_count=0
            0x00,
            0x00,
            // question 0
            // name=aaa.example.com
            0x03,
            97,
            97,
            97,
            0x07,
            101,
            120,
            97,
            109,
            112,
            108,
            101,
            0x03,
            99,
            111,
            109,
            0x00,
            // type=255 ANY
            0x00,
            0xFF,
            // class=1 IN
            0x00,
            0x01,
        ])
        .unwrap();
    let response_len = client_sock.recv(&mut buf).unwrap();
    let response = &buf[0..response_len];
    assert_eq!(
        vec![
            // ID
            0x9A,
            0x9A,
            // is_response=0b1, opcode=0b0000 QUERY, authoritative=0b1, truncated=0b0, recursion_desired=0b1
            0b1_0000_1_0_1,
            // recursion_available=0b0, reserved=0b000, response_code=0b0000 NOERROR
            0b0_000_0000,
            // question_count=1
            0x00,
            0x01,
            // answer_count=2
            0x00,
            0x02,
            // name_server_count=0
            0x00,
            0x00,
            // additional_count=0
            0x00,
            0x00,
            // question 0
            // name=aaa.example.com
            0x03,
            97,
            97,
            97,
            0x07,
            101,
            120,
            97,
            109,
            112,
            108,
            101,
            0x03,
            99,
            111,
            109,
            0x00,
            // type=255 ANY
            0x00,
            0xFF,
            // class=1 IN
            0x00,
            0x01,
            // answer 0
            // name=aaa.example.com
            0x03,
            97,
            97,
            97,
            0x07,
            101,
            120,
            97,
            109,
            112,
            108,
            101,
            0x03,
            99,
            111,
            109,
            0x00,
            // type=1 A
            0x00,
            0x01,
            // class=1 IN
            0x00,
            0x01,
            // ttl_seconds=300
            0x00,
            0x00,
            0x01,
            0x2C,
            // rdlength=4
            0x00,
            0x04,
            // ipv4_addr=10.0.0.1
            10,
            0,
            0,
            1,
            // answer 1
            // name=aaa.example.com
            0x03,
            97,
            97,
            97,
            0x07,
            101,
            120,
            97,
            109,
            112,
            108,
            101,
            0x03,
            99,
            111,
            109,
            0x00,
            // type=28 AAAA
            0x00,
            0x1C,
            // class=1 IN
            0x00,
            0x01,
            // ttl_seconds=300
            0x00,
            0x00,
            0x01,
            0x2C,
            // rdlength=16
            0x00,
            0x10,
            // ipv6_addr=2606:2800:220:1:248:1893:25c8:1946
            0x26,
            0x06,
            0x28,
            0x00,
            0x02,
            0x20,
            0x00,
            0x01,
            0x02,
            0x48,
            0x18,
            0x93,
            0x25,
            0xc8,
            0x19,
            0x46
        ],
        response
    );
    drop(top_permit);
    join_handle.join().unwrap();
}

// https://github.com/m-ou-se/single-use-dns
// https://crates.io/crates/dns-parser/0.8.0
// https://docs.rs/rusty_dns/0.0.3/rusty_dns/index.html