1use std::io::{Read, Write};
6use std::net::{TcpStream, ToSocketAddrs};
7use std::path::PathBuf;
8use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
9use std::sync::{Arc, Mutex};
10use std::time::Duration;
11
12use anyhow::{Context, Result, bail, ensure};
13
14static CONNECT_SECS: AtomicU64 = AtomicU64::new(10);
19static IO_SECS: AtomicU64 = AtomicU64::new(30);
20
21const MIN_IO_SECS: u64 = 5;
25
26pub fn set_timeouts(connect_secs: u64, io_secs: u64) {
29 CONNECT_SECS.store(connect_secs, Ordering::Relaxed);
30 IO_SECS.store(io_secs.max(MIN_IO_SECS), Ordering::Relaxed);
31}
32
33pub fn io_timeout() -> Duration {
35 Duration::from_secs(IO_SECS.load(Ordering::Relaxed).max(MIN_IO_SECS))
36}
37
38fn connect_timeout() -> Option<Duration> {
39 match CONNECT_SECS.load(Ordering::Relaxed) {
40 0 => None,
41 secs => Some(Duration::from_secs(secs)),
42 }
43}
44
45fn timed_out(err: &std::io::Error) -> bool {
47 matches!(
48 err.kind(),
49 std::io::ErrorKind::WouldBlock | std::io::ErrorKind::TimedOut
50 )
51}
52
53pub(crate) enum Stream {
54 Plain(TcpStream),
55 Tls(Box<rustls::StreamOwned<rustls::ClientConnection, TcpStream>>),
56}
57
58#[derive(Clone, Default)]
62pub struct Cutoff(Arc<CutoffState>);
63
64#[derive(Default)]
65struct CutoffState {
66 socket: Mutex<Option<TcpStream>>,
67 on_purpose: AtomicBool,
70}
71
72impl Cutoff {
73 fn hold(&self, tcp: &TcpStream) {
74 self.0.on_purpose.store(false, Ordering::Relaxed);
75 if let (Ok(mut slot), Ok(clone)) = (self.0.socket.lock(), tcp.try_clone()) {
76 *slot = Some(clone);
77 }
78 }
79
80 pub fn cut(&self) {
83 self.0.on_purpose.store(true, Ordering::Relaxed);
84 if let Ok(slot) = self.0.socket.lock()
85 && let Some(tcp) = slot.as_ref()
86 {
87 let _ = tcp.shutdown(std::net::Shutdown::Both);
88 }
89 }
90
91 pub fn was_cut(&self) -> bool {
94 self.0.on_purpose.swap(false, Ordering::Relaxed)
95 }
96}
97
98impl Stream {
99 pub(crate) fn into_tcp(self) -> Result<TcpStream> {
101 match self {
102 Stream::Plain(tcp) => Ok(tcp),
103 Stream::Tls(_) => bail!("connection is already TLS"),
104 }
105 }
106}
107
108impl Read for Stream {
109 fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
110 match self {
111 Stream::Plain(s) => s.read(buf),
112 Stream::Tls(s) => s.read(buf),
113 }
114 }
115}
116
117impl Write for Stream {
118 fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
119 match self {
120 Stream::Plain(s) => s.write(buf),
121 Stream::Tls(s) => s.write(buf),
122 }
123 }
124
125 fn flush(&mut self) -> std::io::Result<()> {
126 match self {
127 Stream::Plain(s) => s.flush(),
128 Stream::Tls(s) => s.flush(),
129 }
130 }
131}
132
133pub(crate) fn is_connection_error(err: &anyhow::Error) -> bool {
137 err.downcast_ref::<std::io::Error>().is_some()
138 || err
139 .chain()
140 .any(|c| c.to_string().contains("server closed the connection"))
141}
142
143pub(crate) fn is_timeout(err: &anyhow::Error) -> bool {
146 err.downcast_ref::<std::io::Error>().is_some_and(|e| {
147 matches!(
148 e.kind(),
149 std::io::ErrorKind::WouldBlock | std::io::ErrorKind::TimedOut
150 )
151 })
152}
153
154pub(crate) fn connect(host: &str, port: u16, tls: bool, cutoff: &Cutoff) -> Result<Stream> {
155 let tcp = connect_tcp(host, port)?;
156 tcp.set_read_timeout(Some(io_timeout()))?;
157 tcp.set_write_timeout(Some(io_timeout()))?;
158 cutoff.hold(&tcp);
159 if tls {
160 wrap_tls(tcp, host)
161 } else {
162 Ok(Stream::Plain(tcp))
163 }
164}
165
166fn connect_tcp(host: &str, port: u16) -> Result<TcpStream> {
172 let Some(timeout) = connect_timeout() else {
173 return TcpStream::connect((host, port))
174 .with_context(|| format!("connecting to {host}:{port}"));
175 };
176 let addrs: Vec<_> = (host, port)
177 .to_socket_addrs()
178 .with_context(|| format!("resolving {host}"))?
179 .collect();
180 ensure!(!addrs.is_empty(), "{host} resolves to nothing");
181 let mut last = None;
182 for addr in &addrs {
183 match TcpStream::connect_timeout(addr, timeout) {
184 Ok(tcp) => return Ok(tcp),
185 Err(err) => last = Some(err),
186 }
187 }
188 let err = last.expect("at least one address was tried");
189 let secs = timeout.as_secs();
190 let said = match timed_out(&err) {
191 true => format!("connecting to {host}:{port} timed out after {secs}s"),
192 false => format!("connecting to {host}:{port}"),
193 };
194 Err(anyhow::Error::from(err).context(said))
195}
196
197static TRUST: Mutex<Trust> = Mutex::new(Trust {
203 system: true,
204 extra_pem: None,
205});
206
207struct Trust {
208 system: bool,
209 extra_pem: Option<PathBuf>,
210}
211
212pub fn set_trust(system: bool, certificate_file: Option<PathBuf>) {
216 let mut trust = TRUST.lock().unwrap();
217 trust.system = system;
218 trust.extra_pem = certificate_file;
219}
220
221fn root_store() -> Result<rustls::RootCertStore> {
226 let mut roots = rustls::RootCertStore::empty();
227 roots.extend(webpki_roots::TLS_SERVER_ROOTS.iter().cloned());
228 let trust = TRUST.lock().unwrap();
229 if trust.system {
230 let loaded = rustls_native_certs::load_native_certs();
233 for cert in loaded.certs {
234 let _ = roots.add(cert);
235 }
236 }
237 if let Some(path) = &trust.extra_pem {
238 let pem = std::fs::read(path)
239 .with_context(|| format!("reading certificate_file {}", path.display()))?;
240 let (added, _) = roots.add_parsable_certificates(parse_pem_certs(&pem)?);
241 if added == 0 {
242 anyhow::bail!("no certificates in {}", path.display());
243 }
244 }
245 Ok(roots)
246}
247
248fn parse_pem_certs(pem: &[u8]) -> Result<Vec<rustls::pki_types::CertificateDer<'static>>> {
250 let mut cursor = std::io::Cursor::new(pem);
251 rustls_pemfile::certs(&mut cursor)
252 .collect::<std::result::Result<Vec<_>, _>>()
253 .context("parsing certificate_file PEM")
254}
255
256pub(crate) fn wrap_tls(tcp: TcpStream, host: &str) -> Result<Stream> {
257 let roots = root_store()?;
258 let config = rustls::ClientConfig::builder()
259 .with_root_certificates(roots)
260 .with_no_client_auth();
261 let name = rustls::pki_types::ServerName::try_from(host.to_string())
262 .with_context(|| format!("invalid server name {host}"))?;
263 let conn = rustls::ClientConnection::new(Arc::new(config), name)
264 .with_context(|| format!("setting up TLS to {host}"))?;
265 Ok(Stream::Tls(Box::new(rustls::StreamOwned::new(conn, tcp))))
266}
267
268pub(crate) struct Conn {
270 stream: Stream,
271 peer: String,
273 buf: Vec<u8>,
274 start: usize,
275 end: usize,
276}
277
278impl Conn {
279 pub(crate) fn new(stream: Stream, peer: impl Into<String>) -> Conn {
280 Conn {
281 stream,
282 peer: peer.into(),
283 buf: vec![0; 8192],
284 start: 0,
285 end: 0,
286 }
287 }
288
289 fn io_error(&self, err: std::io::Error, doing: &str) -> anyhow::Error {
292 let peer = &self.peer;
293 let said = match timed_out(&err) {
296 true => format!(
297 "{peer} timed out after {}s while {doing}",
298 io_timeout().as_secs()
299 ),
300 false => format!("{doing} {peer}"),
301 };
302 anyhow::Error::from(err).context(said)
303 }
304
305 pub(crate) fn into_stream(self) -> Stream {
307 self.stream
308 }
309
310 fn fill(&mut self) -> Result<()> {
311 self.start = 0;
312 self.end = loop {
313 match self.stream.read(&mut self.buf) {
314 Err(e) if e.kind() == std::io::ErrorKind::Interrupted => continue,
317 Ok(n) => break n,
318 Err(err) => return Err(self.io_error(err, "reading from")),
319 }
320 };
321 if self.end == 0 {
322 bail!("server closed the connection");
323 }
324 Ok(())
325 }
326
327 pub(crate) fn read_byte(&mut self) -> Result<u8> {
328 if self.start == self.end {
329 self.fill()?;
330 }
331 let b = self.buf[self.start];
332 self.start += 1;
333 Ok(b)
334 }
335
336 pub(crate) fn read_exact_to(&mut self, out: &mut Vec<u8>, n: usize) -> Result<()> {
338 let mut left = n;
339 while left > 0 {
340 if self.start == self.end {
341 self.fill()?;
342 }
343 let take = left.min(self.end - self.start);
344 out.extend_from_slice(&self.buf[self.start..self.start + take]);
345 self.start += take;
346 left -= take;
347 }
348 Ok(())
349 }
350
351 pub(crate) fn read_text_line(&mut self) -> Result<String> {
353 let mut bytes = Vec::new();
354 loop {
355 let b = self.read_byte()?;
356 if b == b'\n' {
357 if bytes.last() == Some(&b'\r') {
358 bytes.pop();
359 }
360 break;
361 }
362 bytes.push(b);
363 ensure!(bytes.len() <= 1 << 20, "response line too long");
364 }
365 Ok(String::from_utf8_lossy(&bytes).into_owned())
366 }
367
368 pub(crate) fn write_all(&mut self, bytes: &[u8]) -> Result<()> {
369 self.stream
370 .write_all(bytes)
371 .map_err(|err| self.io_error(err, "writing to"))?;
372 self.stream
373 .flush()
374 .map_err(|err| self.io_error(err, "writing to"))?;
375 Ok(())
376 }
377}
378
379#[cfg(test)]
380mod tests {
381 use super::*;
382
383 #[test]
384 fn timeouts_are_never_off_and_never_too_short() {
385 set_timeouts(10, 1);
388 assert_eq!(io_timeout().as_secs(), MIN_IO_SECS);
389 assert_eq!(connect_timeout(), Some(Duration::from_secs(10)));
390 set_timeouts(0, 45);
392 assert_eq!(connect_timeout(), None);
393 assert_eq!(io_timeout().as_secs(), 45);
394 set_timeouts(10, 30);
395 }
396
397 #[test]
398 fn the_trust_store_only_ever_adds() {
399 set_trust(false, None);
405 let store = root_store().unwrap();
406 assert!(store.len() > 50, "webpki roots present: {}", store.len());
407
408 let dir = std::env::temp_dir().join(format!("rmut-net-{}", std::process::id()));
411 std::fs::create_dir_all(&dir).unwrap();
412 let path = dir.join("garbage.pem");
413 std::fs::write(&path, b"not a certificate\n").unwrap();
414 set_trust(false, Some(path.clone()));
415 let err = root_store().unwrap_err().to_string();
416 assert!(err.contains("no certificates"), "{err}");
417
418 set_trust(true, None);
420 std::fs::remove_dir_all(&dir).ok();
421 }
422}