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