This repository has no description
7.1 kB
199 lines
1//! Binary cache access: GET with retries, and body decoding.
2
3use std::collections::HashMap;
4use std::io::{self, Read};
5use std::net::{SocketAddr, ToSocketAddrs};
6use std::sync::Mutex;
7use std::thread::sleep;
8use std::time::Duration;
9
10const TRIES: u32 = 6;
11
12/// Resolves each host once and then serves it from memory. Once the FUSE
13/// file system is mounted on /nix/store, getaddrinfo() could end up looking
14/// for NSS modules in the store, i.e. in a request to ourselves.
15struct Resolver(Mutex<HashMap<String, Vec<SocketAddr>>>);
16
17impl ureq::Resolver for Resolver {
18 fn resolve(&self, netloc: &str) -> io::Result<Vec<SocketAddr>> {
19 if let Some(a) = self.0.lock().unwrap().get(netloc) {
20 return Ok(a.clone());
21 }
22 let a: Vec<SocketAddr> = netloc.to_socket_addrs()?.collect();
23 self.0.lock().unwrap().insert(netloc.to_string(), a.clone());
24 Ok(a)
25 }
26}
27
28pub struct Http {
29 agent: ureq::Agent,
30}
31
32pub enum Error {
33 /// 403/404, or a missing file for file:// caches
34 NotFound,
35 Other(String),
36}
37
38impl std::fmt::Display for Error {
39 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
40 match self {
41 Error::NotFound => write!(f, "not found"),
42 Error::Other(e) => write!(f, "{e}"),
43 }
44 }
45}
46
47impl From<Error> for io::Error {
48 fn from(e: Error) -> io::Error {
49 io::Error::other(e.to_string())
50 }
51}
52
53impl Http {
54 /// `caches` are base URLs; their hosts are resolved right away.
55 pub fn new<'a>(caches: impl IntoIterator<Item = &'a str>) -> Http {
56 let mut known = HashMap::new();
57 for c in caches {
58 let Some(rest) = c
59 .strip_prefix("https://")
60 .or_else(|| c.strip_prefix("http://"))
61 else {
62 continue;
63 };
64 let host = rest.split('/').next().unwrap_or(rest);
65 let port = if c.starts_with("https:") { 443 } else { 80 };
66 let netloc = if host.contains(':') {
67 host.to_string()
68 } else {
69 format!("{host}:{port}")
70 };
71 match netloc.to_socket_addrs() {
72 Ok(a) => {
73 known.insert(netloc, a.collect());
74 }
75 Err(e) => eprintln!("resolving {netloc}: {e}"),
76 }
77 }
78 let agent = ureq::AgentBuilder::new()
79 .timeout_connect(Duration::from_secs(30))
80 .timeout_read(Duration::from_secs(60))
81 .max_idle_connections_per_host(64)
82 .resolver(Resolver(Mutex::new(known)))
83 .build();
84 Http { agent }
85 }
86
87 /// GET `url`, retrying 429/5xx and transport errors with backoff (the
88 /// caches answer 503 under load). Returns the body reader and its
89 /// Content-Encoding.
90 pub fn open(&self, url: &str) -> Result<(Box<dyn Read + Send>, Option<String>), Error> {
91 if let Some(path) = url.strip_prefix("file://") {
92 return match std::fs::File::open(path) {
93 Ok(f) => Ok((Box::new(f), None)),
94 Err(e) if e.kind() == io::ErrorKind::NotFound => Err(Error::NotFound),
95 Err(e) => Err(Error::Other(format!("{url}: {e}"))),
96 };
97 }
98 let mut n = 0;
99 loop {
100 let err = match self.agent.get(url).call() {
101 Ok(r) => {
102 let enc = r.header("content-encoding").map(str::to_string);
103 return Ok((Box::new(r.into_reader()), enc));
104 }
105 Err(ureq::Error::Status(403 | 404, _)) => return Err(Error::NotFound),
106 Err(ureq::Error::Status(code, _)) if code != 429 && code < 500 => {
107 return Err(Error::Other(format!("{url}: HTTP {code}")))
108 }
109 Err(ureq::Error::Status(code, _)) => format!("HTTP {code}"),
110 Err(e) => e.to_string(),
111 };
112 n += 1;
113 if n == TRIES {
114 return Err(Error::Other(format!("{url}: {err}")));
115 }
116 eprintln!("retry {n}/{} {url}: {err}", TRIES - 1);
117 sleep(Duration::from_secs((1u64 << (n - 1)).min(30)));
118 }
119 }
120
121 /// GET `url` -> decoded body. Transfer errors while reading the body are
122 /// retried like failed requests.
123 pub fn get(&self, url: &str) -> Result<(Vec<u8>, usize), Error> {
124 let mut n = 0;
125 loop {
126 let (mut r, enc) = self.open(url)?;
127 let mut body = Vec::new();
128 match r.read_to_end(&mut body) {
129 Ok(_) => {
130 let len = body.len();
131 let dec = decode(&body, enc.as_deref())
132 .map_err(|e| Error::Other(format!("{url}: decoding: {e}")))?;
133 return Ok((dec, len));
134 }
135 Err(e) => {
136 n += 1;
137 if n == TRIES {
138 return Err(Error::Other(format!("{url}: {e}")));
139 }
140 eprintln!("retry {n}/{} {url}: {e}", TRIES - 1);
141 sleep(Duration::from_secs((1u64 << (n - 1)).min(30)));
142 }
143 }
144 }
145 }
146}
147
148/// Wrap `r` in a decompressor for `method` (narinfo Compression or
149/// Content-Encoding).
150pub fn decoder<'a>(
151 method: &str,
152 r: impl Read + Send + 'a,
153) -> io::Result<Box<dyn Read + Send + 'a>> {
154 Ok(match method {
155 "" | "none" | "identity" => Box::new(r),
156 "zstd" => Box::new(zstd::stream::read::Decoder::new(r)?),
157 "xz" => Box::new(xz2::read::XzDecoder::new_multi_decoder(r)),
158 "bzip2" => Box::new(bzip2::read::MultiBzDecoder::new(r)),
159 "gzip" => Box::new(flate2::read::MultiGzDecoder::new(r)),
160 "br" => Box::new(brotli_decompressor::Decompressor::new(r, 1 << 16)),
161 m => return Err(io::Error::other(format!("unknown compression {m:?}"))),
162 })
163}
164
165const MAGIC: &[(&[u8], &str)] = &[
166 (b"\x28\xb5\x2f\xfd", "zstd"),
167 (b"\xfd7zXZ\x00", "xz"),
168 (b"\x1f\x8b", "gzip"),
169];
170
171/// Both caches store narinfo/.ls zstd-compressed (Content-Encoding), but the
172/// CDN sometimes serves them without that header: sniff the magic too.
173pub fn decode(body: &[u8], encoding: Option<&str>) -> io::Result<Vec<u8>> {
174 let enc = match encoding {
175 None | Some("") | Some("identity") => {
176 match MAGIC.iter().find(|(m, _)| body.starts_with(m)) {
177 Some((_, t)) => *t,
178 None => return Ok(body.to_vec()),
179 }
180 }
181 Some(e) => e,
182 };
183 let mut out = Vec::new();
184 decoder(enc, body)?.read_to_end(&mut out)?;
185 Ok(out)
186}
187
188#[cfg(test)]
189mod tests {
190 use super::*;
191
192 #[test]
193 fn decode_sniffs_zstd() {
194 let z = zstd::encode_all(&b"{\"root\":1}"[..], 3).unwrap();
195 assert_eq!(decode(&z, None).unwrap(), b"{\"root\":1}");
196 assert_eq!(decode(&z, Some("zstd")).unwrap(), b"{\"root\":1}");
197 assert_eq!(decode(b"plain", None).unwrap(), b"plain");
198 }
199}