Skip to main content

kaspa_txscript/
runtime_resource_meter.rs

1use kaspa_consensus_core::mass::ScriptUnits;
2use kaspa_txscript_errors::TxScriptError;
3
4/// Tracks resource consumption (signature operations and pushed bytes, both expressed
5/// as script units) during script execution against a committed budget.
6#[derive(Debug, Clone)]
7pub struct RuntimeResourceMeter {
8    used_sig_ops: u16,
9    sigop_script_units: ScriptUnits,
10    script_units_limit: ScriptUnits,
11    remaining_script_units: ScriptUnits,
12}
13
14impl RuntimeResourceMeter {
15    pub fn new_script_units(sigop_script_units: ScriptUnits, script_units_limit: ScriptUnits) -> Self {
16        Self { used_sig_ops: 0, sigop_script_units, script_units_limit, remaining_script_units: script_units_limit }
17    }
18
19    pub fn used_sig_ops(&self) -> u16 {
20        self.used_sig_ops
21    }
22
23    pub fn used_script_units(&self) -> ScriptUnits {
24        self.script_units_limit - self.remaining_script_units
25    }
26
27    pub fn consume_script_units(&mut self, units: ScriptUnits) -> Result<(), TxScriptError> {
28        match self.remaining_script_units.checked_sub(units) {
29            Some(new_remaining) => {
30                self.remaining_script_units = new_remaining;
31                Ok(())
32            }
33            None => {
34                let overflow = units.0 - self.remaining_script_units.0;
35                let used_units = self.script_units_limit.0.saturating_add(overflow);
36                Err(TxScriptError::ExceededCommittedScriptUnits { used: used_units, limit: self.script_units_limit.0 })
37            }
38        }
39    }
40
41    pub fn consume_sig_op_cost(&mut self, count: u16) -> Result<(), TxScriptError> {
42        self.consume_script_units(ScriptUnits((count as u64).saturating_mul(self.sigop_script_units.0)))?;
43        self.used_sig_ops = self.used_sig_ops.saturating_add(count);
44        Ok(())
45    }
46
47    pub fn charge_newly_pushed_bytes(&mut self, pushed_bytes_delta: u64) -> Result<(), TxScriptError> {
48        // Pushed bytes are charged 1:1 as script units.
49        self.consume_script_units(pushed_bytes_delta.into())
50    }
51}
52
53#[cfg(test)]
54mod tests {
55    use kaspa_core::assert_match;
56
57    use super::*;
58
59    #[test]
60    fn script_units_meter_charges_sigops_in_script_units() {
61        let mut meter = RuntimeResourceMeter::new_script_units(ScriptUnits(100), ScriptUnits(250));
62
63        assert_eq!(meter.consume_sig_op_cost(2), Ok(()));
64        assert_eq!(meter.used_sig_ops(), 2);
65        assert_eq!(meter.used_script_units(), ScriptUnits(200));
66
67        assert_eq!(meter.consume_sig_op_cost(1), Err(TxScriptError::ExceededCommittedScriptUnits { used: 300, limit: 250 }));
68        assert_eq!(meter.used_sig_ops(), 2);
69        assert_eq!(meter.used_script_units(), ScriptUnits(200));
70    }
71
72    #[test]
73    fn script_units_meter_saturates_exceeded_used_units() {
74        let mut meter = RuntimeResourceMeter::new_script_units(ScriptUnits(0), ScriptUnits(100));
75
76        assert_eq!(meter.consume_script_units(ScriptUnits(60)), Ok(()));
77        assert_eq!(
78            meter.consume_script_units(ScriptUnits(u64::MAX)),
79            Err(TxScriptError::ExceededCommittedScriptUnits { used: u64::MAX, limit: 100 })
80        );
81        assert_eq!(meter.used_script_units(), ScriptUnits(60));
82    }
83
84    #[test]
85    fn script_units_meter_rejects_u64_max_charge_without_panicking() {
86        let mut meter = RuntimeResourceMeter::new_script_units(ScriptUnits(0), ScriptUnits(100));
87
88        assert_eq!(
89            meter.consume_script_units(ScriptUnits(u64::MAX)),
90            Err(TxScriptError::ExceededCommittedScriptUnits { used: u64::MAX, limit: 100 })
91        );
92        assert_eq!(meter.used_script_units(), ScriptUnits(0));
93    }
94
95    #[test]
96    fn script_units_meter_charges_only_newly_pushed_bytes() {
97        let mut meter = RuntimeResourceMeter::new_script_units(ScriptUnits(0), ScriptUnits(20));
98
99        assert_eq!(meter.charge_newly_pushed_bytes(7), Ok(()));
100        assert_eq!(meter.used_script_units(), ScriptUnits(7));
101
102        // Charging zero is a no-op because only newly pushed bytes are charged.
103        assert_eq!(meter.charge_newly_pushed_bytes(0), Ok(()));
104        assert_eq!(meter.used_script_units(), ScriptUnits(7));
105
106        assert_eq!(meter.charge_newly_pushed_bytes(9), Ok(()));
107        assert_eq!(meter.used_script_units(), ScriptUnits(16));
108    }
109
110    #[test]
111    fn meter_bounds_do_not_panic() {
112        let mut max_used_sig_ops_meter = RuntimeResourceMeter::new_script_units(ScriptUnits(0), ScriptUnits(0));
113        assert_eq!(max_used_sig_ops_meter.consume_sig_op_cost(u16::MAX), Ok(()));
114        assert_eq!(max_used_sig_ops_meter.used_sig_ops(), u16::MAX);
115        assert_eq!(max_used_sig_ops_meter.consume_sig_op_cost(1), Ok(()));
116        assert_eq!(max_used_sig_ops_meter.used_sig_ops(), u16::MAX);
117
118        let mut max_units_meter = RuntimeResourceMeter::new_script_units(ScriptUnits(0), ScriptUnits(u64::MAX));
119        assert_eq!(max_units_meter.charge_newly_pushed_bytes(u64::MAX), Ok(()));
120        assert_eq!(max_units_meter.used_script_units(), ScriptUnits(u64::MAX));
121        assert_match!(
122            max_units_meter.charge_newly_pushed_bytes(u64::MAX),
123            Err(TxScriptError::ExceededCommittedScriptUnits { used: _, limit: _ })
124        ); // Charging u64::MAX again should excceed budget and not panic due to overflow.
125        assert_eq!(max_units_meter.used_script_units(), ScriptUnits(u64::MAX)); // On overflow, used_script_units are not affected.
126
127        // Now checking overflow behaviour when starting non-max u64 values.
128        let mut max_units_meter = RuntimeResourceMeter::new_script_units(ScriptUnits(0), ScriptUnits(u64::MAX));
129        assert_eq!(max_units_meter.charge_newly_pushed_bytes(u64::MAX - 1), Ok(()));
130        assert_eq!(max_units_meter.used_script_units(), ScriptUnits(u64::MAX - 1));
131        assert_match!(
132            max_units_meter.charge_newly_pushed_bytes(2),
133            Err(TxScriptError::ExceededCommittedScriptUnits { used: _, limit: _ })
134        ); // Charging 2 should excceed budget and not panic due to overflow.
135        assert_eq!(max_units_meter.used_script_units(), ScriptUnits(u64::MAX - 1)); // On overflow, used_script_units are not affected.
136    }
137}