tg_utils/
preauth.rs

1use thiserror::Error;
2
3use cosmwasm_std::{StdError, Storage};
4use cw_storage_plus::Item;
5
6#[derive(Error, Debug, PartialEq)]
7pub enum PreauthError {
8    #[error("{0}")]
9    Std(#[from] StdError),
10
11    #[error("No preauthorization available to add hook")]
12    NoPreauth {},
13}
14
15// store all hook addresses in one item. We cannot have many of them before the contract becomes unusable anyway.
16pub struct Preauth<'a>(Item<'a, u64>);
17
18impl<'a> Preauth<'a> {
19    pub const fn new(preauth_key: &'a str) -> Self {
20        Preauth(Item::new(preauth_key))
21    }
22
23    pub fn set_auth(&self, storage: &mut dyn Storage, count: u64) -> Result<(), StdError> {
24        self.0.save(storage, &count)
25    }
26
27    pub fn get_auth(&self, storage: &dyn Storage) -> Result<u64, StdError> {
28        Ok(self.0.may_load(storage)?.unwrap_or_default())
29    }
30
31    pub fn add_auth(&self, storage: &mut dyn Storage) -> Result<(), StdError> {
32        let count = self.get_auth(storage)?;
33        self.set_auth(storage, count + 1)
34    }
35
36    pub fn use_auth(&self, storage: &mut dyn Storage) -> Result<(), PreauthError> {
37        let count = self.get_auth(storage)?;
38        let count = count.checked_sub(1).ok_or(PreauthError::NoPreauth {})?;
39        Ok(self.set_auth(storage, count)?)
40    }
41}
42
43#[cfg(test)]
44mod test {
45    use super::*;
46    use cosmwasm_std::testing::MockStorage;
47
48    const PREAUTH: Preauth = Preauth::new("preauth");
49
50    #[test]
51    fn count_preauth() {
52        let mut storage = MockStorage::new();
53
54        // start and cannot consume
55        assert_eq!(PREAUTH.get_auth(&storage).unwrap(), 0);
56        let err = PREAUTH.use_auth(&mut storage).unwrap_err();
57        assert_eq!(err, PreauthError::NoPreauth {});
58
59        // add one and use it (only once)
60        PREAUTH.add_auth(&mut storage).unwrap();
61        assert_eq!(PREAUTH.get_auth(&storage).unwrap(), 1);
62        PREAUTH.use_auth(&mut storage).unwrap();
63        assert_eq!(PREAUTH.get_auth(&storage).unwrap(), 0);
64        let err = PREAUTH.use_auth(&mut storage).unwrap_err();
65        assert_eq!(err, PreauthError::NoPreauth {});
66
67        // set to a higher value
68        PREAUTH.set_auth(&mut storage, 27).unwrap();
69        assert_eq!(PREAUTH.get_auth(&storage).unwrap(), 27);
70        PREAUTH.use_auth(&mut storage).unwrap();
71        assert_eq!(PREAUTH.get_auth(&storage).unwrap(), 26);
72    }
73}
tg_utils/preauth.rs

tg_utils/
preauth.rs
