add rate limiting for download

Cargo.toml

@@ -7,26 +7,28 @@ edition = "2021"
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 
 [dependencies]
-actix-web = { version = "4.1.0", default-features = false, features = [
+actix-web = { version = "4.2.1", default-features = false, features = [
   "macros",
   "compress-gzip",
   "compress-zstd",
 ] }
-sqlx = { version = "0.6.0", default-features = false, features = [
+sqlx = { version = "0.6.2", default-features = false, features = [
   "runtime-tokio-rustls",
   "postgres",
   "time",
 ] }
-env_logger = "0.9.0"
+env_logger = "0.9.1"
 log = "0.4.17"
-actix-files = "0.6.1"
+actix-files = "0.6.2"
-tokio = { version = "1.19.2", features = ["rt", "macros", "sync"] }
+tokio = { version = "1.21.2", features = ["rt", "macros", "sync"] }
 actix-multipart = "0.4.0"
-futures-util = "0.3.21"
+futures-util = "0.3.24"
 rand = "0.8.5"
-time = "0.3.11"
+time = "0.3.14"
 htmlescape = "0.3.1"
-urlencoding = "2.1.0"
+urlencoding = "2.1.2"
 tree_magic_mini = { version = "3.0.3", features = ["with-gpl-data"] }
 mime = "0.3.16"
-url = "2.2.2"
+url = "2.3.1"
+actix-governor = "0.3.2"
+governor = "0.4.2"

README.md

@@ -18,12 +18,16 @@ To run the software directly, use the compiling instructions below.
 
 ### General configuration
 
-| environment variable | default value  |
+| environment variable  | default value  | description                                    |
-| -------------------- | -------------- |
+| --------------------- | -------------- | ---------------------------------------------- |
-| STATIC_DIR           | ./static       |
+| STATIC_DIR            | ./static       | directory to generate "static" files into      |
-| FILES_DIR            | ./files        |
+| FILES_DIR             | ./files        | directory to save uploaded files into          |
-| UPLOAD_MAX_BYTES     | 8388608 (8MiB) |
+| UPLOAD_MAX_BYTES      | 8388608 (8MiB) | maximum size for uploaded files                |
-| BIND_ADDRESS         | 0.0.0.0:8000   |
+| BIND_ADDRESS          | 0.0.0.0:8000   | address to bind the server to                  |
+| RATE_LIMIT            | true           | whether download rate should be limited        |
+| RATE_LIMIT_PROXIED    | false          | whether rate limit should read x-forwarded-for |
+| RATE_LIMIT_PER_SECOND | 60             | seconds to wait between requests               |
+| RATE_LIMIT_BURST      | 1440           | allowed request burst                          |
 
 ### Database configuration
 

src/config.rs

@@ -10,6 +10,10 @@ pub struct Config {
     pub files_dir: PathBuf,
     pub max_file_size: Option<u64>,
     pub no_auth_limits: Option<NoAuthLimits>,
+    pub enable_rate_limit: bool,
+    pub proxied: bool,
+    pub rate_limit_per_second: u64,
+    pub rate_limit_burst: u32,
 }
 
 #[derive(Clone)]
@@ -36,11 +40,26 @@ pub async fn get_config() -> Config {
 
     let no_auth_limits = get_no_auth_limits();
 
+    let enable_rate_limit = matches!(env::var("RATE_LIMIT").as_deref(), Ok("true") | Err(_));
+    let proxied = env::var("PROXIED").as_deref() == Ok("true");
+    let rate_limit_per_second = env::var("RATE_LIMIT_PER_SECOND")
+        .ok()
+        .and_then(|rate_limit| rate_limit.parse().ok())
+        .unwrap_or(60);
+    let rate_limit_burst = env::var("RATE_LIMIT_BURST")
+        .ok()
+        .and_then(|rate_limit| rate_limit.parse().ok())
+        .unwrap_or(1440);
+
     Config {
         static_dir,
         files_dir,
         max_file_size,
         no_auth_limits,
+        enable_rate_limit,
+        proxied,
+        rate_limit_per_second,
+        rate_limit_burst,
     }
 }
 

src/main.rs

@@ -3,10 +3,13 @@ mod db;
 mod deleter;
 mod download;
 mod multipart;
+mod rate_limit;
 mod template;
 mod upload;
 
+use crate::rate_limit::ForwardedPeerIpKeyExtractor;
 use actix_files::Files;
+use actix_governor::{Governor, GovernorConfigBuilder};
 use actix_web::{
     http::header::{HeaderName, CONTENT_SECURITY_POLICY},
     middleware::{self, DefaultHeaders, Logger},
@@ -52,9 +55,19 @@ async fn main() -> std::io::Result<()> {
     template::write_prefillable_templates(&config).await;
     let config = Data::new(config);
 
+    let governor_conf = GovernorConfigBuilder::default()
+        .per_second(config.rate_limit_per_second)
+        .burst_size(config.rate_limit_burst)
+        .key_extractor(ForwardedPeerIpKeyExtractor {
+            proxied: config.proxied,
+        })
+        .use_headers()
+        .finish()
+        .unwrap();
+
     HttpServer::new({
         move || {
-            App::new()
+            let app = App::new()
                 .wrap(Logger::new(r#"%{r}a "%r" =%s %bbytes %Tsec"#))
                 .wrap(DefaultHeaders::new().add(DEFAULT_CSP))
                 .wrap(middleware::Compress::default())
@@ -68,7 +81,19 @@ async fn main() -> std::io::Result<()> {
                         .route(web::get().to(upload::uploaded)),
                 )
                 .service(Files::new("/static", "static").disable_content_disposition())
-                .service(
+                .default_service(web::route().to(not_found));
+            if config.enable_rate_limit {
+                app.service(
+                    web::resource([
+                        "/{id:[a-z0-9]{5}}",
+                        "/{id:[a-z0-9]{5}}/",
+                        "/{id:[a-z0-9]{5}}/{name}",
+                    ])
+                    .wrap(Governor::new(&governor_conf))
+                    .route(web::get().to(download::download)),
+                )
+            } else {
+                app.service(
                     web::resource([
                         "/{id:[a-z0-9]{5}}",
                         "/{id:[a-z0-9]{5}}/",
@@ -76,7 +101,7 @@ async fn main() -> std::io::Result<()> {
                     ])
                     .route(web::get().to(download::download)),
                 )
-                .default_service(web::route().to(not_found))
+            }
         }
     })
     .bind(bind_address)?

src/rate_limit.rs

@@ -0,0 +1,51 @@
+use actix_governor::KeyExtractor;
+use actix_governor::PeerIpKeyExtractor;
+use actix_web::{dev::ServiceRequest, http::header::ContentType};
+use governor::clock::{Clock, DefaultClock, QuantaInstant};
+use governor::NotUntil;
+use std::net::IpAddr;
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub struct ForwardedPeerIpKeyExtractor {
+    pub proxied: bool,
+}
+
+impl KeyExtractor for ForwardedPeerIpKeyExtractor {
+    type Key = IpAddr;
+    type KeyExtractionError = &'static str;
+
+    #[cfg(feature = "log")]
+    fn name(&self) -> &'static str {
+        "Forwarded peer IP"
+    }
+
+    fn extract(&self, req: &ServiceRequest) -> Result<Self::Key, Self::KeyExtractionError> {
+        let forwarded_for = req.headers().get("x-forwarded-for");
+        if !self.proxied && forwarded_for.is_some() {
+            let forwarded_for = forwarded_for
+                .unwrap()
+                .to_str()
+                .map_err(|_| "x-forwarded-for contains invalid header value")?;
+            forwarded_for
+                .parse::<IpAddr>()
+                .map_err(|_| "x-forwarded-for contains invalid ip adress")
+        } else {
+            PeerIpKeyExtractor.extract(req)
+        }
+    }
+
+    fn response_error_content(&self, negative: &NotUntil<QuantaInstant>) -> (String, ContentType) {
+        let wait_time = negative
+            .wait_time_from(DefaultClock::default().now())
+            .as_secs();
+        (
+            format!("too many requests, retry in {}s", wait_time),
+            ContentType::plaintext(),
+        )
+    }
+
+    #[cfg(feature = "log")]
+    fn key_name(&self, key: &Self::Key) -> Option<String> {
+        Some(key.to_string())
+    }
+}

src/upload.rs

@@ -17,8 +17,8 @@ const UPLOAD_HTML: &str = include_str!("../template/upload.html");
 const UPLOAD_SHORT_HTML: &str = include_str!("../template/upload-short.html");
 
 const ID_CHARS: &[char] = &[
-    'a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'j', 'k', 'm', 'n', 'p', 'q', 'r', 's', 't', 'u', 'v',
+    'a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j', 'k', 'm', 'n', 'p', 'q', 'r', 's', 't', 'u',
-    'w', 'x', 'y', 'z', '1', '2', '3', '4', '5', '6', '7', '8', '9',
+    'v', 'w', 'x', 'y', 'z', '1', '2', '3', '4', '5', '6', '7', '8', '9',
 ];
 
 pub async fn index(config: web::Data<Config>) -> Result<NamedFile, Error> {