atuin_dotfiles/
store.rs

1use std::collections::BTreeMap;
2
3use atuin_client::record::sqlite_store::SqliteStore;
4// Sync aliases
5// This will be noticeable similar to the kv store, though I expect the two shall diverge
6// While we will support a range of shell config, I'd rather have a larger number of small records
7// + stores, rather than one mega config store.
8use atuin_common::record::{DecryptedData, Host, HostId};
9use atuin_common::utils::unquote;
10use eyre::{bail, ensure, eyre, Result};
11
12use atuin_client::record::encryption::PASETO_V4;
13use atuin_client::record::store::Store;
14
15use crate::shell::Alias;
16
17const CONFIG_SHELL_ALIAS_VERSION: &str = "v0";
18const CONFIG_SHELL_ALIAS_TAG: &str = "config-shell-alias";
19const CONFIG_SHELL_ALIAS_FIELD_MAX_LEN: usize = 20000; // 20kb max total len, way more than should be needed.
20
21mod alias;
22pub mod var;
23
24#[derive(Debug, Clone, PartialEq, Eq)]
25pub enum AliasRecord {
26    Create(Alias),  // create a full record
27    Delete(String), // delete by name
28}
29
30impl AliasRecord {
31    pub fn serialize(&self) -> Result<DecryptedData> {
32        use rmp::encode;
33
34        let mut output = vec![];
35
36        match self {
37            AliasRecord::Create(alias) => {
38                encode::write_u8(&mut output, 0)?; // create
39                encode::write_array_len(&mut output, 2)?; // 2 fields
40
41                encode::write_str(&mut output, alias.name.as_str())?;
42                encode::write_str(&mut output, alias.value.as_str())?;
43            }
44            AliasRecord::Delete(name) => {
45                encode::write_u8(&mut output, 1)?; // delete
46                encode::write_array_len(&mut output, 1)?; // 1 field
47
48                encode::write_str(&mut output, name.as_str())?;
49            }
50        }
51
52        Ok(DecryptedData(output))
53    }
54
55    pub fn deserialize(data: &DecryptedData, version: &str) -> Result<Self> {
56        use rmp::decode;
57
58        fn error_report<E: std::fmt::Debug>(err: E) -> eyre::Report {
59            eyre!("{err:?}")
60        }
61
62        match version {
63            CONFIG_SHELL_ALIAS_VERSION => {
64                let mut bytes = decode::Bytes::new(&data.0);
65
66                let record_type = decode::read_u8(&mut bytes).map_err(error_report)?;
67
68                match record_type {
69                    // create
70                    0 => {
71                        let nfields = decode::read_array_len(&mut bytes).map_err(error_report)?;
72                        ensure!(
73                            nfields == 2,
74                            "too many entries in v0 shell alias create record"
75                        );
76
77                        let bytes = bytes.remaining_slice();
78
79                        let (key, bytes) =
80                            decode::read_str_from_slice(bytes).map_err(error_report)?;
81                        let (value, bytes) =
82                            decode::read_str_from_slice(bytes).map_err(error_report)?;
83
84                        if !bytes.is_empty() {
85                            bail!("trailing bytes in encoded shell alias record. malformed")
86                        }
87
88                        Ok(AliasRecord::Create(Alias {
89                            name: key.to_owned(),
90                            value: value.to_owned(),
91                        }))
92                    }
93
94                    // delete
95                    1 => {
96                        let nfields = decode::read_array_len(&mut bytes).map_err(error_report)?;
97                        ensure!(
98                            nfields == 1,
99                            "too many entries in v0 shell alias delete record"
100                        );
101
102                        let bytes = bytes.remaining_slice();
103
104                        let (key, bytes) =
105                            decode::read_str_from_slice(bytes).map_err(error_report)?;
106
107                        if !bytes.is_empty() {
108                            bail!("trailing bytes in encoded shell alias record. malformed")
109                        }
110
111                        Ok(AliasRecord::Delete(key.to_owned()))
112                    }
113
114                    n => {
115                        bail!("unknown AliasRecord type {n}")
116                    }
117                }
118            }
119            _ => {
120                bail!("unknown version {version:?}")
121            }
122        }
123    }
124}
125
126#[derive(Debug, Clone)]
127pub struct AliasStore {
128    pub store: SqliteStore,
129    pub host_id: HostId,
130    pub encryption_key: [u8; 32],
131}
132
133impl AliasStore {
134    // will want to init the actual kv store when that is done
135    pub fn new(store: SqliteStore, host_id: HostId, encryption_key: [u8; 32]) -> AliasStore {
136        AliasStore {
137            store,
138            host_id,
139            encryption_key,
140        }
141    }
142
143    pub async fn posix(&self) -> Result<String> {
144        let aliases = self.aliases().await?;
145
146        let mut config = String::new();
147
148        for alias in aliases {
149            // If it's quoted, remove the quotes. If it's not quoted, do nothing.
150            let value = unquote(alias.value.as_str()).unwrap_or(alias.value.clone());
151
152            // we're about to quote it ourselves anyway!
153            config.push_str(&format!("alias {}='{}'\n", alias.name, value));
154        }
155
156        Ok(config)
157    }
158
159    pub async fn xonsh(&self) -> Result<String> {
160        let aliases = self.aliases().await?;
161
162        let mut config = String::new();
163
164        for alias in aliases {
165            config.push_str(&format!("aliases['{}'] ='{}'\n", alias.name, alias.value));
166        }
167
168        Ok(config)
169    }
170
171    pub async fn build(&self) -> Result<()> {
172        let dir = atuin_common::utils::dotfiles_cache_dir();
173        tokio::fs::create_dir_all(dir.clone()).await?;
174
175        // Build for all supported shells
176        let posix = self.posix().await?;
177        let xonsh = self.xonsh().await?;
178
179        // All the same contents, maybe optimize in the future or perhaps there will be quirks
180        // per-shell
181        // I'd prefer separation atm
182        let zsh = dir.join("aliases.zsh");
183        let bash = dir.join("aliases.bash");
184        let fish = dir.join("aliases.fish");
185        let xsh = dir.join("aliases.xsh");
186
187        tokio::fs::write(zsh, &posix).await?;
188        tokio::fs::write(bash, &posix).await?;
189        tokio::fs::write(fish, &posix).await?;
190        tokio::fs::write(xsh, &xonsh).await?;
191
192        Ok(())
193    }
194
195    pub async fn set(&self, name: &str, value: &str) -> Result<()> {
196        if name.len() + value.len() > CONFIG_SHELL_ALIAS_FIELD_MAX_LEN {
197            return Err(eyre!(
198                "alias record too large: max len {} bytes",
199                CONFIG_SHELL_ALIAS_FIELD_MAX_LEN
200            ));
201        }
202
203        let record = AliasRecord::Create(Alias {
204            name: name.to_string(),
205            value: value.to_string(),
206        });
207
208        let bytes = record.serialize()?;
209
210        let idx = self
211            .store
212            .last(self.host_id, CONFIG_SHELL_ALIAS_TAG)
213            .await?
214            .map_or(0, |entry| entry.idx + 1);
215
216        let record = atuin_common::record::Record::builder()
217            .host(Host::new(self.host_id))
218            .version(CONFIG_SHELL_ALIAS_VERSION.to_string())
219            .tag(CONFIG_SHELL_ALIAS_TAG.to_string())
220            .idx(idx)
221            .data(bytes)
222            .build();
223
224        self.store
225            .push(&record.encrypt::<PASETO_V4>(&self.encryption_key))
226            .await?;
227
228        // set mutates shell config, so build again
229        self.build().await?;
230
231        Ok(())
232    }
233
234    pub async fn delete(&self, name: &str) -> Result<()> {
235        if name.len() > CONFIG_SHELL_ALIAS_FIELD_MAX_LEN {
236            return Err(eyre!(
237                "alias record too large: max len {} bytes",
238                CONFIG_SHELL_ALIAS_FIELD_MAX_LEN
239            ));
240        }
241
242        let record = AliasRecord::Delete(name.to_string());
243
244        let bytes = record.serialize()?;
245
246        let idx = self
247            .store
248            .last(self.host_id, CONFIG_SHELL_ALIAS_TAG)
249            .await?
250            .map_or(0, |entry| entry.idx + 1);
251
252        let record = atuin_common::record::Record::builder()
253            .host(Host::new(self.host_id))
254            .version(CONFIG_SHELL_ALIAS_VERSION.to_string())
255            .tag(CONFIG_SHELL_ALIAS_TAG.to_string())
256            .idx(idx)
257            .data(bytes)
258            .build();
259
260        self.store
261            .push(&record.encrypt::<PASETO_V4>(&self.encryption_key))
262            .await?;
263
264        // delete mutates shell config, so build again
265        self.build().await?;
266
267        Ok(())
268    }
269
270    pub async fn aliases(&self) -> Result<Vec<Alias>> {
271        let mut build = BTreeMap::new();
272
273        // this is sorted, oldest to newest
274        let tagged = self.store.all_tagged(CONFIG_SHELL_ALIAS_TAG).await?;
275
276        for record in tagged {
277            let version = record.version.clone();
278
279            let decrypted = match version.as_str() {
280                CONFIG_SHELL_ALIAS_VERSION => record.decrypt::<PASETO_V4>(&self.encryption_key)?,
281                version => bail!("unknown version {version:?}"),
282            };
283
284            let ar = AliasRecord::deserialize(&decrypted.data, version.as_str())?;
285
286            match ar {
287                AliasRecord::Create(a) => {
288                    build.insert(a.name.clone(), a);
289                }
290                AliasRecord::Delete(d) => {
291                    build.remove(&d);
292                }
293            }
294        }
295
296        Ok(build.into_values().collect())
297    }
298}
299
300#[cfg(test)]
301pub(crate) fn test_local_timeout() -> f64 {
302    std::env::var("ATUIN_TEST_LOCAL_TIMEOUT")
303        .ok()
304        .and_then(|x| x.parse().ok())
305        // this hardcoded value should be replaced by a simple way to get the
306        // default local_timeout of Settings if possible
307        .unwrap_or(2.0)
308}
309
310#[cfg(test)]
311mod tests {
312    use rand::rngs::OsRng;
313
314    use atuin_client::record::sqlite_store::SqliteStore;
315
316    use crate::shell::Alias;
317
318    use super::{test_local_timeout, AliasRecord, AliasStore, CONFIG_SHELL_ALIAS_VERSION};
319    use crypto_secretbox::{KeyInit, XSalsa20Poly1305};
320
321    #[test]
322    fn encode_decode() {
323        let record = Alias {
324            name: "k".to_owned(),
325            value: "kubectl".to_owned(),
326        };
327        let record = AliasRecord::Create(record);
328
329        let snapshot = [204, 0, 146, 161, 107, 167, 107, 117, 98, 101, 99, 116, 108];
330
331        let encoded = record.serialize().unwrap();
332        let decoded = AliasRecord::deserialize(&encoded, CONFIG_SHELL_ALIAS_VERSION).unwrap();
333
334        assert_eq!(encoded.0, &snapshot);
335        assert_eq!(decoded, record);
336    }
337
338    #[tokio::test]
339    async fn build_aliases() {
340        let store = SqliteStore::new(":memory:", test_local_timeout())
341            .await
342            .unwrap();
343        let key: [u8; 32] = XSalsa20Poly1305::generate_key(&mut OsRng).into();
344        let host_id = atuin_common::record::HostId(atuin_common::utils::uuid_v7());
345
346        let alias = AliasStore::new(store, host_id, key);
347
348        alias.set("k", "kubectl").await.unwrap();
349        alias.set("gp", "git push").await.unwrap();
350        alias
351            .set("kgap", "'kubectl get pods --all-namespaces'")
352            .await
353            .unwrap();
354
355        let mut aliases = alias.aliases().await.unwrap();
356
357        aliases.sort_by_key(|a| a.name.clone());
358
359        assert_eq!(aliases.len(), 3);
360
361        assert_eq!(
362            aliases[0],
363            Alias {
364                name: String::from("gp"),
365                value: String::from("git push")
366            }
367        );
368
369        assert_eq!(
370            aliases[1],
371            Alias {
372                name: String::from("k"),
373                value: String::from("kubectl")
374            }
375        );
376
377        assert_eq!(
378            aliases[2],
379            Alias {
380                name: String::from("kgap"),
381                value: String::from("'kubectl get pods --all-namespaces'")
382            }
383        );
384
385        let build = alias.posix().await.expect("failed to build aliases");
386
387        assert_eq!(
388            build,
389            "alias gp='git push'
390alias k='kubectl'
391alias kgap='kubectl get pods --all-namespaces'
392"
393        )
394    }
395}
atuin_dotfiles/store.rs

atuin_dotfiles/
store.rs
