cw_hooks/
lib.rs

1#![doc = include_str!(concat!(env!("CARGO_MANIFEST_DIR"), "/README.md"))]
2
3use thiserror::Error;
4
5use cosmwasm_schema::cw_serde;
6use cosmwasm_std::{Addr, CustomQuery, Deps, StdError, StdResult, Storage, SubMsg};
7use cw_storage_plus::Item;
8
9#[cw_serde]
10pub struct HooksResponse {
11    pub hooks: Vec<String>,
12}
13
14#[derive(Error, Debug, PartialEq)]
15pub enum HookError {
16    #[error("{0}")]
17    Std(#[from] StdError),
18
19    #[error("Given address already registered as a hook")]
20    HookAlreadyRegistered {},
21
22    #[error("Given address not registered as a hook")]
23    HookNotRegistered {},
24}
25
26// store all hook addresses in one item. We cannot have many of them before the contract becomes unusable anyway.
27pub struct Hooks<'a>(Item<'a, Vec<Addr>>);
28
29impl<'a> Hooks<'a> {
30    pub const fn new(storage_key: &'a str) -> Self {
31        Hooks(Item::new(storage_key))
32    }
33
34    pub fn add_hook(&self, storage: &mut dyn Storage, addr: Addr) -> Result<(), HookError> {
35        let mut hooks = self.0.may_load(storage)?.unwrap_or_default();
36        if !hooks.iter().any(|h| h == addr) {
37            hooks.push(addr);
38        } else {
39            return Err(HookError::HookAlreadyRegistered {});
40        }
41        Ok(self.0.save(storage, &hooks)?)
42    }
43
44    pub fn remove_hook(&self, storage: &mut dyn Storage, addr: Addr) -> Result<(), HookError> {
45        let mut hooks = self.0.load(storage)?;
46        if let Some(p) = hooks.iter().position(|h| h == addr) {
47            hooks.remove(p);
48        } else {
49            return Err(HookError::HookNotRegistered {});
50        }
51        Ok(self.0.save(storage, &hooks)?)
52    }
53
54    pub fn remove_hook_by_index(
55        &self,
56        storage: &mut dyn Storage,
57        index: u64,
58    ) -> Result<Addr, HookError> {
59        let mut hooks = self.0.load(storage)?;
60        let hook = hooks.remove(index as usize);
61        self.0.save(storage, &hooks)?;
62        Ok(hook)
63    }
64
65    pub fn prepare_hooks<F: FnMut(Addr) -> StdResult<SubMsg>>(
66        &self,
67        storage: &dyn Storage,
68        prep: F,
69    ) -> StdResult<Vec<SubMsg>> {
70        self.0
71            .may_load(storage)?
72            .unwrap_or_default()
73            .into_iter()
74            .map(prep)
75            .collect()
76    }
77
78    pub fn prepare_hooks_custom_msg<F: FnMut(Addr) -> StdResult<SubMsg<T>>, T>(
79        &self,
80        storage: &dyn Storage,
81        prep: F,
82    ) -> StdResult<Vec<SubMsg<T>>> {
83        self.0
84            .may_load(storage)?
85            .unwrap_or_default()
86            .into_iter()
87            .map(prep)
88            .collect::<Result<Vec<SubMsg<T>>, _>>()
89    }
90
91    pub fn hook_count(&self, storage: &dyn Storage) -> StdResult<u32> {
92        // The WASM VM (as of version 1) is 32 bit and sets limits for
93        // memory accordingly:
94        // <https://webassembly.github.io/spec/core/syntax/types.html#syntax-limits>. We
95        // can safely return a u32 here as that's the biggest size in
96        // the WASM VM.
97        Ok(self.0.may_load(storage)?.unwrap_or_default().len() as u32)
98    }
99
100    pub fn query_hooks<Q: CustomQuery>(&self, deps: Deps<Q>) -> StdResult<HooksResponse> {
101        let hooks = self.0.may_load(deps.storage)?.unwrap_or_default();
102        let hooks = hooks.into_iter().map(String::from).collect();
103        Ok(HooksResponse { hooks })
104    }
105}
106
107#[cfg(test)]
108mod tests {
109    use super::*;
110    use cosmwasm_std::{coins, testing::mock_dependencies, BankMsg, Empty};
111
112    // Shorthand for an unchecked address.
113    macro_rules! addr {
114        ($x:expr ) => {
115            Addr::unchecked($x)
116        };
117    }
118
119    #[test]
120    fn test_hooks() {
121        let mut deps = mock_dependencies();
122        let storage = &mut deps.storage;
123        let hooks = Hooks::new("hooks");
124
125        // Prepare hooks doesn't through error if no hooks added
126        let msgs = hooks
127            .prepare_hooks(storage, |a| {
128                Ok(SubMsg::reply_always(
129                    BankMsg::Burn {
130                        amount: coins(a.as_str().len() as u128, "uekez"),
131                    },
132                    2,
133                ))
134            })
135            .unwrap();
136        assert_eq!(msgs, vec![]);
137
138        hooks.add_hook(storage, addr!("ekez")).unwrap();
139        hooks.add_hook(storage, addr!("meow")).unwrap();
140
141        assert_eq!(hooks.hook_count(storage).unwrap(), 2);
142
143        hooks.remove_hook_by_index(storage, 0).unwrap();
144
145        assert_eq!(hooks.hook_count(storage).unwrap(), 1);
146
147        let msgs = hooks
148            .prepare_hooks(storage, |a| {
149                Ok(SubMsg::reply_always(
150                    BankMsg::Burn {
151                        amount: coins(a.as_str().len() as u128, "uekez"),
152                    },
153                    2,
154                ))
155            })
156            .unwrap();
157
158        assert_eq!(
159            msgs,
160            vec![SubMsg::reply_always(
161                BankMsg::Burn {
162                    amount: coins(4, "uekez"),
163                },
164                2,
165            )]
166        );
167
168        // Test prepare hooks with custom messages.
169        // In a real world scenario, you would be using something like
170        // TokenFactoryMsg.
171        let msgs = hooks
172            .prepare_hooks_custom_msg(storage, |a| {
173                Ok(SubMsg::<Empty>::reply_always(
174                    BankMsg::Burn {
175                        amount: coins(a.as_str().len() as u128, "uekez"),
176                    },
177                    2,
178                ))
179            })
180            .unwrap();
181
182        assert_eq!(
183            msgs,
184            vec![SubMsg::<Empty>::reply_always(
185                BankMsg::Burn {
186                    amount: coins(4, "uekez"),
187                },
188                2,
189            )]
190        );
191
192        // Query hooks returns all hooks added
193        let HooksResponse { hooks: the_hooks } = hooks.query_hooks(deps.as_ref()).unwrap();
194        assert_eq!(the_hooks, vec![addr!("meow")]);
195
196        // Remove last hook
197        hooks.remove_hook(&mut deps.storage, addr!("meow")).unwrap();
198
199        // Query hooks returns empty vector if no hooks added
200        let HooksResponse { hooks: the_hooks } = hooks.query_hooks(deps.as_ref()).unwrap();
201        let no_hooks: Vec<String> = vec![];
202        assert_eq!(the_hooks, no_hooks);
203    }
204}