openmoat_proxy/
local.rs

1//! A service of the proxy's own process on the proxy's port (ADR-028).
2//!
3//! `moat run`'s sandbox lets the agent reach one TCP port: the proxy's. The
4//! hooks the agent starts run in that sandbox too, so their decisions are made
5//! by the process outside it, over the same port. A connection whose first
6//! bytes are a [`LocalService::prefix`] is handed to that service instead of
7//! being read as an HTTP request; no HTTP request line starts with it.
8
9use std::net::TcpStream;
10use std::thread;
11use std::time::{Duration, Instant};
12
13/// A service that shares the proxy's port.
14pub trait LocalService: Sync {
15    /// The bytes a connection for this service starts with. Not empty, and
16    /// not the start of any HTTP method.
17    fn prefix(&self) -> &[u8];
18
19    /// Serve one connection. Its first bytes are [`Self::prefix`], still
20    /// unread. The service records what it decides itself.
21    fn serve(&self, client: &TcpStream);
22}
23
24/// How long to wait between looks at a connection that sent part of a prefix.
25const PEEK_PAUSE: Duration = Duration::from_millis(2);
26
27/// Whether `client` starts with `prefix`, without consuming anything. `false`
28/// as soon as its bytes differ, at end of data, on an error or at `deadline`.
29pub(crate) fn starts_with(client: &TcpStream, prefix: &[u8], deadline: Instant) -> bool {
30    let mut seen = vec![0u8; prefix.len()];
31    loop {
32        let n = match client.peek(&mut seen) {
33            Ok(0) | Err(_) => return false,
34            Ok(n) => n,
35        };
36        if seen[..n] != prefix[..n] {
37            return false;
38        }
39        if n == prefix.len() {
40            return true;
41        }
42        if Instant::now() >= deadline {
43            return false;
44        }
45        // `peek` returns what is there at once; wait for the rest.
46        thread::sleep(PEEK_PAUSE);
47    }
48}
49
50#[cfg(test)]
51mod tests {
52    use std::io::{Read as _, Write as _};
53    use std::net::{Shutdown, TcpListener};
54
55    use super::*;
56
57    fn pair() -> (TcpStream, TcpStream) {
58        let listener = TcpListener::bind("127.0.0.1:0").unwrap();
59        let client = TcpStream::connect(listener.local_addr().unwrap()).unwrap();
60        let (server, _) = listener.accept().unwrap();
61        (server, client)
62    }
63
64    fn soon() -> Instant {
65        Instant::now() + Duration::from_millis(500)
66    }
67
68    #[test]
69    fn a_prefix_is_matched_without_consuming_it() {
70        let (mut server, mut client) = pair();
71        client.write_all(b"MOAT-X/1\nrest").unwrap();
72        assert!(starts_with(&server, b"MOAT-X/1\n", soon()));
73        let mut all = Vec::new();
74        client.shutdown(Shutdown::Write).unwrap();
75        server.read_to_end(&mut all).unwrap();
76        assert_eq!(all, b"MOAT-X/1\nrest");
77    }
78
79    #[test]
80    fn a_prefix_sent_in_pieces_is_matched() {
81        let (server, mut client) = pair();
82        let writer = thread::spawn(move || {
83            client.write_all(b"MOAT").unwrap();
84            thread::sleep(Duration::from_millis(30));
85            client.write_all(b"-X/1\n").unwrap();
86            client
87        });
88        assert!(starts_with(&server, b"MOAT-X/1\n", soon()));
89        drop(writer.join());
90    }
91
92    #[test]
93    fn http_a_short_stream_and_silence_are_not_matched() {
94        let (server, mut client) = pair();
95        client.write_all(b"GET http://a/ HTTP/1.1\r\n").unwrap();
96        assert!(!starts_with(&server, b"MOAT-X/1\n", soon()));
97        let (server, mut client) = pair();
98        client.write_all(b"MOA").unwrap();
99        client.shutdown(Shutdown::Write).unwrap();
100        thread::sleep(Duration::from_millis(20));
101        assert!(!starts_with(&server, b"MOAT-X/1\n", soon()));
102        let (server, _client) = pair();
103        server
104            .set_read_timeout(Some(Duration::from_millis(50)))
105            .unwrap();
106        let deadline = Instant::now() + Duration::from_millis(100);
107        assert!(!starts_with(&server, b"MOAT-X/1\n", deadline));
108    }
109}
openmoat_proxy/local.rs

openmoat_proxy/
local.rs
