Parse TTL and max-bytes parameters in a human-friendly way

Author: sandhose

server/Cargo.toml

@@ -9,10 +9,12 @@ repository = "https://github.com/matrix-org/rust-http-rendezvous-server/"
 rust-version = "1.61"
 
 [dependencies]
+bytesize = "1.1.0"
+clap = { version = "4.0.14", features = ["derive"] }
+humantime = "2.1.0"
 hyper = { version = "0.14.20", features = ["server"] }
 tokio = { version = "1.21.2", features = ["macros", "rt-multi-thread"] }
 tracing = "0.1.37"
 tracing-subscriber = "0.3.16"
-clap = { version = "4.0.14", features = ["derive"] }
 
 matrix-http-rendezvous = { path = "../", version = "0.1.3" }

server/src/main.rs

@@ -16,6 +16,7 @@
 #![deny(clippy::all)]
 #![warn(clippy::pedantic)]
 
+use bytesize::ByteSize;
 use clap::Parser;
 use std::{
     net::{IpAddr, Ipv4Addr, SocketAddr},
@@ -37,12 +38,12 @@ struct Options {
     prefix: Option<String>,
 
     /// Time to live of entries, in seconds
-    #[arg(short, long, default_value_t = 60)]
-    ttl: u64,
+    #[arg(short, long, default_value_t = Duration::from_secs(60).into())]
+    ttl: humantime::Duration,
 
     /// Maximum payload size, in bytes
-    #[arg(short, long, default_value_t = 4096)]
-    max_bytes: usize,
+    #[arg(short, long, default_value = "4KiB")]
+    max_bytes: ByteSize,
 }
 
 #[tokio::main]
@@ -51,17 +52,22 @@ async fn main() {
 
     let options = Options::parse();
     let prefix = options.prefix.unwrap_or_default();
-    let ttl = Duration::from_secs(options.ttl);
+    let ttl = options.ttl.into();
+    let max_bytes = options
+        .max_bytes
+        .0
+        .try_into()
+        .expect("Max bytes size too large");
 
     let addr = SocketAddr::from((options.address, options.port));
 
-    let service = matrix_http_rendezvous::router(&prefix, ttl, options.max_bytes);
+    let service = matrix_http_rendezvous::router(&prefix, ttl, max_bytes);
 
     tracing::info!("Listening on http://{addr}");
     tracing::info!(
-        "TTL: {ttl}s – Maximum payload size: {max_bytes} bytes",
-        ttl = ttl.as_secs(),
-        max_bytes = options.max_bytes
+        "TTL: {ttl} – Maximum payload size: {max_bytes}",
+        ttl = humantime::format_duration(ttl),
+        max_bytes = options.max_bytes.to_string_as(true)
     );
 
     hyper::Server::bind(&addr)

src/lib.rs

@@ -26,7 +26,7 @@ use std::{
 
 use axum::{
     body::HttpBody,
-    extract::{FromRef, Path, State},
+    extract::{DefaultBodyLimit, FromRef, Path, State},
     http::{
         header::{CONTENT_TYPE, ETAG, IF_MATCH, IF_NONE_MATCH, LOCATION},
         StatusCode,
@@ -253,6 +253,7 @@ where
 
     Router::new()
         .nest(prefix, router)
+        .layer(DefaultBodyLimit::disable())
         .layer(RequestBodyLimitLayer::new(max_bytes))
         .layer(SetResponseHeaderLayer::if_not_present(
             HeaderName::from_static("x-max-bytes"),
@@ -270,14 +271,71 @@ where
 
 #[cfg(test)]
 mod tests {
+    use std::convert::Infallible;
+
     use super::*;
 
     use axum::http::{
         header::{CONTENT_LENGTH, CONTENT_TYPE},
         Request,
     };
+    use bytes::Buf;
     use tower::util::ServiceExt;
 
+    /// A slow body, which sends the bytes in small chunks (1 byte per chunk by default)
+    #[derive(Clone)]
+    struct SlowBody {
+        body: Bytes,
+        chunk_size: usize,
+    }
+
+    impl SlowBody {
+        const fn from_static(bytes: &'static [u8]) -> Self {
+            Self {
+                body: Bytes::from_static(bytes),
+                chunk_size: 1,
+            }
+        }
+
+        const fn from_bytes(body: Bytes) -> Self {
+            Self {
+                body,
+                chunk_size: 1,
+            }
+        }
+
+        const fn with_chunk_size(mut self, chunk_size: usize) -> Self {
+            self.chunk_size = chunk_size;
+            self
+        }
+    }
+
+    impl HttpBody for SlowBody {
+        type Data = Bytes;
+        type Error = Infallible;
+
+        fn poll_data(
+            self: std::pin::Pin<&mut Self>,
+            _cx: &mut std::task::Context<'_>,
+        ) -> std::task::Poll<Option<Result<Self::Data, Self::Error>>> {
+            if self.body.is_empty() {
+                std::task::Poll::Ready(None)
+            } else {
+                let size = self.chunk_size.min(self.body.len());
+                let ret = self.body.slice(0..size);
+                self.get_mut().body.advance(size);
+                std::task::Poll::Ready(Some(Ok(ret)))
+            }
+        }
+
+        fn poll_trailers(
+            self: std::pin::Pin<&mut Self>,
+            _cx: &mut std::task::Context<'_>,
+        ) -> std::task::Poll<Result<Option<headers::HeaderMap>, Self::Error>> {
+            std::task::Poll::Ready(Ok(None))
+        }
+    }
+
     async fn advance_time(duration: Duration) {
         tokio::task::yield_now().await;
         tokio::time::pause();
@@ -325,6 +383,63 @@ mod tests {
         assert_eq!(response.status(), StatusCode::NOT_FOUND);
     }
 
+    #[tokio::test]
+    async fn test_post_max_bytes() {
+        let ttl = Duration::from_secs(60);
+
+        let body = br#"{"hello": "world"}"#;
+
+        // It doesn't work with a way too small size
+        let slow_body = SlowBody::from_static(body);
+        let request = Request::post("/")
+            .header(CONTENT_TYPE, "application/json")
+            .body(slow_body)
+            .unwrap();
+        let response = router("/", ttl, 8).oneshot(request).await.unwrap();
+        assert_eq!(response.status(), StatusCode::PAYLOAD_TOO_LARGE);
+
+        // It works with exactly the right size
+        let slow_body = SlowBody::from_static(body);
+        let request = Request::post("/")
+            .header(CONTENT_TYPE, "application/json")
+            .body(slow_body)
+            .unwrap();
+        let response = router("/", ttl, body.len()).oneshot(request).await.unwrap();
+        assert_eq!(response.status(), StatusCode::CREATED);
+
+        // It doesn't work even if the size is one too short
+        let slow_body = SlowBody::from_static(body);
+        let request = Request::post("/")
+            .header(CONTENT_TYPE, "application/json")
+            .body(slow_body)
+            .unwrap();
+        let response = router("/", ttl, body.len() - 1)
+            .oneshot(request)
+            .await
+            .unwrap();
+        assert_eq!(response.status(), StatusCode::PAYLOAD_TOO_LARGE);
+
+        // Try with a big body (4MB), sent in small 128 bytes chunks
+        let body = vec![42; 4 * 1024 * 1024].into_boxed_slice();
+        let slow_body = SlowBody::from_bytes(Bytes::from(body)).with_chunk_size(128);
+        let request = Request::post("/").body(slow_body).unwrap();
+        let response = router("/", ttl, 4 * 1024 * 1024)
+            .oneshot(request)
+            .await
+            .unwrap();
+        assert_eq!(response.status(), StatusCode::CREATED);
+
+        // Try with a big body (4MB + 1B), sent in small 128 bytes chunks
+        let body = vec![42; 4 * 1024 * 1024 + 1].into_boxed_slice();
+        let slow_body = SlowBody::from_bytes(Bytes::from(body)).with_chunk_size(128);
+        let request = Request::post("/").body(slow_body).unwrap();
+        let response = router("/", ttl, 4 * 1024 * 1024)
+            .oneshot(request)
+            .await
+            .unwrap();
+        assert_eq!(response.status(), StatusCode::PAYLOAD_TOO_LARGE);
+    }
+
     #[tokio::test]
     async fn test_post_and_get_if_none_match() {
         let ttl = Duration::from_secs(60);

synapse/Cargo.toml

@@ -15,7 +15,10 @@ crate-type = ["cdylib"]
 
 [dependencies]
 anyhow = "1.0.65"
+bytesize = { version = "1.1.0", features = ["serde"] }
 http-body = "0.4.5"
+humantime = "2.1.0"
+humantime-serde = "1.1.1"
 pyo3 = { version = "0.17.2", features = ["extension-module", "abi3-py37", "anyhow"] }
 pyo3-log = "0.7.0"
 pyo3-matrix-synapse-module = "0.1.0"

synapse/src/lib.rs

@@ -19,32 +19,33 @@
 
 use std::time::Duration;
 
-use anyhow::anyhow;
+use anyhow::{anyhow, Context};
+use bytesize::ByteSize;
 use http_body::Body;
 use pyo3::prelude::*;
 use serde::Deserialize;
 use tower::ServiceExt;
 
 use pyo3_matrix_synapse_module::{parse_config, ModuleApi};
 
-fn default_ttl() -> u64 {
-    60
+fn default_ttl() -> Duration {
+    Duration::from_secs(60)
 }
 
-fn default_max_bytes() -> usize {
-    4096
+fn default_max_bytes() -> ByteSize {
+    ByteSize::kib(4)
 }
 
 #[pyclass]
 #[derive(Deserialize)]
 struct Config {
     prefix: String,
 
-    #[serde(default = "default_ttl")]
-    ttl: u64,
+    #[serde(default = "default_ttl", with = "humantime_serde")]
+    ttl: Duration,
 
     #[serde(default = "default_max_bytes")]
-    max_bytes: usize,
+    max_bytes: ByteSize,
 }
 
 #[pyclass]
@@ -54,10 +55,21 @@ pub struct SynapseRendezvousModule;
 impl SynapseRendezvousModule {
     #[new]
     fn new(config: &Config, module_api: ModuleApi) -> PyResult<Self> {
-        let ttl = Duration::from_secs(config.ttl);
-        let service = matrix_http_rendezvous::router(&config.prefix, ttl, config.max_bytes)
-            .map_response(|res| res.map(|b| b.map_err(|e| anyhow!(e))));
+        tracing::info!(
+            "Mounting rendez-vous server on {prefix}, with a TTL of {ttl} and max payload size of {max_bytes}",
+            prefix = config.prefix,
+            ttl = humantime::format_duration(config.ttl),
+            max_bytes = config.max_bytes.to_string_as(true),
+        );
+
+        let max_bytes = config
+            .max_bytes
+            .0
+            .try_into()
+            .context("Could not convert max_bytes from config")?;
 
+        let service = matrix_http_rendezvous::router(&config.prefix, config.ttl, max_bytes)
+            .map_response(|res| res.map(|b| b.map_err(|e| anyhow!(e))));
         module_api.register_web_service(&config.prefix, service)?;
         Ok(Self)
     }