networker_rs/
lib.rs

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
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
use std::io::{self, BufRead, Read, Write, BufReader};
use std::net::{TcpListener, TcpStream, UdpSocket, ToSocketAddrs};
use std::collections::VecDeque;

pub mod net {
    use super::*;

    /// Dial a TCP connection to the specified address.
    pub fn dial<A: ToSocketAddrs>(address: A) -> io::Result<TcpStream> {
        TcpStream::connect(address)
    }

    /// Start a TCP listener at the specified address.
    pub fn listen<A: ToSocketAddrs>(address: A) -> io::Result<TcpListener> {
        TcpListener::bind(address)
    }

    /// Send a UDP message to the specified address.
    pub fn udp_send<A: ToSocketAddrs>(address: A, message: &[u8]) -> io::Result<()> {
        let socket = UdpSocket::bind("0.0.0.0:0")?;
        socket.send_to(message, address)?;
        Ok(())
    }

    /// Receive a UDP message on the specified address.
    pub fn udp_receive<A: ToSocketAddrs>(address: A, buffer: &mut [u8]) -> io::Result<(usize, String)> {
        let socket = UdpSocket::bind(address)?;
        let (size, src) = socket.recv_from(buffer)?;
        Ok((size, src.to_string()))
    }

    /// Utility function to write data to a TCP stream.
    pub fn write_to_stream(stream: &mut TcpStream, data: &[u8]) -> io::Result<()> {
        stream.write_all(data)
    }

    /// Utility function to read data from a TCP stream.
    pub fn read_from_stream(stream: &mut TcpStream, buffer: &mut [u8]) -> io::Result<usize> {
        stream.read(buffer)
    }

    /// Read the latest message from the TCP stream until the last newline character.
    pub fn read_latest_message(stream: &mut TcpStream) -> io::Result<String> {
        let mut reader = BufReader::new(stream);
        let mut buffer = VecDeque::new();
        let mut line = String::new();

        while reader.read_line(&mut line)? > 0 {
            buffer.push_back(line.clone());
            line.clear();
        }

        Ok(buffer.pop_back().unwrap_or_default().trim().to_string())
    }

    pub mod http {
        use super::*;

        /// Make a simple HTTP GET request to the specified address.
        pub fn get<A: ToSocketAddrs>(address: A, path: &str) -> io::Result<String> {
            let mut stream = TcpStream::connect(address)?;
            let request = format!("GET {} HTTP/1.1\r\nHost: localhost\r\nConnection: close\r\n\r\n", path);
            stream.write_all(request.as_bytes())?;

            let mut response = String::new();
            stream.read_to_string(&mut response)?;
            Ok(response)
        }

        /// Make a simple HTTP POST request to the specified address with a body.
        pub fn post<A: ToSocketAddrs>(address: A, path: &str, body: &str) -> io::Result<String> {
            let mut stream = TcpStream::connect(address)?;
            let request = format!(
                "POST {} HTTP/1.1\r\nHost: localhost\r\nContent-Length: {}\r\nContent-Type: text/plain\r\nConnection: close\r\n\r\n{}",
                path,
                body.len(),
                body
            );
            stream.write_all(request.as_bytes())?;

            let mut response = String::new();
            stream.read_to_string(&mut response)?;
            Ok(response)
        }
    }
}

#[cfg(test)]
mod tests {
    use super::net;
    use std::thread;

    #[test]
    fn test_tcp_connection() {
        let address = "127.0.0.1:7878";

        // Start a server in a separate thread
        thread::spawn(move || {
            let listener = net::listen(address).unwrap();
            for stream in listener.incoming() {
                let mut stream = stream.unwrap();
                let mut buffer = [0; 512];
                net::read_from_stream(&mut stream, &mut buffer).unwrap();
                net::write_to_stream(&mut stream, b"Hello, client!").unwrap();
            }
        });

        // Simulate a client
        let mut client = net::dial(address).unwrap();
        net::write_to_stream(&mut client, b"Hello, server!").unwrap();

        let mut buffer = [0; 512];
        let size = net::read_from_stream(&mut client, &mut buffer).unwrap();
        assert_eq!(String::from_utf8_lossy(&buffer[..size]), "Hello, client!");
    }

    #[test]
    fn test_udp_message() {
        let server_address = "127.0.0.1:8888";
        let client_address = "127.0.0.1:0";

        let handle = thread::spawn(move || {
            let mut buffer = [0; 512];
            let (size, src) = net::udp_receive(server_address, &mut buffer).unwrap();
            assert_eq!(src, "127.0.0.1:0");
            assert_eq!(String::from_utf8_lossy(&buffer[..size]), "Hello, UDP server!");
        });

        net::udp_send(server_address, b"Hello, UDP server!").unwrap();
        handle.join().unwrap();
    }

    #[test]
    fn test_read_latest_message() {
        let address = "127.0.0.1:8989";

        thread::spawn(move || {
            let listener = net::listen(address).unwrap();
            for stream in listener.incoming() {
                let mut stream = stream.unwrap();
                net::write_to_stream(&mut stream, b"yes\nno\nye\n").unwrap();
            }
        });

        let mut client = net::dial(address).unwrap();
        let latest_message = net::read_latest_message(&mut client).unwrap();
        assert_eq!(latest_message, "ye");
    }

    #[test]
    fn test_http_get() {
        // Simulate a basic HTTP server
        let address = "127.0.0.1:8080";
        thread::spawn(move || {
            let listener = net::listen(address).unwrap();
            for stream in listener.incoming() {
                let mut stream = stream.unwrap();
                let mut buffer = [0; 512];
                net::read_from_stream(&mut stream, &mut buffer).unwrap();
                net::write_to_stream(&mut stream, b"HTTP/1.1 200 OK\r\nContent-Length: 13\r\n\r\nHello, world!").unwrap();
            }
        });

        let response = net::http::get(address, "/").unwrap();
        assert!(response.contains("Hello, world!"));
    }

    #[test]
    fn test_http_post() {
        // Simulate a basic HTTP server
        let address = "127.0.0.1:8081";
        thread::spawn(move || {
            let listener = net::listen(address).unwrap();
            for stream in listener.incoming() {
                let mut stream = stream.unwrap();
                let mut buffer = [0; 512];
                net::read_from_stream(&mut stream, &mut buffer).unwrap();
                net::write_to_stream(&mut stream, b"HTTP/1.1 200 OK\r\nContent-Length: 7\r\n\r\nSuccess").unwrap();
            }
        });

        let response = net::http::post(address, "/submit", "data").unwrap();
        assert!(response.contains("Success"));
    }
}