nix_dev_env/
nix_profile_cache.rs

1use std::{
2    ffi::OsStr,
3    fs,
4    io::{self, Write},
5    path::{Path, PathBuf},
6    process,
7    time::SystemTime,
8};
9
10use serde_json::Value;
11use sha1::{Digest, Sha1};
12
13use crate::nix_command;
14
15#[derive(Debug, Clone)]
16pub struct NixProfileCache {
17    cache_dir: PathBuf,
18    flake_inputs_dir: PathBuf,
19    flake_reference: FlakeReference,
20    evaluation_mode: EvaluationMode,
21    files_to_watch: Vec<PathBuf>,
22    profile_symlink: PathBuf,
23    profile_rc_file: PathBuf,
24}
25
26impl NixProfileCache {
27    pub fn new(
28        cache_dir: PathBuf,
29        flake_reference: &str,
30        evaluation_mode: EvaluationMode,
31    ) -> anyhow::Result<Self> {
32        let flake_inputs_dir = cache_dir.join("flake-inputs");
33
34        let flake_reference = FlakeReference::parse(flake_reference)?;
35
36        let mut files_to_watch = vec![];
37        let hash = if let Some(flake_dir) = &flake_reference.flake_dir {
38            files_to_watch.extend_from_slice(&[
39                flake_dir.join("flake.nix"),
40                flake_dir.join("flake.lock"),
41                flake_dir.join("devshell.toml"),
42            ]);
43            hash_files(&files_to_watch)?
44        } else {
45            hash_flake_reference(&flake_reference.flake_reference_string)?
46        };
47
48        let profile_symlink = cache_dir.join(format!("flake-profile-{}", hash));
49        let profile_rc_file = profile_symlink.with_extension("rc");
50        Ok(Self {
51            cache_dir,
52            flake_inputs_dir,
53            flake_reference,
54            evaluation_mode,
55            files_to_watch,
56            profile_symlink,
57            profile_rc_file,
58        })
59    }
60
61    pub fn needs_update(&self) -> anyhow::Result<bool> {
62        let mut need_update = true;
63
64        if self.profile_rc_file.is_file() && self.profile_symlink.is_symlink() {
65            let profile_rc_mtime = fs::metadata(&self.profile_rc_file)?.modified()?;
66
67            need_update = self.files_to_watch.iter().any(|file| {
68                fs::metadata(file)
69                    .and_then(|meta| meta.modified())
70                    .unwrap_or(SystemTime::UNIX_EPOCH)
71                    > profile_rc_mtime
72            });
73        }
74
75        Ok(need_update)
76    }
77
78    pub fn update(&self) -> anyhow::Result<()> {
79        clean_old_gcroots(&self.cache_dir, &self.flake_inputs_dir)?;
80
81        let tmp_profile = self
82            .cache_dir
83            .join(format!("flake-tmp-profile.{}", process::id()));
84
85        let stdout_content = nix_command::nix([
86            OsStr::new("print-dev-env"),
87            OsStr::new(self.impure_arg()),
88            OsStr::new("--no-write-lock-file"),
89            OsStr::new("--profile"),
90            tmp_profile.as_os_str(),
91            OsStr::new(&self.flake_reference.flake_reference_string),
92        ])?;
93
94        fs::File::create(&self.profile_rc_file)?.write_all(stdout_content.as_bytes())?;
95
96        self.add_gcroot(&tmp_profile, &self.profile_symlink)?;
97        fs::remove_file(&tmp_profile)?;
98
99        if self.flake_reference.flake_dir.is_some() {
100            for input in self.get_flake_input_paths()? {
101                let store_path = PathBuf::from("/nix/store").join(&input);
102                let symlink_path = self.flake_inputs_dir.join(&input);
103                self.add_gcroot(&store_path, &symlink_path)?;
104            }
105        }
106
107        Ok(())
108    }
109
110    pub fn profile_rc(&self) -> &Path {
111        &self.profile_rc_file
112    }
113
114    fn impure_arg(&self) -> &str {
115        match self.evaluation_mode {
116            EvaluationMode::Impure => "--impure",
117            EvaluationMode::Pure => "",
118        }
119    }
120
121    fn add_gcroot(&self, store_path: &Path, symlink: &Path) -> anyhow::Result<()> {
122        nix_command::nix([
123            OsStr::new("build"),
124            OsStr::new(&self.impure_arg()),
125            OsStr::new("--out-link"),
126            symlink.as_os_str(),
127            store_path.as_os_str(),
128        ])?;
129        Ok(())
130    }
131
132    fn get_flake_input_paths(&self) -> anyhow::Result<Vec<PathBuf>> {
133        let stdout_content = nix_command::nix([
134            "flake",
135            "archive",
136            self.impure_arg(),
137            "--json",
138            "--no-write-lock-file",
139            &self.flake_reference.flake_reference_string,
140        ])?;
141        let json = serde_json::from_str::<Value>(&stdout_content)?;
142        Ok(get_paths_from_doc(&json))
143    }
144}
145
146#[derive(Debug, Clone, Copy)]
147pub enum EvaluationMode {
148    Impure,
149    Pure,
150}
151
152#[derive(Debug, Clone)]
153struct FlakeReference {
154    pub flake_reference_string: String,
155    pub flake_dir: Option<PathBuf>,
156}
157
158impl FlakeReference {
159    pub fn parse(flake_reference: &str) -> anyhow::Result<Self> {
160        let mut flake_reference_iter = flake_reference.split('#');
161        let flake_uri = flake_reference_iter
162            .next()
163            .ok_or_else(|| anyhow::anyhow!("Missing flake URI"))?;
164        let flake_specifier = flake_reference_iter.next();
165
166        let expanded_flake_reference_and_flake_dir =
167            if FlakeReference::is_path_type(flake_reference) {
168                let flake_dir_str =
169                    shellexpand::full(flake_uri.strip_prefix("path:").unwrap_or(flake_uri))?;
170                let expanded_flake_reference = format!(
171                    "{}{}{}",
172                    &flake_dir_str,
173                    flake_specifier.map(|_| "#").unwrap_or(""),
174                    flake_specifier.unwrap_or("")
175                );
176                Some((expanded_flake_reference, flake_dir_str))
177            } else {
178                None
179            };
180
181        Ok(Self {
182            flake_dir: expanded_flake_reference_and_flake_dir
183                .as_ref()
184                .map(|x| PathBuf::from(x.1.to_string())),
185            flake_reference_string: expanded_flake_reference_and_flake_dir
186                .map(|x| x.0)
187                .unwrap_or_else(|| String::from(flake_reference)),
188        })
189    }
190
191    fn is_path_type(flake_reference: &str) -> bool {
192        flake_reference.starts_with("path:")
193            || flake_reference.starts_with('~')
194            || flake_reference.starts_with('/')
195            || flake_reference.starts_with("./")
196            || flake_reference.starts_with("../")
197    }
198}
199
200fn hash_files(filenames: impl AsRef<[PathBuf]>) -> anyhow::Result<String> {
201    let (hasher, no_files) = filenames
202        .as_ref()
203        .iter()
204        .filter(|f| {
205            // TODO: figure out what to do if the file doesn't exist
206            f.exists()
207        })
208        .try_fold((Sha1::new(), true), |(mut acc, ..), f| {
209            acc.update(fs::read(f)?);
210            anyhow::Result::<(Sha1, bool)>::Ok((acc, false))
211        })?;
212
213    if no_files {
214        return Err(anyhow::anyhow!("No files found to hash"));
215    }
216
217    Ok(format!("{:x}", hasher.finalize()))
218}
219
220fn hash_flake_reference(flake_reference: &str) -> anyhow::Result<String> {
221    let mut hasher = Sha1::new();
222    hasher.update(flake_reference);
223    Ok(format!("{:x}", hasher.finalize()))
224}
225
226fn clean_old_gcroots(cache_dir: &Path, flake_inputs_dir: &Path) -> anyhow::Result<()> {
227    let res = fs::remove_dir_all(cache_dir);
228    if let Err(e) = &res
229        && e.kind() != io::ErrorKind::NotFound
230    {
231        res?;
232    }
233    fs::create_dir_all(flake_inputs_dir)?;
234    Ok(())
235}
236
237fn get_paths_from_doc(doc: &Value) -> Vec<PathBuf> {
238    let mut result = Vec::new();
239
240    if let Some(p) = get_path(doc) {
241        result.push(p);
242    }
243
244    if let Some(inputs) = doc.get("inputs").and_then(|i| i.as_object()) {
245        for (_k, v) in inputs {
246            let sub_paths = get_paths_from_doc(v);
247            result.extend(sub_paths);
248        }
249    }
250
251    result
252}
253
254fn get_path(doc: &Value) -> Option<PathBuf> {
255    doc.get("path")
256        .and_then(|value| value.as_str())
257        .map(|path| {
258            if path.len() > 11 {
259                PathBuf::from(&path[11..])
260            } else {
261                PathBuf::from(path)
262            }
263        })
264}
265
266#[cfg(test)]
267mod tests {
268    use std::{io::Write, path::PathBuf};
269
270    use once_cell::sync::Lazy;
271    use serde_json::json;
272    use tempfile::NamedTempFile;
273
274    use super::{get_path, get_paths_from_doc, hash_files};
275
276    static TEST_FILE: Lazy<NamedTempFile> = Lazy::new(|| {
277        let mut test_file = tempfile::NamedTempFile::new().unwrap();
278        writeln!(test_file.as_file_mut(), r#"echo "1.1.1";"#).unwrap();
279        test_file
280    });
281
282    #[test]
283    fn test_hash_one() {
284        assert_eq!(
285            hash_files([TEST_FILE.path().to_path_buf()]).unwrap(),
286            "6ead949bf4bcae230b9ed9cd11e578e34ce9f9ea"
287        );
288    }
289
290    #[test]
291    fn test_hash_multiple() {
292        assert_eq!(
293            hash_files([
294                TEST_FILE.path().to_path_buf(),
295                TEST_FILE.path().to_path_buf(),
296            ])
297            .unwrap(),
298            "f109b7892a541ed1e3cf39314cd25d21042b984f"
299        );
300    }
301
302    #[test]
303    fn test_hash_filters_nonexistent() {
304        assert_eq!(
305            hash_files([TEST_FILE.path().to_path_buf(), PathBuf::from("FOOBARBAZ"),]).unwrap(),
306            "6ead949bf4bcae230b9ed9cd11e578e34ce9f9ea"
307        );
308    }
309
310    #[test]
311    fn test_get_path_removes_prefix() {
312        let input = json!({
313            "path": "aaaaaaaaaaabbbbb"
314        });
315        let result = get_path(&input);
316        assert_eq!(result, Some(PathBuf::from("bbbbb")));
317    }
318
319    #[test]
320    fn test_get_paths_from_doc() {
321        let input = json!({
322            "path": "aaaaaaaaaaabbbbb",
323            "inputs": {
324                "foo": {
325                    "path": "aaaaaaaaaaaccccc",
326                    "inputs": {
327                        "bar": {
328                            "path": "aaaaaaaaaaaddddd",
329                            "inputs": {}
330                        }
331                    }
332                }
333            }
334        });
335        let result = get_paths_from_doc(&input);
336        assert_eq!(
337            result,
338            vec![
339                "bbbbb".to_string(),
340                "ccccc".to_string(),
341                "ddddd".to_string()
342            ]
343        );
344    }
345}
nix_dev_env/nix_profile_cache.rs

nix_dev_env/
nix_profile_cache.rs
