1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132

//! Edge rate limiter and supporting functions.
//!
pub(crate) mod handle;

use std::time::Duration;

pub use handle::ERLError;
use handle::{ERLHandle, PenaltyboxHandle, RateCounterHandle};

/// An edge rate limiter which includes a ratecounter and a penaltybox.
///
pub struct ERL {
    handle: ERLHandle,
}

/// A rate counter that can be used with a edge rate limiter or stand alone for counting and rate calculations
///
pub struct RateCounter {
    handle: RateCounterHandle,
}

/// To be used for picking the window in a rate counter lookup_rate or a ERL check_rate
///
#[repr(u32)]
pub enum RateWindow {
    /// A one secound rate window
    OneSec = 1,
    /// A ten secound rate window
    TenSecs = 10,
    /// A sixty secound rate window
    SixtySecs = 60,
}

/// To be used for picking the duration in a rate counter lookup_count call
///
#[repr(u32)]
pub enum CounterDuration {
    /// A one second counter window
    TenSec = 10,
    /// A twenty second counter window
    TwentySecs = 20,
    /// A thirty second counter window
    ThirtySecs = 30,
    /// A forty second counter window
    FortySecs = 40,
    /// A fifty second counter window
    FiftySecs = 50,
    /// A sixty second counter window
    SixtySecs = 60,
}

/// A penaltybox that can be used with the edge rate limiter or stand alone for adding and checking if some entry is in the data set.
///
pub struct Penaltybox {
    handle: PenaltyboxHandle,
}

impl ERL {
    /// Open a ERL with the given ratecounter and penaltybox.
    ///
    pub fn open(ratecounter: RateCounter, penaltybox: Penaltybox) -> Self {
        Self {
            handle: ERLHandle::open(ratecounter, penaltybox),
        }
    }

    /// Increment an entry in a rate counter and check if the client has exceeded some average number of requests per second (RPS) over the window.
    /// If the client is over the rps limit for the window, add to the penaltybox for ttl.
    /// Valid ttl span is 1m to 1h and TTL value is truncated to the nearest minute.
    ///
    pub fn check_rate(
        &self,
        entry: &str,
        delta: u32,
        window: RateWindow,
        limit: u32,
        ttl: Duration,
    ) -> Result<bool, ERLError> {
        self.handle.check_rate(entry, delta, window, limit, ttl)
    }
}

impl RateCounter {
    /// Open a RateCounter with the given name for a ratecounter.
    ///
    pub fn open(ratecounter_name: &str) -> Self {
        Self {
            handle: RateCounterHandle::open(ratecounter_name),
        }
    }

    /// Increment an entry in the ratecounter by delta.
    ///
    pub fn increment(&self, entry: &str, delta: u32) -> Result<(), ERLError> {
        self.handle.increment(entry, delta)
    }

    /// Lookup the current rate for entry in the ratecounter for a window.
    ///
    pub fn lookup_rate(&self, entry: &str, window: RateWindow) -> Result<u32, ERLError> {
        self.handle.lookup_rate(entry, window)
    }

    /// Lookup the current count for entry in the ratecounter for duration.
    ///
    pub fn lookup_count(&self, entry: &str, duration: CounterDuration) -> Result<u32, ERLError> {
        self.handle.lookup_count(entry, duration)
    }
}

impl Penaltybox {
    /// Open a Penaltybox identified by the given name.
    ///
    pub fn open(penaltybox_name: &str) -> Self {
        Self {
            handle: PenaltyboxHandle::open(penaltybox_name),
        }
    }

    /// Add entry to a the penaltybox for the duration of ttl.
    /// Valid ttl span is 1m to 1h and TTL value is truncated to the nearest minute.
    ///
    pub fn add(&self, entry: &str, ttl: Duration) -> Result<(), ERLError> {
        self.handle.add(entry, ttl)
    }

    /// Check if entry is in the penaltybox.
    ///  
    pub fn has(&self, entry: &str) -> Result<bool, ERLError> {
        self.handle.has(entry)
    }
}