|
| 1 | +#![cfg_attr(not(feature = "std"), no_std)] |
| 2 | + |
| 3 | +use ink::prelude::string::String; |
| 4 | + |
| 5 | +#[derive(Debug, PartialEq, Eq, scale::Encode, scale::Decode)] |
| 6 | +#[cfg_attr(feature = "std", derive(scale_info::TypeInfo))] |
| 7 | +pub enum ReentrancyError { |
| 8 | + ReentrantCall, |
| 9 | +} |
| 10 | + |
| 11 | +/// Simple mutex-based reentrancy guard (OpenZeppelin-style) |
| 12 | +#[derive(Default)] |
| 13 | +pub struct ReentrancyGuard { |
| 14 | + locked: bool, |
| 15 | +} |
| 16 | + |
| 17 | +impl ReentrancyGuard { |
| 18 | + pub fn new() -> Self { |
| 19 | + Self { locked: false } |
| 20 | + } |
| 21 | + |
| 22 | + /// Enter protected section |
| 23 | + pub fn enter(&mut self) -> Result<(), ReentrancyError> { |
| 24 | + if self.locked { |
| 25 | + return Err(ReentrancyError::ReentrantCall); |
| 26 | + } |
| 27 | + self.locked = true; |
| 28 | + Ok(()) |
| 29 | + } |
| 30 | + |
| 31 | + /// Exit protected section |
| 32 | + pub fn exit(&mut self) { |
| 33 | + self.locked = false; |
| 34 | + } |
| 35 | +} |
| 36 | + |
| 37 | +/// Helper macro to simplify usage |
| 38 | +#[macro_export] |
| 39 | +macro_rules! non_reentrant { |
| 40 | + ($self:ident, $body:block) => {{ |
| 41 | + $self.reentrancy_guard.enter().map_err(|_| ())?; |
| 42 | + let result = (|| $body)(); |
| 43 | + $self.reentrancy_guard.exit(); |
| 44 | + result |
| 45 | + }}; |
| 46 | +} |
| 47 | + |
| 48 | +/// Optional: Gas limit wrapper for external calls |
| 49 | +pub fn safe_external_call<F, T>(call: F, gas_limit: u64) -> Result<T, String> |
| 50 | +where |
| 51 | + F: FnOnce() -> Result<T, String>, |
| 52 | +{ |
| 53 | + // In real ink!, gas control is limited, but we simulate safety check |
| 54 | + if gas_limit == 0 { |
| 55 | + return Err("Gas limit too low".into()); |
| 56 | + } |
| 57 | + |
| 58 | + call() |
| 59 | +} |
0 commit comments