brakes 0.1.0

A distributed rate limiting library
Documentation 
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

pub mod fixed_window;
pub mod leaky_bucket;
pub mod sliding_window;
pub mod token_bucket;

use crate::backend::BackendError;
use fixed_window::FixedWindowInstance;
use leaky_bucket::LeakyBucketInstance;
use serde::{Deserialize, Serialize};
use sliding_window::SlidingWindowInstance;
use std::{
    error::Error,
    fmt::{self, Debug, Display},
};
use token_bucket::TokenBucketInstance;

pub trait LimiterType: Clone {
    fn is_ratelimited(&self, value: Option<Vec<u8>>) -> Result<Vec<u8>, RateLimiterError>;
    fn window_instance(&self, value: Vec<u8>) -> Result<LimiterInstance, RateLimiterError>;
}

#[derive(Debug)]
pub enum LimiterInstance {
    FixedWindowInstance(FixedWindowInstance),
    SlidingWindowInstance(SlidingWindowInstance),
    TokenBucketInstance(TokenBucketInstance),
    LeakyBucketInstance(LeakyBucketInstance),
}

impl LimiterInstance {
    pub fn as_fixed_window_instance(self) -> Result<FixedWindowInstance, RateLimiterError> {
        match self {
            Self::FixedWindowInstance(i) => Ok(i),
            _ => Err(RateLimiterError::MalformedValue(None)),
        }
    }

    pub fn as_sliding_window_instance(self) -> Result<SlidingWindowInstance, RateLimiterError> {
        match self {
            Self::SlidingWindowInstance(i) => Ok(i),
            _ => Err(RateLimiterError::MalformedValue(None)),
        }
    }

    pub fn as_token_bucket_instance(self) -> Result<TokenBucketInstance, RateLimiterError> {
        match self {
            Self::TokenBucketInstance(i) => Ok(i),
            _ => Err(RateLimiterError::MalformedValue(None)),
        }
    }

    pub fn as_leaky_bucket_instance(self) -> Result<LeakyBucketInstance, RateLimiterError> {
        match self {
            Self::LeakyBucketInstance(i) => Ok(i),
            _ => Err(RateLimiterError::MalformedValue(None)),
        }
    }
}

trait SerializableInstance: Debug + Serialize + for<'de> Deserialize<'de> {
    fn from_bytes(bytes: Vec<u8>) -> Result<Self, RateLimiterError> {
        bincode::deserialize(&bytes).map_err(|e| RateLimiterError::MalformedValue(Some(e)))
    }
    fn to_bytes(self) -> Result<Vec<u8>, RateLimiterError> {
        bincode::serialize(&self).map_err(|e| RateLimiterError::MalformedValue(Some(e)))
    }
}

#[derive(Debug)]
pub enum RateLimiterError {
    MalformedValue(Option<bincode::Error>),
    RateExceeded,
    BackendError(BackendError),
    BackendConflict,
}

impl Display for RateLimiterError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            RateLimiterError::MalformedValue(Some(e)) => std::fmt::Display::fmt(&e, f),
            RateLimiterError::MalformedValue(None) => write!(f, "malformed value"),
            RateLimiterError::RateExceeded => write!(f, "rate exceeded"),
            RateLimiterError::BackendError(e) => std::fmt::Display::fmt(&e, f),
            RateLimiterError::BackendConflict => write!(f, "backend value conflict"),
        }
    }
}

impl Error for RateLimiterError {}