Skip to main content

zoi_lua/api/
download.rs

1//! File download utilities with progress reporting for the Lua environment.
2//!
3//! This module provides functions to download files from URLs, including
4//! support for progress bars and hash verification of downloaded files.
5
6use std::fs;
7use std::io::{Read, Write};
8use std::path::Path;
9
10use colored::Colorize;
11use indicatif::{ProgressBar, ProgressStyle};
12use zoi_core::utils;
13
14/// Downloads a file from the given URL to the destination path with a progress
15/// bar.
16///
17/// # Errors
18///
19/// Returns an `mlua::Error` if the download fails after several attempts or if
20/// there is an error creating the destination file.
21pub fn download_with_progress(
22    url: &str,
23    dest_path: &Path,
24    quiet: bool
25) -> Result<(), mlua::Error> {
26    if url.starts_with("http://") && !quiet {
27        println!(
28            "{}: downloading over insecure HTTP: {url}",
29            "Warning:".yellow()
30        );
31    }
32
33    let client = utils::get_http_client()
34        .map_err(|e| mlua::Error::RuntimeError(e.to_string()))?;
35
36    let mut attempt = 0u32;
37    let mut response = loop {
38        attempt += 1;
39        match client.get(url).send() {
40            Ok(resp) => {
41                if !resp.status().is_success() {
42                    return Err(mlua::Error::RuntimeError(format!(
43                        "Failed to download {url}: {}",
44                        resp.status()
45                    )));
46                }
47                break resp;
48            }
49            Err(e) => {
50                if attempt < 3 {
51                    if !quiet {
52                        eprintln!("Download failed ({e}). Retrying...");
53                    }
54                    zoi_core::utils::retry_backoff_sleep(attempt);
55                    continue;
56                }
57                return Err(mlua::Error::RuntimeError(e.to_string()));
58            }
59        }
60    };
61
62    let total_size = response.content_length().unwrap_or(0);
63
64    let pb = if quiet {
65        None
66    } else {
67        let pb = ProgressBar::new(total_size);
68        pb.set_style(
69            ProgressStyle::default_bar()
70                .template(
71                    "{spinner:.green} {msg:30.cyan} [{bar:40.cyan/blue}] \
72                     {bytes}/{total_bytes} ({bytes_per_sec}, \
73                     {elapsed_precise})"
74                )
75                .map_err(|e| mlua::Error::RuntimeError(e.to_string()))?
76                .progress_chars("=>-")
77        );
78
79        let filename = url.split('/').next_back().unwrap_or("file");
80        pb.set_message(format!("Downloading {filename}"));
81        Some(pb)
82    };
83
84    let mut dest_file = fs::File::create(dest_path)
85        .map_err(|e| mlua::Error::RuntimeError(e.to_string()))?;
86
87    let mut buffer = [0; 8192];
88    let mut downloaded = 0;
89
90    while let Ok(n) = response.read(&mut buffer) {
91        if n == 0 {
92            break;
93        }
94        dest_file
95            .write_all(buffer.get(..n).ok_or_else(|| {
96                mlua::Error::RuntimeError(
97                    "buffer slice out of bounds".to_string()
98                )
99            })?)
100            .map_err(|e| mlua::Error::RuntimeError(e.to_string()))?;
101        downloaded += n as u64;
102        if let Some(ref p) = pb {
103            p.set_position(downloaded);
104        }
105    }
106
107    if let Some(p) = pb {
108        p.finish_and_clear();
109    }
110
111    Ok(())
112}
113
114/// Registers the `UTILS.DOWNLOAD` function in the Lua environment.
115///
116/// # Errors
117///
118/// Returns an `mlua::Error` if the `UTILS` table cannot be found or if there is
119/// an error setting the `DOWNLOAD` function.
120pub fn add_download_util(
121    lua: &mlua::Lua,
122    quiet: bool
123) -> Result<(), mlua::Error> {
124    let download_fn = lua.create_function(
125        move |lua,
126              (url, out_name, hash): (
127            String,
128            Option<String>,
129            Option<String>
130        )| {
131            let build_dir_str: String = lua.globals().get("BUILD_DIR")?;
132            let build_dir = Path::new(&build_dir_str);
133
134            let filename = out_name.unwrap_or_else(|| {
135                url.split('/')
136                    .next_back()
137                    .unwrap_or("download.tmp")
138                    .to_string()
139            });
140
141            let dest_path = build_dir.join(&filename);
142
143            download_with_progress(&url, &dest_path, quiet)?;
144
145            if let Some(hash_spec) = hash {
146                let parts: Vec<&str> = hash_spec.splitn(2, '-').collect();
147                let (algo, expected_hash) = if parts.len() == 2 {
148                    (
149                        parts.first().copied().unwrap_or_default(),
150                        parts.get(1).copied().unwrap_or_default()
151                    )
152                } else {
153                    ("sha512", hash_spec.as_str())
154                };
155
156                let Some(hash_algo) =
157                    zoi_core::hash::HashAlgorithm::from_name(algo)
158                else {
159                    return Err(mlua::Error::RuntimeError(format!(
160                        "Unsupported hash algorithm: {algo}"
161                    )));
162                };
163
164                let actual_hash =
165                    zoi_core::hash::calculate_file_hash(&dest_path, hash_algo)
166                        .map_err(|e| {
167                            mlua::Error::RuntimeError(format!(
168                                "Failed to calculate hash: {e}"
169                            ))
170                        })?;
171
172                if !actual_hash.eq_ignore_ascii_case(expected_hash) {
173                    return Err(mlua::Error::RuntimeError(format!(
174                        "Hash mismatch for {filename}. Expected: \
175                         {algo}-{expected_hash}, Got: {algo}-{actual_hash}"
176                    )));
177                }
178            }
179
180            Ok(filename)
181        }
182    )?;
183
184    let utils_table: mlua::Table = lua.globals().get("UTILS")?;
185    utils_table.set("DOWNLOAD", download_fn)?;
186
187    Ok(())
188}