lightning_signer/util/
velocity.rs

1use serde::Deserialize;
2
3use crate::prelude::*;
4use core::cmp::min;
5use core::fmt::{self, Debug, Formatter};
6
7/// Limit velocity per unit time.
8///
9/// We track velocity in intervals instead of tracking each send, to keep
10/// storage requirements constant.
11#[derive(Clone)]
12pub struct VelocityControl {
13    /// start second for the current velocity epoch
14    pub start_sec: u64,
15    /// the number of seconds represented by a bucket
16    pub bucket_interval: u32,
17    /// each bucket entry is the total velocity detected in that interval, in satoshi
18    pub buckets: Vec<u64>,
19    /// the limit, or MAX if the control is disabled
20    pub limit: u64,
21}
22
23impl Debug for VelocityControl {
24    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
25        f.debug_struct("VelocityControl")
26            .field("start_sec", &self.start_sec)
27            .field("bucket_interval", &self.bucket_interval)
28            .field("buckets", &format_args!("{:?}", self.buckets))
29            .field("limit", &self.limit)
30            .finish()
31    }
32}
33
34/// The total interval in which to track velocity
35#[derive(Clone, Copy, Debug, Deserialize)]
36pub enum VelocityControlIntervalType {
37    /// Tracked in 5 minute sub-intervals
38    Hourly,
39    /// Tracked in 1 hour sub-intervals
40    Daily,
41    /// Unlimited velocity
42    Unlimited,
43}
44
45/// A specifier for creating velocity controls
46#[derive(Clone, Copy, Debug, Deserialize)]
47pub struct VelocityControlSpec {
48    /// The limit per interval in msat
49    pub limit_msat: u64,
50    /// The interval type
51    pub interval_type: VelocityControlIntervalType,
52}
53
54impl VelocityControlSpec {
55    /// A velocity control spec for controls which don't limit velocity
56    pub const UNLIMITED: VelocityControlSpec = VelocityControlSpec {
57        limit_msat: 0,
58        interval_type: VelocityControlIntervalType::Unlimited,
59    };
60}
61
62impl VelocityControl {
63    /// Create a velocity control with arbitrary specified intervals
64    /// current_sec: the current second
65    /// num_buckets: how many buckets to keep track of
66    /// bucket_interval: each bucket represents this number of seconds
67    /// limit: the total velocity limit when summing the buckets
68    pub fn new_with_intervals(limit_msat: u64, bucket_interval: u32, num_buckets: usize) -> Self {
69        assert!(bucket_interval > 0 && num_buckets > 0);
70        let mut buckets = Vec::new();
71        buckets.resize(num_buckets, 0);
72        VelocityControl { start_sec: 0, bucket_interval, buckets, limit: limit_msat }
73    }
74
75    /// Create an unlimited velocity control (i.e. no actual control)
76    pub fn new_unlimited(bucket_interval: u32, num_buckets: usize) -> Self {
77        assert!(bucket_interval > 0 && num_buckets > 0);
78        let mut buckets = Vec::new();
79        buckets.resize(num_buckets, 0);
80        VelocityControl { start_sec: 0, bucket_interval, buckets, limit: u64::MAX }
81    }
82
83    /// Create a velocity control with the given interval type
84    pub fn new(spec: VelocityControlSpec) -> Self {
85        let (limit, bucket_interval, num_buckets) = Self::spec_to_triple(&spec);
86        Self::new_with_intervals(limit, bucket_interval, num_buckets)
87    }
88
89    /// Whether the spec matches this control
90    pub fn spec_matches(&self, spec: &VelocityControlSpec) -> bool {
91        let (limit, bucket_interval, num_buckets) = Self::spec_to_triple(spec);
92        self.limit == limit
93            && self.bucket_interval == bucket_interval
94            && self.buckets.len() == num_buckets
95    }
96
97    /// Update this control to match the given spec.  If the spec does not
98    /// match the previous spec, the control is reset.
99    pub fn update_spec(&mut self, spec: &VelocityControlSpec) {
100        if !self.spec_matches(spec) {
101            let (limit, bucket_interval, num_buckets) = Self::spec_to_triple(spec);
102            self.limit = limit;
103            self.bucket_interval = bucket_interval;
104            self.buckets = Vec::new();
105            self.buckets.resize(num_buckets, 0);
106            self.start_sec = 0;
107        }
108    }
109
110    // Convert a spec to a (limit, bucket_interval, num_buckets) triple
111    fn spec_to_triple(spec: &VelocityControlSpec) -> (u64, u32, usize) {
112        match spec.interval_type {
113            VelocityControlIntervalType::Hourly => (spec.limit_msat, 300, 12),
114            VelocityControlIntervalType::Daily => (spec.limit_msat, 3600, 24),
115            VelocityControlIntervalType::Unlimited => (u64::MAX, 300, 12),
116        }
117    }
118
119    /// Load from persistence
120    pub fn load_from_state(spec: VelocityControlSpec, state: (u64, Vec<u64>)) -> Self {
121        let control = Self::new(spec);
122        control.with_state(state)
123    }
124
125    fn with_state(mut self, state: (u64, Vec<u64>)) -> VelocityControl {
126        self.start_sec = state.0;
127        self.buckets = state.1;
128        self
129    }
130
131    /// Get the state for persistence
132    pub fn get_state(&self) -> (u64, Vec<u64>) {
133        (self.start_sec, self.buckets.clone())
134    }
135
136    /// Whether this instance is unlimited (no control)
137    pub fn is_unlimited(&self) -> bool {
138        self.limit == u64::MAX
139    }
140
141    /// Update the velocity given the passage of time given by `current_sec`
142    /// and the given velocity.  If the limit would be exceeded, the given velocity
143    /// is not inserted and false is returned.
144    pub fn insert(&mut self, current_sec: u64, velocity_msat: u64) -> bool {
145        let nshift = (current_sec - self.start_sec) / self.bucket_interval as u64;
146        let len = self.buckets.len();
147        let nshift = min(len, nshift as usize);
148        self.buckets.resize(len - nshift, 0);
149        for _ in 0..nshift {
150            self.buckets.insert(0, 0);
151        }
152        self.start_sec = current_sec - (current_sec % self.bucket_interval as u64);
153        let current_velocity = self.velocity();
154        if current_velocity.saturating_add(velocity_msat) > self.limit {
155            false
156        } else {
157            self.buckets[0] = self.buckets[0].saturating_add(velocity_msat);
158            true
159        }
160    }
161
162    /// Clear the control (e.g. when the user manually approves)
163    pub fn clear(&mut self) {
164        for bucket in self.buckets.iter_mut() {
165            *bucket = 0;
166        }
167    }
168
169    /// The total msat velocity in the tracked interval.
170    ///
171    /// If this is an unlimited control, zero is always returned.
172    pub fn velocity(&self) -> u64 {
173        let mut sum = 0u64;
174        for bucket in self.buckets.iter() {
175            sum = sum.saturating_add(*bucket)
176        }
177        sum
178    }
179}
180
181#[cfg(test)]
182mod tests {
183    use crate::util::velocity::VelocityControl;
184
185    #[test]
186    fn test_velocity() {
187        let mut c = VelocityControl::new_with_intervals(100, 10, 4);
188        assert_eq!(c.velocity(), 0);
189        assert!(c.insert(1100, 90));
190        assert_eq!(c.velocity(), 90);
191        assert!(!c.insert(1101, 11));
192        assert_eq!(c.velocity(), 90);
193        assert!(c.insert(1101, 10));
194        assert_eq!(c.velocity(), 100);
195        assert!(!c.insert(1139, 90));
196        assert_eq!(c.velocity(), 100);
197        assert!(c.insert(1140, 90));
198        assert_eq!(c.velocity(), 90);
199        assert!(c.insert(1150, 5));
200        assert_eq!(c.velocity(), 95);
201        assert!(c.insert(1180, 80));
202        assert_eq!(c.velocity(), 85);
203        assert!(c.insert(1190, 1));
204        assert_eq!(c.velocity(), 81);
205    }
206
207    #[test]
208    fn test_unlimited() {
209        let mut c = VelocityControl::new_unlimited(10, 4);
210        assert!(c.insert(0, u64::MAX - 1));
211        assert_eq!(c.velocity(), u64::MAX - 1);
212        assert!(c.insert(0, 1));
213        assert_eq!(c.velocity(), u64::MAX);
214        assert!(c.insert(0, 1));
215        assert_eq!(c.velocity(), u64::MAX);
216    }
217}
lightning_signer/util/velocity.rs

lightning_signer/util/
velocity.rs
