Request size limiter

zealsham · zealsham · commit d3bcd0a31842 · 2026-03-19T00:51:12.000+01:00
This pr addresses #941, it implements a body size limit layer to reject request whose size is greater than 7168. This also ensures both v1 and v2 have the same request size which makes it hard for an attacker to tell if a request is v1/v2 baed on payload.
diff --git a/Cargo-minimal.lock b/Cargo-minimal.lock
@@ -2872,6 +2872,7 @@ dependencies = [
  "tokio-stream",
  "tokio-tungstenite",
  "tower",
+ "tower-http",
  "tracing",
  "tracing-subscriber",
  "uuid",
diff --git a/Cargo-recent.lock b/Cargo-recent.lock
@@ -2872,6 +2872,7 @@ dependencies = [
  "tokio-stream",
  "tokio-tungstenite",
  "tower",
+ "tower-http",
  "tracing",
  "tracing-subscriber",
  "uuid",
diff --git a/payjoin-mailroom/Cargo.toml b/payjoin-mailroom/Cargo.toml
@@ -72,6 +72,7 @@ tokio-rustls-acme = { version = "0.9.0", features = ["axum"], optional = true }
 tokio-stream = { version = "0.1.17", features = ["net"] }
 tokio-tungstenite = { version = "0.27.0", optional = true }
 tower = "0.5"
+tower-http = { version = "0.6", features = ["limit"] }
 tracing = "0.1"
 tracing-subscriber = { version = "0.3", features = ["env-filter", "json"] }
 
diff --git a/payjoin-mailroom/src/directory.rs b/payjoin-mailroom/src/directory.rs
@@ -18,7 +18,7 @@ const CHACHA20_POLY1305_NONCE_LEN: usize = 32; // chacha20poly1305 n_k
 const POLY1305_TAG_SIZE: usize = 16;
 pub const BHTTP_REQ_BYTES: usize =
     ENCAPSULATED_MESSAGE_BYTES - (CHACHA20_POLY1305_NONCE_LEN + POLY1305_TAG_SIZE);
-const V1_MAX_BUFFER_SIZE: usize = 65536;
+pub const V1_MAX_BUFFER_SIZE: usize = 7168;
 
 const V1_REJECT_RES_JSON: &str =
     r#"{{"errorCode": "original-psbt-rejected ", "message": "Body is not a string"}}"#;
@@ -276,9 +276,7 @@ impl<D: Db> Service<D> {
             .await
             .map_err(|e| HandlerError::InternalServerError(e.into()))?
             .to_bytes();
-        if req.len() > V1_MAX_BUFFER_SIZE {
-            return Err(HandlerError::PayloadTooLarge);
-        }
+
         match self.db.post_v2_payload(&id, req.into()).await {
             Ok(_) => Ok(none_response),
             Err(e) => Err(HandlerError::InternalServerError(e.into())),
@@ -322,9 +320,6 @@ impl<D: Db> Service<D> {
             .await
             .map_err(|e| HandlerError::InternalServerError(e.into()))?
             .to_bytes();
-        if req.len() > V1_MAX_BUFFER_SIZE {
-            return Err(HandlerError::PayloadTooLarge);
-        }
 
         let body_str = std::str::from_utf8(&req).map_err(|e| HandlerError::BadRequest(e.into()))?;
         self.check_v1_blocklist(body_str).await?;
@@ -455,7 +450,6 @@ async fn handle_directory_home_path() -> Result<Response<Body>, HandlerError> {
 
 #[derive(Debug)]
 enum HandlerError {
-    PayloadTooLarge,
     InternalServerError(anyhow::Error),
     ServiceUnavailable(anyhow::Error),
     SenderGone(anyhow::Error),
@@ -470,11 +464,11 @@ impl HandlerError {
     fn to_response(&self) -> Response<Body> {
         let mut res = Response::new(empty());
         match self {
-            HandlerError::PayloadTooLarge => *res.status_mut() = StatusCode::PAYLOAD_TOO_LARGE,
             HandlerError::InternalServerError(e) => {
                 error!("Internal server error: {}", e);
                 *res.status_mut() = StatusCode::INTERNAL_SERVER_ERROR
             }
+
             HandlerError::ServiceUnavailable(e) => {
                 error!("Service temporarily unavailable: {}", e);
                 *res.status_mut() = StatusCode::SERVICE_UNAVAILABLE
diff --git a/payjoin-mailroom/src/lib.rs b/payjoin-mailroom/src/lib.rs
@@ -11,6 +11,7 @@ use opentelemetry_sdk::metrics::SdkMeterProvider;
 use rand::Rng;
 use tokio_listener::{Listener, SystemOptions, UserOptions};
 use tower::{Service, ServiceBuilder};
+use tower_http::limit::RequestBodyLimitLayer;
 use tracing::info;
 
 use crate::ohttp_relay::SentinelTag;
@@ -360,7 +361,8 @@ fn build_app(services: Services) -> Router {
         .layer(
             ServiceBuilder::new()
                 .layer(axum::middleware::from_fn_with_state(metrics.clone(), track_metrics))
-                .layer(axum::middleware::from_fn_with_state(metrics, track_connections)),
+                .layer(axum::middleware::from_fn_with_state(metrics, track_connections))
+                .layer(RequestBodyLimitLayer::new(crate::directory::V1_MAX_BUFFER_SIZE)),
         )
         .with_state(services);
 
@@ -565,4 +567,97 @@ mod tests {
         assert!(metric_names.contains(&TOTAL_CONNECTIONS), "missing total_connections");
         assert!(metric_names.contains(&ACTIVE_CONNECTIONS), "missing active_connections");
     }
+
+    /// Exercises the full middleware stack (track_metrics, track_connections, RequestBodyLimit)
+    /// with a request that has a body, ensuring the layer ordering compiles and runs correctly.
+    #[tokio::test]
+    async fn middleware_stack_accepts_request_with_body() {
+        use axum::body::Body;
+        use axum::http::header::CONTENT_TYPE;
+        use axum::http::Request;
+        use tower::ServiceExt;
+
+        let tempdir = tempdir().unwrap();
+        let config = Config::new(
+            "[::]:0".parse().expect("valid listener address"),
+            tempdir.path().to_path_buf(),
+            Duration::from_secs(2),
+            None,
+        );
+
+        let sentinel_tag = generate_sentinel_tag();
+        let services = Services {
+            directory: init_directory(&config, sentinel_tag).await.unwrap(),
+            relay: ohttp_relay::Service::new(sentinel_tag).await,
+            metrics: MetricsService::new(None),
+            #[cfg(feature = "access-control")]
+            geoip: None,
+        };
+
+        let app = build_app(services);
+
+        let body = b"small payload under 65k limit";
+        let request = Request::builder()
+            .method("POST")
+            .uri("/some-path")
+            .header(CONTENT_TYPE, ohttp_relay::EXPECTED_MEDIA_TYPE.clone())
+            .body(Body::from(body.as_slice()))
+            .unwrap();
+
+        let response = ServiceExt::<Request<Body>>::oneshot(app, request).await.unwrap();
+
+        assert!(
+            !response.status().is_server_error(),
+            "middleware stack should accept request with body (no 5xx)"
+        );
+        assert_ne!(
+            response.status(),
+            axum::http::StatusCode::PAYLOAD_TOO_LARGE,
+            "small body should not be rejected by body limit"
+        );
+    }
+
+    /// Ensures RequestBodyLimitLayer (65_536 bytes) rejects oversized bodies (413 or 415).
+    #[tokio::test]
+    async fn request_body_limit_rejects_oversized_body() {
+        use axum::body::Body;
+        use axum::http::header::CONTENT_TYPE;
+        use axum::http::Request;
+        use tower::ServiceExt;
+
+        let tempdir = tempdir().unwrap();
+        let config = Config::new(
+            "[::]:0".parse().expect("valid listener address"),
+            tempdir.path().to_path_buf(),
+            Duration::from_secs(2),
+            None,
+        );
+
+        let sentinel_tag = generate_sentinel_tag();
+        let services = Services {
+            directory: init_directory(&config, sentinel_tag).await.unwrap(),
+            relay: ohttp_relay::Service::new(sentinel_tag).await,
+            metrics: MetricsService::new(None),
+            #[cfg(feature = "access-control")]
+            geoip: None,
+        };
+
+        let app = build_app(services);
+
+        let oversized: Vec<u8> = (0..65_537).map(|_| 0u8).collect();
+        let request = Request::builder()
+            .method("POST")
+            .uri("/")
+            .header(CONTENT_TYPE, ohttp_relay::EXPECTED_MEDIA_TYPE.clone())
+            .body(Body::from(oversized))
+            .unwrap();
+
+        let response = ServiceExt::<Request<Body>>::oneshot(app, request).await.unwrap();
+
+        assert!(
+            response.status().is_client_error(),
+            "body over 65_536 bytes should be rejected with a 4xx status, got {}",
+            response.status()
+        );
+    }
 }
