1use std::collections::BTreeMap;
10use std::io::Write;
11use std::path::PathBuf;
12
13pub const MAX_KEY: usize = 128;
14pub const MAX_VALUE: usize = 1 << 20;
15pub const QUOTA: usize = 16 << 20;
16
17#[derive(Debug)]
19pub enum StorageError {
20 Key(String),
21 Size(usize),
22 Quota,
23 Io(String),
24}
25
26impl StorageError {
27 pub fn code(&self) -> i32 {
28 match self {
29 StorageError::Key(_) => -1,
30 StorageError::Size(_) => -2,
31 StorageError::Quota => -3,
32 StorageError::Io(_) => -4,
33 }
34 }
35}
36
37impl std::fmt::Display for StorageError {
38 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
39 match self {
40 StorageError::Key(k) => write!(f, "invalid key {k:?}"),
41 StorageError::Size(n) => write!(f, "value is {n} bytes (max {MAX_VALUE})"),
42 StorageError::Quota => write!(f, "storage quota of {QUOTA} bytes exceeded"),
43 StorageError::Io(e) => write!(f, "{e}"),
44 }
45 }
46}
47
48pub struct Storage {
49 dir: Option<PathBuf>,
50 values: BTreeMap<String, Vec<u8>>,
52 used: usize,
54 keys: Option<Vec<String>>,
56}
57
58pub fn valid_key(k: &str) -> bool {
59 !k.is_empty()
60 && k.len() <= MAX_KEY
61 && k.bytes().all(|b| b.is_ascii_alphanumeric() || b == b'.' || b == b'_' || b == b'-')
62 && k != "."
63 && k != ".."
64}
65
66pub fn default_dir(game_id: &str) -> Option<PathBuf> {
69 data_root().map(|r| r.join(game_id))
70}
71
72pub fn data_root() -> Option<PathBuf> {
74 let base = if cfg!(target_os = "macos") {
75 std::env::var_os("HOME").map(|h| PathBuf::from(h).join("Library/Application Support"))
76 } else if cfg!(windows) {
77 std::env::var_os("APPDATA").map(PathBuf::from)
78 } else {
79 std::env::var_os("XDG_DATA_HOME")
80 .map(PathBuf::from)
81 .or_else(|| std::env::var_os("HOME").map(|h| PathBuf::from(h).join(".local/share")))
82 };
83 base.map(|b| b.join("gasm"))
84}
85
86fn temp_name(key: &str) -> String {
87 format!(".{key}~tmp")
88}
89
90impl Storage {
91 pub fn memory() -> Storage {
92 Storage { dir: None, values: BTreeMap::new(), used: 0, keys: None }
93 }
94
95 pub fn open(dir: PathBuf) -> Result<Storage, String> {
98 std::fs::create_dir_all(&dir).map_err(|e| format!("{}: {e}", dir.display()))?;
99 let mut s = Storage { dir: Some(dir.clone()), ..Storage::memory() };
100 for entry in std::fs::read_dir(&dir).map_err(|e| e.to_string())?.flatten() {
101 let name = entry.file_name().to_string_lossy().into_owned();
102 if name.ends_with("~tmp") {
103 let _ = std::fs::remove_file(entry.path());
104 } else if valid_key(&name) && entry.path().is_file() {
105 if let Ok(v) = std::fs::read(entry.path()) {
106 s.used += name.len() + v.len();
107 s.values.insert(name, v);
108 }
109 }
110 }
111 Ok(s)
112 }
113
114 pub fn location(&self) -> String {
115 self.dir.as_ref().map_or("memory".into(), |d| d.display().to_string())
116 }
117
118 pub fn get(&self, key: &str) -> Option<&[u8]> {
119 self.values.get(key).map(Vec::as_slice)
120 }
121
122 pub fn set(&mut self, key: &str, value: &[u8]) -> Result<(), StorageError> {
123 if !valid_key(key) {
124 return Err(StorageError::Key(key.to_owned()));
125 }
126 if value.len() > MAX_VALUE {
127 return Err(StorageError::Size(value.len()));
128 }
129 let old = self.values.get(key).map_or(0, |v| key.len() + v.len());
130 let used = self.used - old + key.len() + value.len();
131 if used > QUOTA {
132 return Err(StorageError::Quota);
133 }
134 if let Some(dir) = &self.dir {
135 let tmp = dir.join(temp_name(key));
136 let write = || -> std::io::Result<()> {
137 let mut f = std::fs::File::create(&tmp)?;
138 f.write_all(value)?;
139 f.sync_all()?;
140 std::fs::rename(&tmp, dir.join(key))
141 };
142 if let Err(e) = write() {
143 let _ = std::fs::remove_file(&tmp);
144 return Err(StorageError::Io(e.to_string()));
145 }
146 }
147 if !self.values.contains_key(key) {
148 self.keys = None;
149 }
150 self.values.insert(key.to_owned(), value.to_vec());
151 self.used = used;
152 Ok(())
153 }
154
155 pub fn keys(&mut self) -> &[String] {
157 self.keys.get_or_insert_with(|| self.values.keys().cloned().collect())
158 }
159
160 pub fn delete(&mut self, key: &str) -> bool {
161 let Some(v) = self.values.remove(key) else { return false };
162 self.used -= key.len() + v.len();
163 self.keys = None;
164 if let Some(dir) = &self.dir {
165 let _ = std::fs::remove_file(dir.join(key));
166 }
167 true
168 }
169}
170
171#[cfg(test)]
172mod tests {
173 use super::*;
174
175 #[test]
176 fn temp_files_never_collide_with_keys() {
177 let d = std::env::temp_dir().join(format!("gasm-storage-test-{}", std::process::id()));
178 let _ = std::fs::remove_dir_all(&d);
179 let mut s = Storage::open(d.clone()).unwrap();
180 s.set(".foo.tmp", b"guest's own key").unwrap();
181 s.set("foo", b"value").unwrap();
182 assert_eq!(s.get(".foo.tmp"), Some(&b"guest's own key"[..]));
183 std::fs::write(d.join(temp_name("bar")), b"left by a crash").unwrap();
184 let mut s = Storage::open(d.clone()).unwrap();
185 assert_eq!(s.keys(), [".foo.tmp", "foo"]);
186 assert!(!d.join(temp_name("bar")).exists());
187 let _ = std::fs::remove_dir_all(&d);
188 }
189
190 #[test]
191 fn error_codes_and_quota_accounting() {
192 let mut s = Storage::memory();
193 assert_eq!(s.set("bad/key", b"x").unwrap_err().code(), -1);
194 assert_eq!(s.set("big", &vec![0; MAX_VALUE + 1]).unwrap_err().code(), -2);
195 for i in 0..15 {
196 s.set(&format!("k{i:02}"), &vec![0; MAX_VALUE]).unwrap();
197 }
198 assert_eq!(s.set("one-more", &vec![0; MAX_VALUE]).unwrap_err().code(), -3);
199 s.set("k00", b"small").unwrap(); s.set("one-more", &vec![0; MAX_VALUE]).unwrap();
201 assert!(s.delete("k01") && !s.delete("k01"));
202 assert_eq!(s.keys().len(), 15);
203 }
204}