forked from
lass.pds.lassul.us/superconfig
This repository has no description
7.7 kB
212 lines
1//! Binary cache access: GET with retries, and body decoding.
2
3use std::io;
4use std::pin::Pin;
5use std::time::Duration;
6
7use async_compression::tokio::bufread::{
8 BrotliDecoder, BzDecoder, GzipDecoder, XzDecoder, ZstdDecoder,
9};
10use futures_util::TryStreamExt;
11use tokio::io::{AsyncBufRead, AsyncRead, AsyncReadExt, BufReader};
12use tokio::time::sleep;
13use tokio_util::io::StreamReader;
14
15const TRIES: u32 = 6;
16
17pub type Body = Pin<Box<dyn AsyncRead + Send>>;
18
19pub struct Http {
20 client: reqwest::Client,
21}
22
23pub enum Error {
24 /// 403/404, or a missing file for file:// caches
25 NotFound,
26 Other(String),
27}
28
29impl std::fmt::Display for Error {
30 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
31 match self {
32 Error::NotFound => write!(f, "not found"),
33 Error::Other(e) => write!(f, "{e}"),
34 }
35 }
36}
37
38impl From<Error> for io::Error {
39 fn from(e: Error) -> io::Error {
40 io::Error::other(e.to_string())
41 }
42}
43
44impl Http {
45 /// `caches` are base URLs. Their hosts are resolved right away and then
46 /// served from memory: once the FUSE file system is mounted on
47 /// /nix/store, getaddrinfo() could end up looking for NSS modules in
48 /// the store, i.e. in a request to ourselves.
49 pub async fn new<'a>(caches: impl IntoIterator<Item = &'a str>) -> Http {
50 // rustls with bundled roots: nothing read from /etc/ssl or the store
51 let roots = webpki_root_certs::TLS_SERVER_ROOT_CERTS
52 .iter()
53 .filter_map(|c| reqwest::Certificate::from_der(c).ok());
54 let mut b = reqwest::Client::builder()
55 .connect_timeout(Duration::from_secs(30))
56 .read_timeout(Duration::from_secs(60))
57 .pool_max_idle_per_host(64)
58 .tls_certs_only(roots);
59 for c in caches {
60 let Ok(u) = reqwest::Url::parse(c) else {
61 continue;
62 };
63 let (Some(host), Some(port)) = (u.host_str(), u.port_or_known_default()) else {
64 continue;
65 };
66 match tokio::net::lookup_host((host, port)).await {
67 Ok(a) => b = b.resolve_to_addrs(host, &a.collect::<Vec<_>>()),
68 Err(e) => eprintln!("resolving {host}:{port}: {e}"),
69 };
70 }
71 Http {
72 client: b.build().expect("building the HTTP client"),
73 }
74 }
75
76 /// GET `url`, retrying 429/5xx and transport errors with backoff (the
77 /// caches answer 503 under load). Returns the body reader and its
78 /// Content-Encoding.
79 pub async fn open(&self, url: &str) -> Result<(Body, Option<String>), Error> {
80 if let Some(path) = url.strip_prefix("file://") {
81 return match tokio::fs::File::open(path).await {
82 Ok(f) => Ok((Box::pin(f), None)),
83 Err(e) if e.kind() == io::ErrorKind::NotFound => Err(Error::NotFound),
84 Err(e) => Err(Error::Other(format!("{url}: {e}"))),
85 };
86 }
87 let mut n = 0;
88 loop {
89 let err = match self.client.get(url).send().await {
90 Ok(r) if r.status().is_success() => {
91 let enc = r
92 .headers()
93 .get("content-encoding")
94 .and_then(|v| v.to_str().ok())
95 .map(str::to_string);
96 let body = r.bytes_stream().map_err(io::Error::other);
97 return Ok((Box::pin(StreamReader::new(body)), enc));
98 }
99 Ok(r) => match r.status().as_u16() {
100 403 | 404 => return Err(Error::NotFound),
101 code if code != 429 && code < 500 => {
102 return Err(Error::Other(format!("{url}: HTTP {code}")))
103 }
104 code => format!("HTTP {code}"),
105 },
106 Err(e) => e.to_string(),
107 };
108 n += 1;
109 if n == TRIES {
110 return Err(Error::Other(format!("{url}: {err}")));
111 }
112 eprintln!("retry {n}/{} {url}: {err}", TRIES - 1);
113 sleep(Duration::from_secs((1u64 << (n - 1)).min(30))).await;
114 }
115 }
116
117 /// GET `url` -> decoded body. Transfer errors while reading the body are
118 /// retried like failed requests.
119 pub async fn get(&self, url: &str) -> Result<(Vec<u8>, usize), Error> {
120 let mut n = 0;
121 loop {
122 let (mut r, enc) = self.open(url).await?;
123 let mut body = Vec::new();
124 match r.read_to_end(&mut body).await {
125 Ok(_) => {
126 let len = body.len();
127 let dec = decode(&body, enc.as_deref())
128 .await
129 .map_err(|e| Error::Other(format!("{url}: decoding: {e}")))?;
130 return Ok((dec, len));
131 }
132 Err(e) => {
133 n += 1;
134 if n == TRIES {
135 return Err(Error::Other(format!("{url}: {e}")));
136 }
137 eprintln!("retry {n}/{} {url}: {e}", TRIES - 1);
138 sleep(Duration::from_secs((1u64 << (n - 1)).min(30))).await;
139 }
140 }
141 }
142 }
143}
144
145/// Wrap `r` in a decompressor for `method` (narinfo Compression or
146/// Content-Encoding). Concatenated streams are decoded as one, like Nix.
147pub fn decoder<'a>(
148 method: &str,
149 r: impl AsyncBufRead + Send + 'a,
150) -> io::Result<Pin<Box<dyn AsyncRead + Send + 'a>>> {
151 fn multi<D: AsyncRead + Send>(mut d: D, set: impl FnOnce(&mut D)) -> D {
152 set(&mut d);
153 d
154 }
155 Ok(match method {
156 "" | "none" | "identity" => Box::pin(r),
157 "zstd" => Box::pin(multi(ZstdDecoder::new(r), |d| d.multiple_members(true))),
158 "xz" => Box::pin(multi(XzDecoder::new(r), |d| d.multiple_members(true))),
159 "bzip2" => Box::pin(multi(BzDecoder::new(r), |d| d.multiple_members(true))),
160 "gzip" => Box::pin(multi(GzipDecoder::new(r), |d| d.multiple_members(true))),
161 "br" => Box::pin(BrotliDecoder::new(r)),
162 m => return Err(io::Error::other(format!("unknown compression {m:?}"))),
163 })
164}
165
166/// `decoder` for an unbuffered body.
167pub fn decode_body<'a>(
168 method: &str,
169 r: impl AsyncRead + Send + 'a,
170) -> io::Result<Pin<Box<dyn AsyncRead + Send + 'a>>> {
171 decoder(method, BufReader::with_capacity(1 << 16, r))
172}
173
174const MAGIC: &[(&[u8], &str)] = &[
175 (b"\x28\xb5\x2f\xfd", "zstd"),
176 (b"\xfd7zXZ\x00", "xz"),
177 (b"\x1f\x8b", "gzip"),
178];
179
180/// Both caches store narinfo/.ls zstd-compressed (Content-Encoding), but the
181/// CDN sometimes serves them without that header: sniff the magic too.
182pub async fn decode(body: &[u8], encoding: Option<&str>) -> io::Result<Vec<u8>> {
183 let enc = match encoding {
184 None | Some("") | Some("identity") => {
185 match MAGIC.iter().find(|(m, _)| body.starts_with(m)) {
186 Some((_, t)) => *t,
187 None => return Ok(body.to_vec()),
188 }
189 }
190 Some(e) => e,
191 };
192 let mut out = Vec::new();
193 decoder(enc, body)?.read_to_end(&mut out).await?;
194 Ok(out)
195}
196
197#[cfg(test)]
198mod tests {
199 use super::*;
200
201 #[tokio::test]
202 async fn decode_sniffs_zstd() {
203 let mut z = Vec::new();
204 async_compression::tokio::bufread::ZstdEncoder::new(&b"{\"root\":1}"[..])
205 .read_to_end(&mut z)
206 .await
207 .unwrap();
208 assert_eq!(decode(&z, None).await.unwrap(), b"{\"root\":1}");
209 assert_eq!(decode(&z, Some("zstd")).await.unwrap(), b"{\"root\":1}");
210 assert_eq!(decode(b"plain", None).await.unwrap(), b"plain");
211 }
212}