Skip to main content

stash_cli/store/
mod.rs

1use crate::preview::build_preview_data;
2use serde::{Deserialize as SerdeDeserialize, Serialize as SerdeSerialize};
3use serde_json::Value;
4use std::collections::BTreeMap;
5use std::error::Error as StdError;
6use std::fs::{self, File};
7use std::io::{self, Read, Write};
8use std::path::{Path, PathBuf};
9use std::time::{SystemTime, UNIX_EPOCH};
10
11mod attr;
12mod cache;
13mod push;
14
15pub use push::{push_from_reader, tee_from_reader_partial};
16
17pub const SHORT_ID_LEN: usize = 8;
18pub const MIN_ID_LEN: usize = 6;
19
20// ---------------------------------------------------------------------------
21// Public types
22// ---------------------------------------------------------------------------
23
24#[derive(Debug)]
25pub struct PartialSavedError {
26    pub id: String,
27    pub cause: std::io::Error,
28    pub signal: Option<i32>,
29}
30
31impl std::fmt::Display for PartialSavedError {
32    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
33        write!(
34            f,
35            "partial entry saved as \"{}\": {}",
36            self.id,
37            self.cause
38        )
39    }
40}
41
42impl StdError for PartialSavedError {
43    fn source(&self) -> Option<&(dyn StdError + 'static)> {
44        Some(&self.cause)
45    }
46}
47
48pub struct UtcDateTime {
49    pub year: i32,
50    pub month: u32,
51    pub day: u32,
52    pub hour: u32,
53    pub min: u32,
54    pub sec: u32,
55}
56
57#[derive(
58    Clone,
59    Debug,
60    SerdeSerialize,
61    SerdeDeserialize,
62    rkyv::Archive,
63    rkyv::Serialize,
64    rkyv::Deserialize,
65)]
66pub struct Meta {
67    pub id: String,
68    pub ts: String,
69    pub size: i64,
70    pub preview: String,
71    pub attrs: BTreeMap<String, String>,
72}
73
74impl Meta {
75    // IDs are always stored lowercase (created via new_ulid which calls to_ascii_lowercase).
76    pub fn short_id(&self) -> &str {
77        &self.id[self.id.len().saturating_sub(SHORT_ID_LEN)..]
78    }
79
80    pub fn display_id(&self) -> &str {
81        &self.id
82    }
83
84    pub fn to_json_value(&self, include_preview: bool) -> Value {
85        let capacity =
86            3 + self.attrs.len() + usize::from(include_preview && !self.preview.is_empty());
87        let mut map = serde_json::Map::with_capacity(capacity);
88        map.insert("id".into(), Value::String(self.id.clone()));
89        map.insert("ts".into(), Value::String(self.ts.clone()));
90        map.insert("size".into(), Value::Number(self.size.into()));
91        for (k, v) in &self.attrs {
92            map.insert(k.clone(), Value::String(v.clone()));
93        }
94        if include_preview && !self.preview.is_empty() {
95            map.insert("preview".into(), Value::String(self.preview.clone()));
96        }
97        Value::Object(map)
98    }
99}
100
101// ---------------------------------------------------------------------------
102// Entry selection / filtering
103// ---------------------------------------------------------------------------
104
105#[derive(Clone, Debug, Default)]
106pub struct MetaSelection {
107    pub show_all: bool,
108    pub display_tags: Vec<String>,
109    pub filter_tags: Vec<String>,
110}
111
112pub fn parse_meta_selection(values: &[String], show_all: bool) -> io::Result<MetaSelection> {
113    let mut out = MetaSelection {
114        show_all,
115        display_tags: Vec::with_capacity(values.len()),
116        filter_tags: Vec::with_capacity(values.len()),
117    };
118    let mut seen_display = std::collections::HashSet::with_capacity(values.len());
119    let mut seen_filter = std::collections::HashSet::with_capacity(values.len());
120    for value in values {
121        if value.contains(',') || value.contains('=') || value.trim().is_empty() {
122            return Err(io::Error::new(
123                io::ErrorKind::InvalidInput,
124                "--attr accepts name, +name, or ++name and is repeatable",
125            ));
126        }
127        if let Some(key) = value.strip_prefix("++") {
128            if key.is_empty() {
129                return Err(io::Error::new(
130                    io::ErrorKind::InvalidInput,
131                    "--attr filter+display must be ++name",
132                ));
133            }
134            if seen_display.insert(key.to_string()) {
135                out.display_tags.push(key.to_string());
136            }
137            if seen_filter.insert(key.to_string()) {
138                out.filter_tags.push(key.to_string());
139            }
140        } else if let Some(key) = value.strip_prefix('+') {
141            if key.is_empty() {
142                return Err(io::Error::new(
143                    io::ErrorKind::InvalidInput,
144                    "--attr filter must be +name",
145                ));
146            }
147            if seen_filter.insert(key.to_string()) {
148                out.filter_tags.push(key.to_string());
149            }
150        } else if seen_display.insert(value.to_string()) {
151            out.display_tags.push(value.clone());
152        }
153    }
154    Ok(out)
155}
156
157pub fn matches_meta(attrs: &BTreeMap<String, String>, sel: &MetaSelection) -> bool {
158    if sel.filter_tags.is_empty() {
159        return true;
160    }
161    sel.filter_tags.iter().all(|tag| attrs.contains_key(tag))
162}
163
164// ---------------------------------------------------------------------------
165// Directory / path helpers
166// ---------------------------------------------------------------------------
167
168fn cached_base_dir() -> &'static PathBuf {
169    use std::sync::OnceLock;
170    static BASE: OnceLock<PathBuf> = OnceLock::new();
171    BASE.get_or_init(|| match std::env::var("STASH_DIR") {
172        Ok(dir) if !dir.trim().is_empty() => PathBuf::from(dir),
173        _ => {
174            let home = std::env::var("HOME")
175                .map(PathBuf::from)
176                .expect("HOME not set");
177            home.join(".stash")
178        }
179    })
180}
181
182pub fn base_dir() -> io::Result<PathBuf> {
183    Ok(cached_base_dir().clone())
184}
185
186pub fn data_dir() -> io::Result<PathBuf> {
187    Ok(cached_base_dir().join("data"))
188}
189
190pub fn attr_dir() -> io::Result<PathBuf> {
191    Ok(cached_base_dir().join("attr"))
192}
193
194fn cache_dir() -> io::Result<PathBuf> {
195    Ok(cached_base_dir().join("cache"))
196}
197
198fn list_cache_path() -> io::Result<PathBuf> {
199    Ok(cache_dir()?.join("list.cache"))
200}
201
202pub fn entry_dir(id: &str) -> io::Result<PathBuf> {
203    Ok(cached_base_dir().join(id))
204}
205
206pub fn entry_data_path(id: &str) -> io::Result<PathBuf> {
207    Ok(data_dir()?.join(id.to_ascii_lowercase()))
208}
209
210pub fn entry_attr_path(id: &str) -> io::Result<PathBuf> {
211    Ok(attr_dir()?.join(id.to_ascii_lowercase()))
212}
213
214fn tmp_dir() -> io::Result<PathBuf> {
215    Ok(cached_base_dir().join("tmp"))
216}
217
218pub fn init() -> io::Result<()> {
219    let base = cached_base_dir();
220    fs::create_dir_all(base.join("data"))?;
221    fs::create_dir_all(base.join("attr"))?;
222    fs::create_dir_all(base.join("tmp"))?;
223    fs::create_dir_all(base.join("cache"))?;
224    Ok(())
225}
226
227// ---------------------------------------------------------------------------
228// Query API
229// ---------------------------------------------------------------------------
230
231pub fn list_entry_ids() -> io::Result<Vec<String>> {
232    let attrs = attr_dir()?;
233    let read_dir = match fs::read_dir(&attrs) {
234        Ok(rd) => rd,
235        Err(err) if err.kind() == io::ErrorKind::NotFound => return Ok(Vec::new()),
236        Err(err) => return Err(err),
237    };
238    let mut ids: Vec<String> = read_dir
239        .filter_map(|item| item.ok())
240        .map(|item| item.file_name().to_string_lossy().into_owned())
241        .collect();
242    ids.sort_unstable();
243    ids.reverse();
244    Ok(ids)
245}
246
247pub fn list() -> io::Result<Vec<Meta>> {
248    if let Ok(items) = cache::read_list_cache() {
249        return Ok(items);
250    }
251    let entry_ids = list_entry_ids()?;
252    let mut out = Vec::with_capacity(entry_ids.len());
253    for id in entry_ids {
254        if let Ok(meta) = get_meta(&id) {
255            out.push(meta);
256        }
257    }
258    cache::write_list_cache(&out)?;
259    Ok(out)
260}
261
262pub fn all_attr_keys() -> io::Result<Vec<(String, usize)>> {
263    if let Ok(keys) = cache::read_attr_keys() {
264        return Ok(keys);
265    }
266    let items = list()?;
267    // Fall back to rebuilding from the list. write_list_cache already stored
268    // the attr key index, so the next call will hit the cache.
269    cache::write_list_cache(&items)?;
270    cache::read_attr_keys()
271}
272
273pub fn newest() -> io::Result<Meta> {
274    list()?
275        .into_iter()
276        .next()
277        .ok_or_else(|| io::Error::other("stash is empty"))
278}
279
280pub fn nth_newest(n: usize) -> io::Result<Meta> {
281    if n == 0 {
282        return Err(io::Error::new(
283            io::ErrorKind::InvalidInput,
284            "n must be >= 1",
285        ));
286    }
287    let items = list()?;
288    items
289        .into_iter()
290        .nth(n - 1)
291        .ok_or_else(|| io::Error::other("entry index out of range"))
292}
293
294pub fn older_than_ids(id: &str) -> io::Result<Vec<String>> {
295    let items = list()?;
296    for (idx, item) in items.iter().enumerate() {
297        if item.id == id {
298            return Ok(items[idx + 1..].iter().map(|m| m.id.clone()).collect());
299        }
300    }
301    Err(io::Error::new(io::ErrorKind::NotFound, "entry not found"))
302}
303
304pub fn newer_than_ids(id: &str) -> io::Result<Vec<String>> {
305    let items = list()?;
306    for (idx, item) in items.iter().enumerate() {
307        if item.id == id {
308            return Ok(items[..idx].iter().map(|m| m.id.clone()).collect());
309        }
310    }
311    Err(io::Error::new(io::ErrorKind::NotFound, "entry not found"))
312}
313
314pub fn resolve(input: &str) -> io::Result<String> {
315    let raw = input.trim();
316    if raw.is_empty() {
317        return newest().map(|m| m.id);
318    }
319    if let Some(rest) = raw.strip_prefix('@') {
320        let n = rest
321            .parse::<usize>()
322            .map_err(|_| io::Error::new(io::ErrorKind::InvalidInput, "invalid stack ref"))?;
323        return nth_newest(n).map(|m| m.id);
324    }
325    let lower = raw.to_ascii_lowercase();
326    if lower.bytes().all(|c| c.is_ascii_digit()) {
327        let n = lower
328            .parse::<usize>()
329            .map_err(|_| io::Error::new(io::ErrorKind::InvalidInput, "invalid index"))?;
330        return nth_newest(n).map(|m| m.id);
331    }
332    if lower.len() < MIN_ID_LEN {
333        return Err(io::Error::new(io::ErrorKind::InvalidInput, "id too short"));
334    }
335    let ids = list_entry_ids()?;
336    if ids.is_empty() {
337        return Err(io::Error::new(io::ErrorKind::NotFound, "stash is empty"));
338    }
339    if let Some(id) = ids.iter().find(|id| **id == lower) {
340        return Ok(id.clone());
341    }
342    let mut prefix_match: Option<&String> = None;
343    let mut suffix_match: Option<&String> = None;
344    let mut prefix_ambig = false;
345    let mut suffix_ambig = false;
346    for id in &ids {
347        if id.starts_with(&lower) {
348            if prefix_match.is_some() {
349                prefix_ambig = true;
350            } else {
351                prefix_match = Some(id);
352            }
353        }
354        if id.ends_with(&lower) {
355            if suffix_match.is_some() {
356                suffix_ambig = true;
357            } else {
358                suffix_match = Some(id);
359            }
360        }
361    }
362    if let Some(id) = prefix_match {
363        if !prefix_ambig {
364            return Ok(id.clone());
365        }
366        return Err(io::Error::other("ambiguous id"));
367    }
368    if let Some(id) = suffix_match {
369        if !suffix_ambig {
370            return Ok(id.clone());
371        }
372        return Err(io::Error::other("ambiguous id"));
373    }
374    Err(io::Error::new(io::ErrorKind::NotFound, "entry not found"))
375}
376
377pub fn get_meta(id: &str) -> io::Result<Meta> {
378    let path = entry_attr_path(id)?;
379    let data = fs::read_to_string(path)?;
380    attr::parse_attr_file(&data).map_err(io::Error::other)
381}
382
383pub fn write_meta(id: &str, meta: &Meta) -> io::Result<()> {
384    let result = fs::write(entry_attr_path(id)?, attr::encode_attr(meta));
385    if result.is_ok() {
386        cache::invalidate_list_cache();
387    }
388    result
389}
390
391pub fn set_attrs(id: &str, attrs: &BTreeMap<String, String>) -> io::Result<()> {
392    let mut meta = get_meta(id)?;
393    for (k, v) in attrs {
394        meta.attrs.insert(k.clone(), v.clone());
395    }
396    write_meta(id, &meta)
397}
398
399pub fn unset_attrs(id: &str, keys: &[String]) -> io::Result<()> {
400    let mut meta = get_meta(id)?;
401    for key in keys {
402        meta.attrs.remove(key);
403    }
404    write_meta(id, &meta)
405}
406
407pub fn cat_to_writer<W: Write>(id: &str, mut writer: W) -> io::Result<()> {
408    let file = File::open(entry_data_path(id)?)?;
409    let mut reader = io::BufReader::with_capacity(65536, file);
410    io::copy(&mut reader, &mut writer)?;
411    Ok(())
412}
413
414pub fn remove(id: &str) -> io::Result<()> {
415    let data_result = fs::remove_file(entry_data_path(id)?);
416    if let Err(ref e) = data_result {
417        if e.kind() != io::ErrorKind::NotFound {
418            return data_result;
419        }
420    }
421    let attr_result = fs::remove_file(entry_attr_path(id)?);
422    if let Err(ref e) = attr_result {
423        if e.kind() != io::ErrorKind::NotFound {
424            return attr_result;
425        }
426    }
427    cache::invalidate_list_cache();
428    Ok(())
429}
430
431// Called by io::run_read_loop and io::save_or_abort_partial via super::
432fn finalize_saved_entry(
433    id: String,
434    data_path: PathBuf,
435    sample: &[u8],
436    total: i64,
437    attrs: BTreeMap<String, String>,
438) -> io::Result<String> {
439    let meta = Meta {
440        id: id.clone(),
441        ts: now_rfc3339ish()?,
442        size: total,
443        preview: build_preview_data(sample, sample.len()),
444        attrs,
445    };
446    let tmp = tmp_dir()?;
447    let attr_path = tmp.join(format!("{id}.attr"));
448    fs::write(&attr_path, attr::encode_attr(&meta))?;
449    fs::rename(&data_path, entry_data_path(&id)?)?;
450    fs::rename(&attr_path, entry_attr_path(&id)?)?;
451    cache::invalidate_list_cache();
452    Ok(id)
453}
454
455// ---------------------------------------------------------------------------
456// Utilities
457// ---------------------------------------------------------------------------
458
459pub fn human_size(n: i64) -> String {
460    match n {
461        n if n < 1024 => format!("{n}B"),
462        n if n < 1024 * 1024 => format!("{:.1}K", n as f64 / 1024.0),
463        n if n < 1024 * 1024 * 1024 => format!("{:.1}M", n as f64 / (1024.0 * 1024.0)),
464        n => format!("{:.1}G", n as f64 / (1024.0 * 1024.0 * 1024.0)),
465    }
466}
467
468fn now_rfc3339ish() -> io::Result<String> {
469    let now = SystemTime::now()
470        .duration_since(UNIX_EPOCH)
471        .map_err(io::Error::other)?;
472    let secs = now.as_secs() as i64;
473    let nanos = now.subsec_nanos();
474    let dt = unix_to_utc(secs);
475    Ok(format!(
476        "{:04}-{:02}-{:02}T{:02}:{:02}:{:02}.{nanos:09}Z",
477        dt.year, dt.month, dt.day, dt.hour, dt.min, dt.sec,
478    ))
479}
480
481pub fn unix_to_utc(secs: i64) -> UtcDateTime {
482    let days = secs.div_euclid(86_400);
483    let rem = secs.rem_euclid(86_400);
484    let (year, month, day) = civil_from_days(days);
485    UtcDateTime {
486        year,
487        month,
488        day,
489        hour: (rem / 3600) as u32,
490        min: ((rem % 3600) / 60) as u32,
491        sec: (rem % 60) as u32,
492    }
493}
494
495fn civil_from_days(days: i64) -> (i32, u32, u32) {
496    let z = days + 719468;
497    let era = if z >= 0 { z } else { z - 146096 } / 146097;
498    let doe = z - era * 146097;
499    let yoe = (doe - doe / 1460 + doe / 36524 - doe / 146096) / 365;
500    let y = yoe + era * 400;
501    let doy = doe - (365 * yoe + yoe / 4 - yoe / 100);
502    let mp = (5 * doy + 2) / 153;
503    let d = doy - (153 * mp + 2) / 5 + 1;
504    let m = mp + if mp < 10 { 3 } else { -9 };
505    let year = y + if m <= 2 { 1 } else { 0 };
506    (year as i32, m as u32, d as u32)
507}
508
509fn new_ulid() -> io::Result<String> {
510    let now = SystemTime::now()
511        .duration_since(UNIX_EPOCH)
512        .map_err(io::Error::other)?
513        .as_millis() as u64;
514    let mut bytes = [0u8; 16];
515    for (i, byte) in bytes.iter_mut().enumerate().take(6) {
516        *byte = ((now >> (8 * (5 - i))) & 0xff) as u8;
517    }
518    let mut rand = File::open("/dev/urandom")?;
519    rand.read_exact(&mut bytes[6..])?;
520    Ok(encode_ulid(bytes).to_ascii_lowercase())
521}
522
523fn encode_ulid(bytes: [u8; 16]) -> String {
524    const ALPHABET: &[u8; 32] = b"0123456789ABCDEFGHJKMNPQRSTVWXYZ";
525    let mut value = 0u128;
526    for byte in bytes {
527        value = (value << 8) | byte as u128;
528    }
529    let mut out = [0u8; 26];
530    for i in (0..26).rev() {
531        out[i] = ALPHABET[(value & 0x1f) as usize];
532        value >>= 5;
533    }
534    // SAFETY: out contains only bytes from ALPHABET, which is ASCII-only
535    unsafe { String::from_utf8_unchecked(out.to_vec()) }
536}
537
538pub fn add_filename_attr(path: &Path, attrs: &mut BTreeMap<String, String>) {
539    if let Some(name) = path.file_name().and_then(|s| s.to_str()) {
540        attrs.insert("filename".into(), name.into());
541    }
542}