table_editor/
probe.rs

1//! One-shot HTTP requests to a server on the loopback interface.
2//!
3//! The editor uses these to ask what is on a port and to shut a server down,
4//! and they are public because a repository whose editor brings up a companion
5//! service needs the same question answered about it.
6//!
7//! Each call opens a connection of its own, sends an HTTP/1.0 request with
8//! `Connection: close`, and reads until the far end closes, so no client state
9//! outlives the call. There is no TLS, no redirect following, no chunked
10//! decoding, and no header parsing beyond the status line: this speaks plain
11//! HTTP to a local server, and it is not a general-purpose HTTP client.
12
13use std::io::{Read, Write};
14use std::net::{Ipv4Addr, SocketAddr, TcpStream};
15use std::time::Duration;
16
17/// The most of a response worth reading. A local health or status body is far
18/// smaller; this only stops a talkative server from filling memory.
19const MAX_RESPONSE: usize = 8 * 1024;
20
21/// Ask a server on `127.0.0.1:port` for `path`, returning its status code and
22/// body.
23///
24/// `None` means nothing accepted a connection. A status of `0` means something
25/// answered but not in HTTP this could read, which still says something is
26/// there. `timeout` bounds the connection and each read and write separately,
27/// not the call as a whole.
28pub fn probe(port: u16, method: &str, path: &str, timeout: Duration) -> Option<(u16, String)> {
29    let addr = SocketAddr::from((Ipv4Addr::LOCALHOST, port));
30    let mut stream = TcpStream::connect_timeout(&addr, timeout).ok()?;
31    if stream.set_read_timeout(Some(timeout)).is_err()
32        || stream.set_write_timeout(Some(timeout)).is_err()
33    {
34        return None;
35    }
36
37    let request =
38        format!("{method} {path} HTTP/1.0\r\nHost: localhost\r\nConnection: close\r\n\r\n");
39    if stream.write_all(request.as_bytes()).is_err() {
40        return None;
41    }
42
43    Some(parse_response(&read_response(&mut stream)).unwrap_or((0, String::new())))
44}
45
46/// The status code alone, for a caller that only wants to know whether a
47/// service is up and answering. `None` and `0` mean what they mean in
48/// [`probe`].
49pub fn probe_status(port: u16, method: &str, path: &str, timeout: Duration) -> Option<u16> {
50    probe(port, method, path, timeout).map(|(status, _)| status)
51}
52
53/// Read until the server closes the connection. A timeout or a reset ends the
54/// read and whatever arrived stands, which is what a server exiting on a
55/// shutdown request leaves behind.
56fn read_response(stream: &mut TcpStream) -> Vec<u8> {
57    let mut raw = Vec::new();
58    let mut chunk = [0u8; 1024];
59    loop {
60        match stream.read(&mut chunk) {
61            Ok(0) | Err(_) => break,
62            Ok(n) => {
63                raw.extend_from_slice(&chunk[..n]);
64                if raw.len() >= MAX_RESPONSE {
65                    break;
66                }
67            }
68        }
69    }
70    raw
71}
72
73fn parse_response(raw: &[u8]) -> Option<(u16, String)> {
74    let text = std::str::from_utf8(raw).ok()?;
75    let status = parse_status(text.as_bytes())?;
76    let body = text.split_once("\r\n\r\n").map_or("", |(_, body)| body);
77    Some((status, body.to_string()))
78}
79
80/// Parse the status code from an HTTP status line like `HTTP/1.0 200 OK`.
81fn parse_status(bytes: &[u8]) -> Option<u16> {
82    let text = std::str::from_utf8(bytes).ok()?;
83    let line = text.lines().next()?;
84    line.split_whitespace().nth(1)?.parse().ok()
85}
86
87#[cfg(test)]
88mod tests {
89    use std::net::TcpListener;
90    use std::thread;
91
92    use super::*;
93
94    /// A listener that answers one request with `raw` and then closes.
95    fn answer_once(raw: &'static str) -> u16 {
96        let listener = TcpListener::bind((Ipv4Addr::LOCALHOST, 0)).unwrap();
97        let port = listener.local_addr().unwrap().port();
98        thread::spawn(move || {
99            if let Ok((mut stream, _)) = listener.accept() {
100                let mut request = [0u8; 1024];
101                let _ = stream.read(&mut request);
102                let _ = stream.write_all(raw.as_bytes());
103            }
104        });
105        port
106    }
107
108    #[test]
109    fn parse_status_reads_code() {
110        assert_eq!(parse_status(b"HTTP/1.0 200 OK\r\n"), Some(200));
111        assert_eq!(parse_status(b"HTTP/1.1 404 Not Found\r\n"), Some(404));
112        assert_eq!(parse_status(b"garbage"), None);
113    }
114
115    #[test]
116    fn parse_response_separates_the_status_from_the_body() {
117        let raw = b"HTTP/1.0 200 OK\r\nContent-Length: 15\r\n\r\n{\"status\":\"ok\"}";
118        assert_eq!(
119            parse_response(raw),
120            Some((200, r#"{"status":"ok"}"#.to_string()))
121        );
122    }
123
124    #[test]
125    fn parse_response_tolerates_a_response_cut_short() {
126        assert_eq!(
127            parse_response(b"HTTP/1.0 200 OK\r\n"),
128            Some((200, String::new()))
129        );
130        assert_eq!(parse_response(b""), None);
131    }
132
133    #[test]
134    fn a_probe_returns_the_status_and_body_it_was_answered_with() {
135        let port = answer_once(
136            "HTTP/1.0 200 OK\r\nContent-Type: application/json\r\n\r\n{\"status\":\"ok\"}",
137        );
138        assert_eq!(
139            probe(port, "GET", "/api/health", Duration::from_secs(5)),
140            Some((200, r#"{"status":"ok"}"#.to_string()))
141        );
142    }
143
144    #[test]
145    fn a_probe_of_an_answer_that_is_not_http_reports_status_zero() {
146        let port = answer_once("hello");
147        assert_eq!(
148            probe_status(port, "GET", "/api/health", Duration::from_secs(5)),
149            Some(0)
150        );
151    }
152
153    #[test]
154    fn a_probe_of_a_port_nothing_is_listening_on_is_none() {
155        // Port 1 is privileged and never has our server; the connection is
156        // refused rather than timing out.
157        assert_eq!(
158            probe_status(1, "GET", "/api/health", Duration::from_millis(300)),
159            None
160        );
161    }
162}
table_editor/probe.rs

table_editor/
probe.rs
