1mod credential;
2mod listing;
3
4use std::path::Path;
5use std::sync::Arc;
6use std::time::Duration;
7
8use base64::Engine;
9use reqwest::{Method, StatusCode};
10use tokio::io::{AsyncReadExt, AsyncSeekExt};
11
12use super::{Entry, content_length, expect_success, read_retrying, unreachable_store};
13use crate::config::AzureCredential;
14use crate::error::Error;
15
16pub(crate) use credential::{Credential, Scope};
17
18const SINGLE_PUT_CEILING: u64 = 5000 * 1024 * 1024;
19const BLOCK: u64 = 64 * 1024 * 1024;
20const LARGEST_BLOCK: u64 = 4000 * 1024 * 1024;
21const MOST_BLOCKS: u64 = 50_000;
22const SERVER_PERMISSIONS: &str = "racwdl";
23
24pub struct AzureConfig {
25 pub endpoint: String,
26 pub account: String,
27 pub container: String,
28 pub credential: AzureCredential,
29 pub lifetime: Duration,
30}
31
32#[derive(Clone)]
33pub struct AzureKeys {
34 container_url: reqwest::Url,
35 container: String,
36 credential: Arc<Credential>,
37 client: reqwest::Client,
38 lifetime: Duration,
39}
40
41fn block_size(length: u64) -> u64 {
42 BLOCK.max(length.div_ceil(MOST_BLOCKS)).min(LARGEST_BLOCK)
43}
44
45fn upload_tag() -> u64 {
46 let mut tag = [0u8; 8];
47 getrandom::fill(&mut tag).expect("the operating system has a random number generator");
48 u64::from_le_bytes(tag)
49}
50
51fn block_id(upload: u64, index: u64) -> String {
52 base64::engine::general_purpose::STANDARD.encode(format!("{upload:016x}{index:08}"))
53}
54
55impl AzureKeys {
56 pub fn new(config: &AzureConfig) -> Result<Self, Error> {
57 crate::tls::install_crypto_provider();
58
59 let credential = Credential::new(&config.account, &config.credential)?;
60 Self::with_credential(config, credential)
61 }
62
63 pub(crate) fn with_credential(
64 config: &AzureConfig,
65 credential: Credential,
66 ) -> Result<Self, Error> {
67 let mut container_url = reqwest::Url::parse(&config.endpoint)
68 .map_err(|_| Error::Misconfigured("LFSX_AZURE_ENDPOINT is not a URL"))?;
69 container_url
70 .path_segments_mut()
71 .map_err(|()| Error::Misconfigured("LFSX_AZURE_ENDPOINT is not a URL"))?
72 .pop_if_empty()
73 .push(&config.container);
74
75 Ok(Self {
76 container_url,
77 container: config.container.clone(),
78 credential: Arc::new(credential),
79 client: reqwest::Client::new(),
80 lifetime: config.lifetime,
81 })
82 }
83
84 fn url(&self, blob: Option<&str>, query: &[(&str, &str)]) -> reqwest::Url {
85 let mut url = self.container_url.clone();
86 if let Some(blob) = blob
87 && let Ok(mut segments) = url.path_segments_mut()
88 {
89 segments.extend(blob.split('/'));
90 }
91 if !query.is_empty() {
92 url.query_pairs_mut().extend_pairs(query);
93 }
94 url
95 }
96
97 fn with_query(mut url: reqwest::Url, extra: &str) -> reqwest::Url {
98 let query = match url.query() {
99 Some(existing) if !existing.is_empty() => format!("{existing}&{extra}"),
100 _ => extra.to_owned(),
101 };
102 url.set_query(Some(&query));
103 url
104 }
105
106 async fn request(
107 &self,
108 method: Method,
109 blob: Option<&str>,
110 query: &[(&str, &str)],
111 ) -> Result<reqwest::RequestBuilder, Error> {
112 let url = self.url(blob, query);
113
114 let request = match self.credential.as_ref() {
115 Credential::Key { .. } => {
116 let sas = self
117 .credential
118 .sas(
119 &self.container,
120 Scope::Container,
121 SERVER_PERMISSIONS,
122 self.lifetime,
123 )
124 .ok_or(Error::Misconfigured("the Azure account key could not sign"))?;
125 self.client.request(method, Self::with_query(url, &sas))
126 }
127 Credential::Sas(token) => self.client.request(method, Self::with_query(url, token)),
128 Credential::Identity(identity) => self
129 .client
130 .request(method, url)
131 .bearer_auth(identity.token(&self.client).await?),
132 };
133
134 Ok(request.header("x-ms-version", credential::VERSION))
135 }
136
137 pub(crate) async fn reachable(&self) -> Result<(), Error> {
138 let request = self
139 .request(
140 Method::GET,
141 None,
142 &[
143 ("restype", "container"),
144 ("comp", "list"),
145 ("maxresults", "1"),
146 ],
147 )
148 .await?;
149 let response = read_retrying(request).await?;
150
151 if !response.status().is_success() {
152 return Err(Error::Storage(std::io::Error::other(format!(
153 "the object store answered {} for the container",
154 response.status()
155 ))));
156 }
157
158 Ok(())
159 }
160
161 pub(crate) fn signed_download(&self, key: &str) -> Option<String> {
162 let sas = self
163 .credential
164 .sas(&self.container, Scope::Blob(key), "r", self.lifetime)?;
165
166 Some(Self::with_query(self.url(Some(key), &[]), &sas).to_string())
167 }
168
169 pub(crate) async fn get_range(
170 &self,
171 key: &str,
172 start: u64,
173 length: u64,
174 ) -> Result<reqwest::Response, Error> {
175 let response = self
176 .request(Method::GET, Some(key), &[])
177 .await?
178 .header(
179 "x-ms-range",
180 format!("bytes={start}-{}", start + length.saturating_sub(1)),
181 )
182 .send()
183 .await
184 .map_err(|_| unreachable_store())?;
185
186 if !response.status().is_success() {
187 return Err(Error::NotFound);
188 }
189
190 Ok(response)
191 }
192
193 pub(crate) async fn head(&self, key: &str) -> Result<u64, Error> {
194 let response = read_retrying(self.request(Method::HEAD, Some(key), &[]).await?).await?;
195
196 if !response.status().is_success() {
197 return Err(Error::NotFound);
198 }
199
200 content_length(&response)
201 }
202
203 async fn put_blob(
204 &self,
205 key: &str,
206 body: reqwest::Body,
207 length: u64,
208 only_if_absent: bool,
209 ) -> Result<reqwest::Response, Error> {
210 let mut request = self
211 .request(Method::PUT, Some(key), &[])
212 .await?
213 .header("x-ms-blob-type", "BlockBlob")
214 .header(reqwest::header::CONTENT_LENGTH, length);
215 if only_if_absent {
216 request = request.header(reqwest::header::IF_NONE_MATCH, "*");
217 }
218
219 request
220 .body(body)
221 .send()
222 .await
223 .map_err(|_| unreachable_store())
224 }
225
226 pub(crate) async fn put(
227 &self,
228 key: &str,
229 body: reqwest::Body,
230 length: u64,
231 ) -> Result<(), Error> {
232 let response = self.put_blob(key, body, length, false).await?;
233 expect_success(response, "write").await?;
234
235 Ok(())
236 }
237
238 pub(crate) async fn put_file(&self, key: &str, staged: &Path) -> Result<(), Error> {
239 let file = tokio::fs::File::open(staged).await?;
240 let length = file.metadata().await?.len();
241
242 if length > SINGLE_PUT_CEILING {
243 drop(file);
244 return self
245 .put_in_blocks(key, staged, length, block_size(length))
246 .await;
247 }
248
249 let stream = tokio_util::io::ReaderStream::new(file);
250 self.put(key, reqwest::Body::wrap_stream(stream), length)
251 .await
252 }
253
254 pub(crate) async fn put_in_blocks(
255 &self,
256 key: &str,
257 staged: &Path,
258 length: u64,
259 size: u64,
260 ) -> Result<(), Error> {
261 let count = length.div_ceil(size);
262 if count > MOST_BLOCKS {
263 return Err(Error::Storage(std::io::Error::other(
264 "this object needs more blocks than one blob may have",
265 )));
266 }
267
268 tracing::info!(
269 key,
270 length,
271 block_size = size,
272 "an object over the single-request ceiling is going up in blocks"
273 );
274
275 let upload = upload_tag();
276 let mut ids = Vec::new();
277 for index in 0..count {
278 let offset = index * size;
279 let this = size.min(length - offset);
280
281 let mut file = tokio::fs::File::open(staged).await?;
282 file.seek(std::io::SeekFrom::Start(offset)).await?;
283 let stream = tokio_util::io::ReaderStream::new(file.take(this));
284
285 let id = block_id(upload, index);
286 let response = self
287 .request(
288 Method::PUT,
289 Some(key),
290 &[("comp", "block"), ("blockid", &id)],
291 )
292 .await?
293 .header(reqwest::header::CONTENT_LENGTH, this)
294 .body(reqwest::Body::wrap_stream(stream))
295 .send()
296 .await
297 .map_err(|_| unreachable_store())?;
298 expect_success(response, "write a block").await?;
299
300 ids.push(id);
301 }
302
303 let list = format!(
304 "<?xml version=\"1.0\" encoding=\"utf-8\"?><BlockList>{}</BlockList>",
305 ids.iter()
306 .map(|id| format!("<Latest>{id}</Latest>"))
307 .collect::<String>()
308 );
309
310 let response = self
311 .request(Method::PUT, Some(key), &[("comp", "blocklist")])
312 .await?
313 .header(reqwest::header::CONTENT_LENGTH, list.len())
314 .body(list)
315 .send()
316 .await
317 .map_err(|_| unreachable_store())?;
318 expect_success(response, "assemble a blob from its blocks").await?;
319
320 Ok(())
321 }
322
323 pub(crate) async fn put_if_absent(&self, key: &str, body: Vec<u8>) -> Result<bool, Error> {
324 let length = body.len() as u64;
325 let response = self
326 .put_blob(key, reqwest::Body::from(body), length, true)
327 .await?;
328
329 if matches!(
330 response.status(),
331 StatusCode::CONFLICT | StatusCode::PRECONDITION_FAILED
332 ) {
333 return Ok(false);
334 }
335
336 expect_success(response, "write").await?;
337
338 Ok(true)
339 }
340
341 pub(crate) async fn get_bytes(&self, key: &str) -> Result<Option<Vec<u8>>, Error> {
342 let response = read_retrying(self.request(Method::GET, Some(key), &[]).await?).await?;
343
344 if response.status() == StatusCode::NOT_FOUND {
345 return Ok(None);
346 }
347
348 expect_success(response, "read")
349 .await?
350 .bytes()
351 .await
352 .map(|bytes| Some(bytes.to_vec()))
353 .map_err(|_| unreachable_store())
354 }
355
356 pub(crate) async fn delete(&self, key: &str) -> Result<bool, Error> {
357 let response = self
358 .request(Method::DELETE, Some(key), &[])
359 .await?
360 .send()
361 .await
362 .map_err(|_| unreachable_store())?;
363
364 if response.status() == StatusCode::NOT_FOUND {
365 return Ok(false);
366 }
367
368 expect_success(response, "delete").await?;
369
370 Ok(true)
371 }
372
373 pub(crate) async fn entries(&self, prefix: &str) -> Result<Vec<Entry>, Error> {
374 let mut out = Vec::new();
375 let mut marker: Option<String> = None;
376
377 loop {
378 let mut query = vec![
379 ("restype", "container"),
380 ("comp", "list"),
381 ("prefix", prefix),
382 ];
383 if let Some(marker) = &marker {
384 query.push(("marker", marker));
385 }
386
387 let response = read_retrying(self.request(Method::GET, None, &query).await?).await?;
388 let body = expect_success(response, "list")
389 .await?
390 .text()
391 .await
392 .map_err(|_| unreachable_store())?;
393
394 let page = listing::parse(&body)?;
395 out.extend(page.entries);
396
397 match page.next {
398 Some(next) => marker = Some(next),
399 None => break,
400 }
401 }
402
403 Ok(out)
404 }
405}
406
407#[cfg(test)]
408mod tests;