1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148

use std::io::{Read, Write};
use std::net::{IpAddr, SocketAddr};
use std::sync::{atomic::AtomicBool, Arc};

use futures_util::stream::StreamExt;

use crate::tcp::{Connect, TcpClientStream, TcpStream};
use crate::xfer::dns_handle::DnsStreamHandle;
use crate::xfer::SerialMessage;
use crate::Executor;

const TEST_BYTES: &[u8; 8] = b"DEADBEEF";
const TEST_BYTES_LEN: usize = 8;
const SEND_RECV_TIMES: usize = 4;

fn tcp_server_setup(
    server_name: &str,
    server_addr: IpAddr,
) -> (Arc<AtomicBool>, std::thread::JoinHandle<()>, SocketAddr) {
    let succeeded = Arc::new(AtomicBool::new(false));
    let succeeded_clone = succeeded.clone();
    std::thread::Builder::new()
        .name("thread_killer".to_string())
        .spawn(move || {
            let succeeded = succeeded_clone;
            for _ in 0..15 {
                std::thread::sleep(std::time::Duration::from_secs(1));
                if succeeded.load(std::sync::atomic::Ordering::Relaxed) {
                    return;
                }
            }

            println!("Thread Killer has been awoken, killing process");
            std::process::exit(-1);
        })
        .expect("Thread spawning failed");

    // TODO: need a timeout on listen
    let server = std::net::TcpListener::bind(SocketAddr::new(server_addr, 0))
        .expect("Unable to bind a TCP socket");
    let server_addr = server.local_addr().unwrap();

    // an in and out server
    let server_handle = std::thread::Builder::new()
        .name(server_name.to_string())
        .spawn(move || {
            let (mut socket, _) = server.accept().expect("accept failed");

            socket
                .set_read_timeout(Some(std::time::Duration::from_secs(5)))
                .unwrap(); // should receive something within 5 seconds...
            socket
                .set_write_timeout(Some(std::time::Duration::from_secs(5)))
                .unwrap(); // should receive something within 5 seconds...

            for _ in 0..SEND_RECV_TIMES {
                // wait for some bytes...
                let mut len_bytes = [0_u8; 2];
                socket
                    .read_exact(&mut len_bytes)
                    .expect("SERVER: receive failed");
                let length =
                    u16::from(len_bytes[0]) << 8 & 0xFF00 | u16::from(len_bytes[1]) & 0x00FF;
                assert_eq!(length as usize, TEST_BYTES_LEN);

                let mut buffer = [0_u8; TEST_BYTES_LEN];
                socket.read_exact(&mut buffer).unwrap();

                // println!("read bytes iter: {}", i);
                assert_eq!(&buffer, TEST_BYTES);

                // bounce them right back...
                socket
                    .write_all(&len_bytes)
                    .expect("SERVER: send length failed");
                socket
                    .write_all(&buffer)
                    .expect("SERVER: send buffer failed");
                // println!("wrote bytes iter: {}", i);
                std::thread::yield_now();
            }
        })
        .unwrap();
    (succeeded, server_handle, server_addr)
}

/// Test tcp_stream.
pub fn tcp_stream_test<S: Connect, E: Executor>(server_addr: IpAddr, mut exec: E) {
    let (succeeded, server_handle, server_addr) =
        tcp_server_setup("test_tcp_stream:server", server_addr);

    // setup the client, which is going to run on the testing thread...

    // the tests should run within 5 seconds... right?
    // TODO: add timeout here, so that test never hangs...
    // let timeout = Timeout::new(Duration::from_secs(5));
    let (stream, mut sender) = TcpStream::<S>::new(server_addr);

    let mut stream = exec.block_on(stream).expect("run failed to get stream");

    for _ in 0..SEND_RECV_TIMES {
        // test once
        sender
            .send(SerialMessage::new(TEST_BYTES.to_vec(), server_addr))
            .expect("send failed");

        let (buffer, stream_tmp) = exec.block_on(stream.into_future());
        stream = stream_tmp;
        let message = buffer
            .expect("no buffer received")
            .expect("error receiving buffer");
        assert_eq!(message.bytes(), TEST_BYTES);
    }

    succeeded.store(true, std::sync::atomic::Ordering::Relaxed);
    server_handle.join().expect("server thread failed");
}

/// Test tcp_client_stream.
pub fn tcp_client_stream_test<S: Connect, E: Executor>(server_addr: IpAddr, mut exec: E) {
    let (succeeded, server_handle, server_addr) =
        tcp_server_setup("test_tcp_client_stream:server", server_addr);

    // setup the client, which is going to run on the testing thread...

    // the tests should run within 5 seconds... right?
    // TODO: add timeout here, so that test never hangs...
    // let timeout = Timeout::new(Duration::from_secs(5));
    let (stream, mut sender) = TcpClientStream::<S>::new(server_addr);

    let mut stream = exec.block_on(stream).expect("run failed to get stream");

    for _ in 0..SEND_RECV_TIMES {
        // test once
        sender
            .send(SerialMessage::new(TEST_BYTES.to_vec(), server_addr))
            .expect("send failed");
        let (buffer, stream_tmp) = exec.block_on(stream.into_future());
        stream = stream_tmp;
        let buffer = buffer
            .expect("no buffer received")
            .expect("error receiving buffer");
        assert_eq!(buffer.bytes(), TEST_BYTES);
    }

    succeeded.store(true, std::sync::atomic::Ordering::Relaxed);
    server_handle.join().expect("server thread failed");
}