Skip to main content

ic_testkit/pic/
runtime.rs

1use flate2::read::GzDecoder;
2use sha2::{Digest, Sha256};
3use std::{
4    env,
5    fs::{self, OpenOptions},
6    io::{self, Read},
7    path::{Path, PathBuf},
8    process,
9    time::{SystemTime, UNIX_EPOCH},
10};
11
12#[cfg(unix)]
13use std::os::unix::fs::PermissionsExt;
14
15use super::startup::PicStartError;
16
17const POCKET_IC_BIN_ENV: &str = "POCKET_IC_BIN";
18const ALLOW_DOWNLOAD_ENV: &str = "IC_TESTKIT_ALLOW_POCKET_IC_DOWNLOAD";
19
20const SERVER_NAME: &str = "pocket-ic";
21
22///
23/// Runtime policy for resolving the PocketIC server binary used by [`PicBuilder`](super::PicBuilder).
24///
25#[derive(Clone, Debug, Default, Eq, PartialEq)]
26pub struct PicRuntimeConfig {
27    pocket_ic_bin: Option<PathBuf>,
28    cache_dir: Option<PathBuf>,
29    allow_download: bool,
30    server_sha256: Option<String>,
31}
32
33impl PicRuntimeConfig {
34    /// Build a runtime config from supported environment variables.
35    ///
36    /// Supported variables:
37    ///
38    /// - `POCKET_IC_BIN`: explicit trusted PocketIC server binary path
39    /// - `IC_TESTKIT_ALLOW_POCKET_IC_DOWNLOAD=1`: allow network download on cache miss
40    #[must_use]
41    pub fn from_env() -> Self {
42        Self {
43            pocket_ic_bin: non_empty_env_path(POCKET_IC_BIN_ENV),
44            cache_dir: None,
45            allow_download: env::var(ALLOW_DOWNLOAD_ENV).is_ok_and(|value| env_truthy(&value)),
46            server_sha256: None,
47        }
48    }
49
50    /// Set an explicit PocketIC server binary path.
51    #[must_use]
52    pub fn pocket_ic_bin(mut self, path: impl Into<PathBuf>) -> Self {
53        self.pocket_ic_bin = Some(path.into());
54        self
55    }
56
57    /// Set the cache root used for resolved PocketIC server binaries.
58    #[must_use]
59    pub fn cache_dir(mut self, path: impl Into<PathBuf>) -> Self {
60        self.cache_dir = Some(path.into());
61        self
62    }
63
64    /// Allow or deny downloading the pinned PocketIC server binary on cache miss.
65    #[must_use]
66    pub const fn allow_download(mut self, allow_download: bool) -> Self {
67        self.allow_download = allow_download;
68        self
69    }
70
71    /// Set an expected SHA-256 digest for the ungzipped PocketIC server binary.
72    #[must_use]
73    pub fn server_sha256(mut self, sha256: impl Into<String>) -> Self {
74        self.server_sha256 = Some(sha256.into().trim().to_ascii_lowercase());
75        self
76    }
77
78    /// Resolve, validate, and optionally download the PocketIC server binary.
79    pub fn ensure_binary(&self) -> Result<PathBuf, PicStartError> {
80        if let Some(path) = self.pocket_ic_bin.as_deref() {
81            return self.validate_binary(path);
82        }
83
84        let path = self.cache_binary_path();
85        if path.exists() {
86            return self.validate_binary(&path);
87        }
88
89        if !self.allow_download {
90            return Err(PicStartError::BinaryUnavailable {
91                message: missing_binary_message(&path),
92            });
93        }
94
95        self.download_binary(&path)?;
96        self.validate_binary(&path)
97    }
98
99    fn cache_binary_path(&self) -> PathBuf {
100        self.cache_root()
101            .join(format!(
102                "pocket-ic-server-{}",
103                pocket_ic::LATEST_SERVER_VERSION
104            ))
105            .join(SERVER_NAME)
106    }
107
108    fn cache_root(&self) -> PathBuf {
109        self.cache_dir.clone().unwrap_or_else(default_cache_root)
110    }
111
112    fn validate_binary(&self, path: &Path) -> Result<PathBuf, PicStartError> {
113        if path.as_os_str().is_empty() {
114            return Err(PicStartError::BinaryUnavailable {
115                message: missing_binary_message(path),
116            });
117        }
118
119        let metadata = fs::metadata(path).map_err(|err| PicStartError::BinaryUnavailable {
120            message: format!(
121                "PocketIC server binary is unavailable at {}: {err}. {}",
122                path.display(),
123                setup_guidance()
124            ),
125        })?;
126
127        if !metadata.is_file() {
128            return Err(PicStartError::BinaryInvalid {
129                message: format!(
130                    "PocketIC server binary path {} is not a file.",
131                    path.display()
132                ),
133            });
134        }
135
136        #[cfg(unix)]
137        if metadata.permissions().mode() & 0o111 == 0 {
138            return Err(PicStartError::BinaryInvalid {
139                message: format!(
140                    "PocketIC server binary at {} is not executable.",
141                    path.display()
142                ),
143            });
144        }
145
146        if let Some(expected) = self.server_sha256.as_deref() {
147            validate_sha256(path, expected)?;
148        }
149
150        Ok(path.to_path_buf())
151    }
152
153    fn download_binary(&self, path: &Path) -> Result<(), PicStartError> {
154        let url = pocket_ic_download_url()?;
155        let parent = path.parent().ok_or_else(|| PicStartError::DownloadFailed {
156            message: format!(
157                "failed to resolve parent directory for PocketIC cache path {}",
158                path.display()
159            ),
160        })?;
161        fs::create_dir_all(parent).map_err(|err| PicStartError::DownloadFailed {
162            message: format!(
163                "failed to create PocketIC cache directory {}: {err}",
164                parent.display()
165            ),
166        })?;
167
168        if path.exists() {
169            return Ok(());
170        }
171
172        let temp_path = download_temp_path(parent);
173        let result = download_gzip_to_file(&url, &temp_path)
174            .and_then(|()| make_executable(&temp_path))
175            .and_then(|()| {
176                if let Some(expected) = self.server_sha256.as_deref() {
177                    validate_sha256(&temp_path, expected)?;
178                }
179                Ok(())
180            })
181            .and_then(|()| {
182                if path.exists() {
183                    fs::remove_file(&temp_path).map_err(download_failed)?;
184                    Ok(())
185                } else {
186                    fs::rename(&temp_path, path).map_err(download_failed)
187                }
188            });
189
190        if result.is_err() {
191            let _ = fs::remove_file(&temp_path);
192        }
193
194        result
195    }
196}
197
198pub(super) fn ensure_pocket_ic_bin_from_env() -> Result<PathBuf, PicStartError> {
199    PicRuntimeConfig::from_env().ensure_binary()
200}
201
202fn non_empty_env_path(name: &str) -> Option<PathBuf> {
203    env::var_os(name)
204        .filter(|value| !value.is_empty())
205        .map(PathBuf::from)
206}
207
208fn env_truthy(value: &str) -> bool {
209    matches!(
210        value.trim().to_ascii_lowercase().as_str(),
211        "1" | "true" | "yes" | "on"
212    )
213}
214
215fn default_cache_root() -> PathBuf {
216    env::temp_dir()
217}
218
219const fn setup_guidance() -> &'static str {
220    "Set POCKET_IC_BIN to an existing ungzipped executable PocketIC server binary, or set IC_TESTKIT_ALLOW_POCKET_IC_DOWNLOAD=1 to let ic-testkit download the pinned server binary into its cache."
221}
222
223fn missing_binary_message(path: &Path) -> String {
224    format!(
225        "PocketIC server binary is unavailable. Checked cache path {}. {}",
226        path.display(),
227        setup_guidance()
228    )
229}
230
231fn pocket_ic_download_url() -> Result<String, PicStartError> {
232    Ok(format!(
233        "https://github.com/dfinity/pocketic/releases/download/{}/pocket-ic-{}-{}.gz",
234        pocket_ic::LATEST_SERVER_VERSION,
235        pocket_ic_arch()?,
236        pocket_ic_os()?
237    ))
238}
239
240fn pocket_ic_arch() -> Result<&'static str, PicStartError> {
241    match env::consts::ARCH {
242        "aarch64" => Ok("arm64"),
243        "x86_64" => Ok("x86_64"),
244        arch => Err(PicStartError::DownloadFailed {
245            message: format!("PocketIC server download is unsupported on architecture {arch}."),
246        }),
247    }
248}
249
250fn pocket_ic_os() -> Result<&'static str, PicStartError> {
251    match env::consts::OS {
252        "linux" => Ok("linux"),
253        "macos" => Ok("darwin"),
254        os => Err(PicStartError::DownloadFailed {
255            message: format!("PocketIC server download is unsupported on operating system {os}."),
256        }),
257    }
258}
259
260fn download_temp_path(parent: &Path) -> PathBuf {
261    let nanos = SystemTime::now()
262        .duration_since(UNIX_EPOCH)
263        .map_or(0, |duration| duration.as_nanos());
264    parent.join(format!(".pocket-ic-download-{}-{nanos}.tmp", process::id()))
265}
266
267fn download_gzip_to_file(url: &str, path: &Path) -> Result<(), PicStartError> {
268    let response = reqwest::blocking::get(url)
269        .and_then(reqwest::blocking::Response::error_for_status)
270        .map_err(|err| PicStartError::DownloadFailed {
271            message: format!("failed to download PocketIC server from {url}: {err}"),
272        })?;
273    let bytes = response
274        .bytes()
275        .map_err(|err| PicStartError::DownloadFailed {
276            message: format!("failed to read PocketIC server download from {url}: {err}"),
277        })?;
278    let mut gz = GzDecoder::new(&bytes[..]);
279    let mut out = OpenOptions::new()
280        .write(true)
281        .create_new(true)
282        .open(path)
283        .map_err(download_failed)?;
284    io::copy(&mut gz, &mut out).map_err(download_failed)?;
285    Ok(())
286}
287
288#[cfg(unix)]
289fn make_executable(path: &Path) -> Result<(), PicStartError> {
290    let mut permissions = fs::metadata(path).map_err(download_failed)?.permissions();
291    permissions.set_mode(0o755);
292    fs::set_permissions(path, permissions).map_err(download_failed)
293}
294
295#[cfg(not(unix))]
296fn make_executable(_path: &Path) -> Result<(), PicStartError> {
297    Ok(())
298}
299
300fn validate_sha256(path: &Path, expected: &str) -> Result<(), PicStartError> {
301    if !is_sha256_hex(expected) {
302        return Err(PicStartError::BinaryInvalid {
303            message:
304                "PocketIC server SHA-256 must be a 64-character lowercase or uppercase hex digest"
305                    .to_string(),
306        });
307    }
308
309    let actual = sha256_file(path).map_err(|err| PicStartError::BinaryInvalid {
310        message: format!(
311            "failed to calculate SHA-256 for PocketIC server binary {}: {err}",
312            path.display()
313        ),
314    })?;
315
316    if actual != expected {
317        return Err(PicStartError::BinaryInvalid {
318            message: format!(
319                "PocketIC server binary {} has SHA-256 {actual}, expected {expected}.",
320                path.display()
321            ),
322        });
323    }
324
325    Ok(())
326}
327
328fn is_sha256_hex(value: &str) -> bool {
329    value.len() == 64 && value.bytes().all(|byte| byte.is_ascii_hexdigit())
330}
331
332fn sha256_file(path: &Path) -> io::Result<String> {
333    let mut file = fs::File::open(path)?;
334    let mut hasher = Sha256::new();
335    let mut buffer = vec![0_u8; 64 * 1024].into_boxed_slice();
336
337    loop {
338        let bytes = file.read(&mut buffer)?;
339        if bytes == 0 {
340            break;
341        }
342        hasher.update(&buffer[..bytes]);
343    }
344
345    Ok(format!("{:x}", hasher.finalize()))
346}
347
348fn download_failed(err: io::Error) -> PicStartError {
349    PicStartError::DownloadFailed {
350        message: err.to_string(),
351    }
352}
353
354#[cfg(test)]
355mod tests {
356    use super::{PicRuntimeConfig, env_truthy, is_sha256_hex, missing_binary_message};
357
358    #[test]
359    fn truthy_env_accepts_common_opt_in_values() {
360        assert!(env_truthy("1"));
361        assert!(env_truthy("true"));
362        assert!(env_truthy("YES"));
363        assert!(!env_truthy("0"));
364        assert!(!env_truthy(""));
365    }
366
367    #[test]
368    fn sha256_validation_requires_hex_digest() {
369        assert!(is_sha256_hex(
370            "0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef"
371        ));
372        assert!(!is_sha256_hex("not-a-sha"));
373    }
374
375    #[test]
376    fn empty_explicit_binary_is_unavailable() {
377        let error = PicRuntimeConfig::default()
378            .pocket_ic_bin("")
379            .ensure_binary()
380            .unwrap_err();
381
382        assert!(matches!(
383            error,
384            super::PicStartError::BinaryUnavailable { .. }
385        ));
386    }
387
388    #[test]
389    fn missing_binary_guidance_mentions_opt_in_download() {
390        let message = missing_binary_message(std::path::Path::new("/tmp/missing-pocket-ic"));
391
392        assert!(message.contains("POCKET_IC_BIN"));
393        assert!(message.contains("IC_TESTKIT_ALLOW_POCKET_IC_DOWNLOAD=1"));
394    }
395}