Upgrade to Actix V2, bump SQLX version, code cleanup, intergrate ratelimiter (#288)

* Upgrade to Actix V2, bump SQLX version, code cleanup, intergrate ratelimiter

* Add pack file path validation

* Fix compilation error caused by incorrect merge

src/ratelimit/errors.rs

@@ -0,0 +1,45 @@
+//! Errors that can occur during middleware processing stage
+use actix_web::ResponseError;
+use log::*;
+use thiserror::Error;
+
+/// Custom error type. Useful for logging and debugging different kinds of errors.
+/// This type can be converted to Actix Error, which defaults to
+/// InternalServerError
+///
+#[derive(Debug, Error)]
+pub enum ARError {
+    /// Read/Write error on store
+    #[error("read/write operatiion failed: {0}")]
+    ReadWriteError(String),
+
+    /// Identifier error
+    #[error("client identification failed")]
+    IdentificationError,
+    /// Limited Error
+    #[error("You are being ratelimited. Please wait {reset} seconds. {remaining}/{max_requests} remaining.")]
+    LimitedError {
+        max_requests: usize,
+        remaining: usize,
+        reset: u64,
+    },
+}
+
+impl ResponseError for ARError {
+    fn error_response(&self) -> actix_web::web::HttpResponse {
+        match self {
+            Self::LimitedError {
+                max_requests,
+                remaining,
+                reset,
+            } => {
+                let mut response = actix_web::web::HttpResponse::TooManyRequests();
+                response.insert_header(("x-ratelimit-limit", max_requests.to_string()));
+                response.insert_header(("x-ratelimit-remaining", remaining.to_string()));
+                response.insert_header(("x-ratelimit-reset", reset.to_string()));
+                response.body(self.to_string())
+            }
+            _ => actix_web::web::HttpResponse::build(self.status_code()).body(self.to_string()),
+        }
+    }
+}

src/ratelimit/memory.rs

@@ -0,0 +1,252 @@
+//! In memory store for rate limiting
+use actix::prelude::*;
+use dashmap::DashMap;
+use futures::future::{self};
+use log::*;
+use std::sync::Arc;
+use std::time::{Duration, SystemTime, UNIX_EPOCH};
+
+use crate::ratelimit::errors::ARError;
+use crate::ratelimit::{ActorMessage, ActorResponse};
+
+/// Type used to create a concurrent hashmap store
+#[derive(Clone)]
+pub struct MemoryStore {
+    inner: Arc<DashMap<String, (usize, Duration)>>,
+}
+
+impl MemoryStore {
+    /// Create a new hashmap
+    ///
+    /// # Example
+    /// ```rust
+    /// use actix_ratelimit::MemoryStore;
+    ///
+    /// let store = MemoryStore::new();
+    /// ```
+    pub fn new() -> Self {
+        debug!("Creating new MemoryStore");
+        MemoryStore {
+            inner: Arc::new(DashMap::<String, (usize, Duration)>::new()),
+        }
+    }
+
+    #[allow(dead_code)]
+    /// Create a new hashmap with the provided capacity
+    pub fn with_capacity(capacity: usize) -> Self {
+        debug!("Creating new MemoryStore");
+        MemoryStore {
+            inner: Arc::new(DashMap::<String, (usize, Duration)>::with_capacity(
+                capacity,
+            )),
+        }
+    }
+}
+
+/// Actor for memory store
+pub struct MemoryStoreActor {
+    inner: Arc<DashMap<String, (usize, Duration)>>,
+}
+
+impl From<MemoryStore> for MemoryStoreActor {
+    fn from(store: MemoryStore) -> Self {
+        MemoryStoreActor { inner: store.inner }
+    }
+}
+
+impl MemoryStoreActor {
+    /// Starts the memory actor and returns it's address
+    pub fn start(self) -> Addr<Self> {
+        debug!("Started memory store");
+        Supervisor::start(|_| self)
+    }
+}
+
+impl Actor for MemoryStoreActor {
+    type Context = Context<Self>;
+}
+
+impl Supervised for MemoryStoreActor {
+    fn restarting(&mut self, _: &mut Self::Context) {
+        debug!("Restarting memory store");
+    }
+}
+
+impl Handler<ActorMessage> for MemoryStoreActor {
+    type Result = ActorResponse;
+    fn handle(&mut self, msg: ActorMessage, ctx: &mut Self::Context) -> Self::Result {
+        match msg {
+            ActorMessage::Set { key, value, expiry } => {
+                debug!("Inserting key {} with expiry {}", &key, &expiry.as_secs());
+                let future_key = String::from(&key);
+                let now = SystemTime::now();
+                let now = now.duration_since(UNIX_EPOCH).unwrap();
+                self.inner.insert(key, (value, now + expiry));
+                ctx.notify_later(ActorMessage::Remove(future_key), expiry);
+                ActorResponse::Set(Box::pin(future::ready(Ok(()))))
+            }
+            ActorMessage::Update { key, value } => match self.inner.get_mut(&key) {
+                Some(mut c) => {
+                    let val_mut: &mut (usize, Duration) = c.value_mut();
+                    if val_mut.0 > value {
+                        val_mut.0 -= value;
+                    } else {
+                        val_mut.0 = 0;
+                    }
+                    let new_val = val_mut.0;
+                    ActorResponse::Update(Box::pin(future::ready(Ok(new_val))))
+                }
+                None => {
+                    return ActorResponse::Update(Box::pin(future::ready(Err(
+                        ARError::ReadWriteError("memory store: read failed!".to_string()),
+                    ))))
+                }
+            },
+            ActorMessage::Get(key) => {
+                if self.inner.contains_key(&key) {
+                    let val = match self.inner.get(&key) {
+                        Some(c) => c,
+                        None => {
+                            return ActorResponse::Get(Box::pin(future::ready(Err(
+                                ARError::ReadWriteError("memory store: read failed!".to_string()),
+                            ))))
+                        }
+                    };
+                    let val = val.value().0;
+                    ActorResponse::Get(Box::pin(future::ready(Ok(Some(val)))))
+                } else {
+                    ActorResponse::Get(Box::pin(future::ready(Ok(None))))
+                }
+            }
+            ActorMessage::Expire(key) => {
+                let c = match self.inner.get(&key) {
+                    Some(d) => d,
+                    None => {
+                        return ActorResponse::Expire(Box::pin(future::ready(Err(
+                            ARError::ReadWriteError("memory store: read failed!".to_string()),
+                        ))))
+                    }
+                };
+                let dur = c.value().1;
+                let now = SystemTime::now().duration_since(UNIX_EPOCH).unwrap();
+                let res = dur.checked_sub(now).unwrap_or_else(|| Duration::new(0, 0));
+                ActorResponse::Expire(Box::pin(future::ready(Ok(res))))
+            }
+            ActorMessage::Remove(key) => {
+                debug!("Removing key: {}", &key);
+                let val = match self.inner.remove::<String>(&key) {
+                    Some(c) => c,
+                    None => {
+                        return ActorResponse::Remove(Box::pin(future::ready(Err(
+                            ARError::ReadWriteError("memory store: remove failed!".to_string()),
+                        ))))
+                    }
+                };
+                let val = val.1;
+                ActorResponse::Remove(Box::pin(future::ready(Ok(val.0))))
+            }
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[actix_rt::test]
+    async fn test_set() {
+        let store = MemoryStore::new();
+        let addr = MemoryStoreActor::from(store.clone()).start();
+        let res = addr
+            .send(ActorMessage::Set {
+                key: "hello".to_string(),
+                value: 30usize,
+                expiry: Duration::from_secs(5),
+            })
+            .await;
+        let res = res.expect("Failed to send msg");
+        match res {
+            ActorResponse::Set(c) => match c.await {
+                Ok(()) => {}
+                Err(e) => panic!("Shouldn't happen {}", &e),
+            },
+            _ => panic!("Shouldn't happen!"),
+        }
+    }
+
+    #[actix_rt::test]
+    async fn test_get() {
+        let store = MemoryStore::new();
+        let addr = MemoryStoreActor::from(store.clone()).start();
+        let expiry = Duration::from_secs(5);
+        let res = addr
+            .send(ActorMessage::Set {
+                key: "hello".to_string(),
+                value: 30usize,
+                expiry: expiry,
+            })
+            .await;
+        let res = res.expect("Failed to send msg");
+        match res {
+            ActorResponse::Set(c) => match c.await {
+                Ok(()) => {}
+                Err(e) => panic!("Shouldn't happen {}", &e),
+            },
+            _ => panic!("Shouldn't happen!"),
+        }
+        let res2 = addr.send(ActorMessage::Get("hello".to_string())).await;
+        let res2 = res2.expect("Failed to send msg");
+        match res2 {
+            ActorResponse::Get(c) => match c.await {
+                Ok(d) => {
+                    let d = d.unwrap();
+                    assert_eq!(d, 30usize);
+                }
+                Err(e) => panic!("Shouldn't happen {}", &e),
+            },
+            _ => panic!("Shouldn't happen!"),
+        };
+    }
+
+    #[actix_rt::test]
+    async fn test_expiry() {
+        let store = MemoryStore::new();
+        let addr = MemoryStoreActor::from(store.clone()).start();
+        let expiry = Duration::from_secs(3);
+        let res = addr
+            .send(ActorMessage::Set {
+                key: "hello".to_string(),
+                value: 30usize,
+                expiry: expiry,
+            })
+            .await;
+        let res = res.expect("Failed to send msg");
+        match res {
+            ActorResponse::Set(c) => match c.await {
+                Ok(()) => {}
+                Err(e) => panic!("Shouldn't happen {}", &e),
+            },
+            _ => panic!("Shouldn't happen!"),
+        }
+        assert_eq!(addr.connected(), true);
+
+        let res3 = addr.send(ActorMessage::Expire("hello".to_string())).await;
+        let res3 = res3.expect("Failed to send msg");
+        match res3 {
+            ActorResponse::Expire(c) => match c.await {
+                Ok(dur) => {
+                    let now = Duration::from_secs(3);
+                    if dur > now {
+                        panic!("Expiry is invalid!");
+                    } else if dur > now + Duration::from_secs(4) {
+                        panic!("Expiry is invalid!");
+                    }
+                }
+                Err(e) => {
+                    panic!("Shouldn't happen: {}", &e);
+                }
+            },
+            _ => panic!("Shouldn't happen!"),
+        };
+    }
+}

src/ratelimit/middleware.rs

@@ -0,0 +1,279 @@
+//! RateLimiter middleware for actix application
+use crate::ratelimit::errors::ARError;
+use crate::ratelimit::{ActorMessage, ActorResponse};
+use actix::dev::*;
+use actix_web::{
+    dev::{Service, ServiceRequest, ServiceResponse, Transform},
+    error::Error as AWError,
+    http::header::{HeaderName, HeaderValue},
+};
+use futures::future::{ok, Ready};
+use log::*;
+use std::{
+    cell::RefCell,
+    future::Future,
+    ops::Fn,
+    pin::Pin,
+    rc::Rc,
+    task::{Context, Poll},
+    time::Duration,
+};
+
+/// Type that implements the ratelimit middleware.
+///
+/// This accepts _interval_ which specifies the
+/// window size, _max_requests_ which specifies the maximum number of requests in that window, and
+/// _store_ which is essentially a data store used to store client access information. Entry is removed from
+/// the store after _interval_.
+///
+/// # Example
+/// ```rust
+/// # use std::time::Duration;
+/// use actix_ratelimit::{MemoryStore, MemoryStoreActor};
+/// use actix_ratelimit::RateLimiter;
+///
+/// #[actix_rt::main]
+/// async fn main() {
+///     let store = MemoryStore::new();
+///     let ratelimiter = RateLimiter::new(
+///                             MemoryStoreActor::from(store.clone()).start())
+///                         .with_interval(Duration::from_secs(60))
+///                         .with_max_requests(100);
+/// }
+/// ```
+pub struct RateLimiter<T>
+where
+    T: Handler<ActorMessage> + Send + Sync + 'static,
+    T::Context: ToEnvelope<T, ActorMessage>,
+{
+    interval: Duration,
+    max_requests: usize,
+    store: Addr<T>,
+    identifier: Rc<Box<dyn Fn(&ServiceRequest) -> Result<String, ARError>>>,
+    ignore_ips: Vec<String>,
+}
+
+impl<T> RateLimiter<T>
+where
+    T: Handler<ActorMessage> + Send + Sync + 'static,
+    <T as Actor>::Context: ToEnvelope<T, ActorMessage>,
+{
+    /// Creates a new instance of `RateLimiter` with the provided address of `StoreActor`.
+    pub fn new(store: Addr<T>) -> Self {
+        let identifier = |req: &ServiceRequest| {
+            let connection_info = req.connection_info();
+            let ip = connection_info
+                .peer_addr()
+                .ok_or(ARError::IdentificationError)?;
+            Ok(String::from(ip))
+        };
+        RateLimiter {
+            interval: Duration::from_secs(0),
+            max_requests: 0,
+            store,
+            identifier: Rc::new(Box::new(identifier)),
+            ignore_ips: Vec::new(),
+        }
+    }
+
+    /// Specify the interval. The counter for a client is reset after this interval
+    pub fn with_interval(mut self, interval: Duration) -> Self {
+        self.interval = interval;
+        self
+    }
+
+    /// Specify the maximum number of requests allowed in the given interval.
+    pub fn with_max_requests(mut self, max_requests: usize) -> Self {
+        self.max_requests = max_requests;
+        self
+    }
+
+    /// Sets IPs that should be ignored by the ratelimiter
+    pub fn with_ignore_ips(mut self, ignore_ips: Vec<String>) -> Self {
+        self.ignore_ips = ignore_ips;
+        self
+    }
+
+    /// Function to get the identifier for the client request
+    pub fn with_identifier<F: Fn(&ServiceRequest) -> Result<String, ARError> + 'static>(
+        mut self,
+        identifier: F,
+    ) -> Self {
+        self.identifier = Rc::new(Box::new(identifier));
+        self
+    }
+}
+
+impl<T, S, B> Transform<S, ServiceRequest> for RateLimiter<T>
+where
+    T: Handler<ActorMessage> + Send + Sync + 'static,
+    T::Context: ToEnvelope<T, ActorMessage>,
+    S: Service<ServiceRequest, Response = ServiceResponse<B>, Error = AWError> + 'static,
+    S::Future: 'static,
+    B: 'static,
+{
+    type Response = ServiceResponse<B>;
+    type Error = S::Error;
+    type Transform = RateLimitMiddleware<S, T>;
+    type InitError = ();
+    type Future = Ready<Result<Self::Transform, Self::InitError>>;
+
+    fn new_transform(&self, service: S) -> Self::Future {
+        ok(RateLimitMiddleware {
+            service: Rc::new(RefCell::new(service)),
+            store: self.store.clone(),
+            max_requests: self.max_requests,
+            interval: self.interval.as_secs(),
+            identifier: self.identifier.clone(),
+            ignore_ips: self.ignore_ips.clone(),
+        })
+    }
+}
+
+/// Service factory for RateLimiter
+pub struct RateLimitMiddleware<S, T>
+where
+    S: 'static,
+    T: Handler<ActorMessage> + 'static,
+{
+    service: Rc<RefCell<S>>,
+    store: Addr<T>,
+    // Exists here for the sole purpose of knowing the max_requests and interval from RateLimiter
+    max_requests: usize,
+    interval: u64,
+    identifier: Rc<Box<dyn Fn(&ServiceRequest) -> Result<String, ARError> + 'static>>,
+    ignore_ips: Vec<String>,
+}
+
+impl<T, S, B> Service<ServiceRequest> for RateLimitMiddleware<S, T>
+where
+    T: Handler<ActorMessage> + 'static,
+    S: Service<ServiceRequest, Response = ServiceResponse<B>, Error = AWError> + 'static,
+    S::Future: 'static,
+    B: 'static,
+    T::Context: ToEnvelope<T, ActorMessage>,
+{
+    type Response = ServiceResponse<B>;
+    type Error = S::Error;
+    type Future = Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>>>>;
+
+    fn poll_ready(&self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
+        self.service.borrow_mut().poll_ready(cx)
+    }
+
+    fn call(&self, req: ServiceRequest) -> Self::Future {
+        let store = self.store.clone();
+        let srv = self.service.clone();
+        let max_requests = self.max_requests;
+        let interval = Duration::from_secs(self.interval);
+        let identifier = self.identifier.clone();
+        let ignore_ips = self.ignore_ips.clone();
+        Box::pin(async move {
+            let identifier: String = (identifier)(&req)?;
+            if ignore_ips.contains(&identifier) {
+                let fut = srv.call(req);
+                let res = fut.await?;
+                return Ok(res);
+            }
+            let remaining: ActorResponse = store
+                .send(ActorMessage::Get(String::from(&identifier)))
+                .await
+                .map_err(|_| ARError::IdentificationError)?;
+            match remaining {
+                ActorResponse::Get(opt) => {
+                    let opt = opt.await?;
+                    if let Some(c) = opt {
+                        // Existing entry in store
+                        let expiry = store
+                            .send(ActorMessage::Expire(String::from(&identifier)))
+                            .await
+                            .map_err(|_| ARError::ReadWriteError("Setting timeout".to_string()))?;
+                        let reset: Duration = match expiry {
+                            ActorResponse::Expire(dur) => dur.await?,
+                            _ => unreachable!(),
+                        };
+                        if c == 0 {
+                            info!("Limit exceeded for client: {}", &identifier);
+                            Err(ARError::LimitedError {
+                                max_requests,
+                                remaining: c,
+                                reset: reset.as_secs(),
+                            }
+                            .into())
+                        } else {
+                            // Decrement value
+                            let res: ActorResponse = store
+                                .send(ActorMessage::Update {
+                                    key: identifier,
+                                    value: 1,
+                                })
+                                .await
+                                .map_err(|_| {
+                                    ARError::ReadWriteError("Decrementing ratelimit".to_string())
+                                })?;
+                            let updated_value: usize = match res {
+                                ActorResponse::Update(c) => c.await?,
+                                _ => unreachable!(),
+                            };
+                            // Execute the request
+                            let fut = srv.call(req);
+                            let mut res = fut.await?;
+                            let headers = res.headers_mut();
+                            // Safe unwraps, since usize is always convertible to string
+                            headers.insert(
+                                HeaderName::from_static("x-ratelimit-limit"),
+                                HeaderValue::from_str(max_requests.to_string().as_str())?,
+                            );
+                            headers.insert(
+                                HeaderName::from_static("x-ratelimit-remaining"),
+                                HeaderValue::from_str(updated_value.to_string().as_str())?,
+                            );
+                            headers.insert(
+                                HeaderName::from_static("x-ratelimit-reset"),
+                                HeaderValue::from_str(reset.as_secs().to_string().as_str())?,
+                            );
+                            Ok(res)
+                        }
+                    } else {
+                        // New client, create entry in store
+                        let current_value = max_requests - 1;
+                        let res = store
+                            .send(ActorMessage::Set {
+                                key: String::from(&identifier),
+                                value: current_value,
+                                expiry: interval,
+                            })
+                            .await
+                            .map_err(|_| {
+                                ARError::ReadWriteError("Creating store entry".to_string())
+                            })?;
+                        match res {
+                            ActorResponse::Set(c) => c.await?,
+                            _ => unreachable!(),
+                        }
+                        let fut = srv.call(req);
+                        let mut res = fut.await?;
+                        let headers = res.headers_mut();
+                        // Safe unwraps, since usize is always convertible to string
+                        headers.insert(
+                            HeaderName::from_static("x-ratelimit-limit"),
+                            HeaderValue::from_str(max_requests.to_string().as_str()).unwrap(),
+                        );
+                        headers.insert(
+                            HeaderName::from_static("x-ratelimit-remaining"),
+                            HeaderValue::from_str(current_value.to_string().as_str()).unwrap(),
+                        );
+                        headers.insert(
+                            HeaderName::from_static("x-ratelimit-reset"),
+                            HeaderValue::from_str(interval.as_secs().to_string().as_str()).unwrap(),
+                        );
+                        Ok(res)
+                    }
+                }
+                _ => {
+                    unreachable!();
+                }
+            }
+        })
+    }
+}

src/ratelimit/mod.rs

@@ -0,0 +1,64 @@
+use std::future::Future;
+use std::marker::Send;
+use std::pin::Pin;
+use std::time::Duration;
+
+use crate::ratelimit::errors::ARError;
+use actix::dev::*;
+
+pub mod errors;
+pub mod memory;
+/// The code for this module was directly taken from https://github.com/TerminalWitchcraft/actix-ratelimit
+/// with some modifications including upgrading it to Actix 4!
+pub mod middleware;
+
+/// Represents message that can be handled by a `StoreActor`
+pub enum ActorMessage {
+    /// Get the remaining count based on the provided identifier
+    Get(String),
+    /// Set the count of the client identified by `key` to `value` valid for `expiry`
+    Set {
+        key: String,
+        value: usize,
+        expiry: Duration,
+    },
+    /// Change the value of count for the client identified by `key` by `value`
+    Update { key: String, value: usize },
+    /// Get the expiration time for the client.
+    Expire(String),
+    /// Remove the client from the store
+    Remove(String),
+}
+
+impl Message for ActorMessage {
+    type Result = ActorResponse;
+}
+
+/// Wrapper type for `Pin<Box<dyn Future>>` type
+pub type Output<T> = Pin<Box<dyn Future<Output = Result<T, ARError>> + Send>>;
+
+/// Represents data returned in response to `Messages` by a `StoreActor`
+pub enum ActorResponse {
+    /// Returned in response to [Messages::Get](enum.Messages.html)
+    Get(Output<Option<usize>>),
+    /// Returned in response to [Messages::Set](enum.Messages.html)
+    Set(Output<()>),
+    /// Returned in response to [Messages::Update](enum.Messages.html)
+    Update(Output<usize>),
+    /// Returned in response to [Messages::Expire](enum.Messages.html)
+    Expire(Output<Duration>),
+    /// Returned in response to [Messages::Remove](enum.Messages.html)
+    Remove(Output<usize>),
+}
+
+impl<A, M> MessageResponse<A, M> for ActorResponse
+where
+    A: Actor,
+    M: actix::Message<Result = ActorResponse>,
+{
+    fn handle(self, _: &mut A::Context, tx: Option<OneshotSender<Self>>) {
+        if let Some(tx) = tx {
+            let _ = tx.send(self);
+        }
+    }
+}