nd_util/
download_to_path.rs

1use crate::DropRemovePathBlocking;
2use anyhow::Context;
3use std::io::Write;
4use std::path::Path;
5use std::path::PathBuf;
6use tracing::warn;
7
8const LOCKING_SUPPORTED: bool = cfg!(unix) || cfg!(windows);
9
10fn download_to_path_blocking(
11    handle: tokio::runtime::Handle,
12    client: &reqwest::Client,
13    url: &str,
14    path: PathBuf,
15) -> anyhow::Result<()> {
16    // Create temporary path.
17    let temporary_path = path.with_added_extension("part");
18
19    // Setup to open the temporary file.
20    //
21    // We do NOT use mandatory locking on Windows.
22    // This is because the file would need to be dropped to be renamed,
23    // which leads to a race as we must release ALL locks to do so.
24    //
25    // TODO: On linux, consider probing for O_TMPFILE support somehow and create an unnamed tmp file and use linkat.
26    let mut open_options = std::fs::OpenOptions::new();
27    open_options.write(true);
28
29    // If we don't have a mechanism to prevent stomping,
30    // at least ensure that we can't stomp.
31    if LOCKING_SUPPORTED {
32        // We prevent stomping by locking somehow.
33        // Create and overwrite the temporary file.
34        open_options.create(true);
35    } else {
36        // If the temporary file exists, return an error.
37        open_options.create_new(true);
38    }
39
40    // Open the temporary file.
41    let mut temporary_file = open_options
42        .open(&temporary_path)
43        .context("failed to create temporary file")?;
44
45    // Create the remove handle for the temporary path.
46    let mut temporary_path = DropRemovePathBlocking::new(temporary_path);
47
48    // We fall back to create_new for temp file creation if we know we are on an unsupported platform.
49    // Therefore, we don't need to do anything special here if we are on an unsupported platform.
50    // TODO: Check for and handle unsupported errors.
51    if LOCKING_SUPPORTED {
52        temporary_file
53            .try_lock()
54            .context("failed to lock temporary file")?;
55    }
56
57    let result = (|| {
58        // Send the request
59        let mut response = handle
60            .block_on(client.get(url).send())
61            .context("failed to get headers")?
62            .error_for_status()?;
63
64        // Download the file chunk-by-chunk
65        while let Some(chunk) = handle
66            .block_on(response.chunk())
67            .context("failed to get next chunk")?
68        {
69            temporary_file
70                .write_all(&chunk)
71                .context("failed to write to file")?;
72        }
73
74        // Sync data
75        temporary_file.flush().context("failed to flush file")?;
76        temporary_file
77            .sync_all()
78            .context("failed to sync file data")?;
79
80        // Perform rename from temporary file path to actual file path.
81        std::fs::rename(&temporary_path, &path).context("failed to rename temporary file")?;
82
83        Ok(())
84    })();
85
86    // Since we renamed (or failed), we can unlock the file and drop it.
87    if LOCKING_SUPPORTED {
88        temporary_file.unlock()?;
89    }
90    drop(temporary_file);
91
92    match result.as_ref() {
93        Ok(()) => {
94            // Persist the file,
95            // since it was renamed and we don't want to remove a non-existent file.
96            temporary_path.persist();
97        }
98        Err(_error) => {
99            // Try to clean up the temporary file before returning.
100            if let Err((mut temporary_path, error)) = temporary_path.try_drop() {
101                // Don't try to delete the file again.
102                temporary_path.persist();
103
104                // Returning the original error is more important,
105                // so we just log the temporary file error here.
106                warn!("failed to delete temporary file: \"{error}\"");
107            }
108        }
109    }
110
111    result
112}
113
114/// Using the given client, download the file at a url to a given path.
115///
116/// Note that this function will overwrite the file at the given path.
117///
118/// # Temporary Files
119/// This will create a temporary ".part" file in the same directory while downloading.
120/// On failure, this file will try to be cleaned up.
121/// On success, this temporary file will be renamed to the actual file name.
122/// As a result, it may be assumed that the file created at the given path is the complete, non-erroneus download.
123///
124/// # Locking
125/// During downloads, the temporary file is locked via advisory locking on platforms that support it.
126/// If locking is not supported, overwriting a pre-existing temporary file causes an error.
127/// Currently, Unix and Windows support advisory locking.
128pub async fn download_to_path<P>(client: &reqwest::Client, url: &str, path: P) -> anyhow::Result<()>
129where
130    P: AsRef<Path>,
131{
132    let handle = tokio::runtime::Handle::try_current()?;
133    let client = client.clone();
134    let url = url.to_string();
135    let path = path.as_ref().to_path_buf();
136    tokio::task::spawn_blocking(move || download_to_path_blocking(handle, &client, &url, path))
137        .await??;
138    Ok(())
139}
140
141#[cfg(test)]
142mod test {
143    use super::*;
144
145    #[tokio::test]
146    async fn it_works() {
147        tokio::fs::create_dir_all("test_tmp")
148            .await
149            .expect("failed to create tmp dir");
150
151        let client = reqwest::Client::new();
152        download_to_path(&client, "http://google.com", "test_tmp/google.html")
153            .await
154            .expect("failed to download");
155    }
156}