From a4750f3241f46d047da8a5519a2305a5a3bd9159 Mon Sep 17 00:00:00 2001
From: Tom Harding <i.am.tom.harding@gmail.com>
Date: Tue, 20 Aug 2024 14:42:41 +0200
Subject: [PATCH] Add a request size limit (#29)

There is a current issue in `sqlx` being investigated:
https://github.com/launchbadge/sqlx/issues/3440

In Engine, we already have request size limits so this isn't an issue,
but no such limits exist in the connectors. This PR adds a 100MB limit
to connector requests, in order to avoid problems if bad actors find
ways to query NDCs directly.
---
 crates/sdk/Cargo.toml               |  2 +-
 crates/sdk/src/connector/example.rs |  2 +-
 crates/sdk/src/default_main.rs      | 14 +++++++++++++-
 3 files changed, 15 insertions(+), 3 deletions(-)

diff --git a/crates/sdk/Cargo.toml b/crates/sdk/Cargo.toml
index e41be11..c9b09b8 100644
--- a/crates/sdk/Cargo.toml
+++ b/crates/sdk/Cargo.toml
@@ -46,7 +46,7 @@ serde = { version = "1.0.197", features = ["derive"] }
 serde_json = { version = "1.0.115", features = ["raw_value"] }
 thiserror = "1.0"
 tokio = { version = "1.36.0", features = ["fs", "macros", "rt-multi-thread", "signal"] }
-tower-http = { version = "0.4.4", features = ["cors", "trace", "validate-request"] }
+tower-http = { version = "0.4.4", features = ["cors", "limit", "trace", "validate-request"] }
 tracing = "0.1.40"
 tracing-opentelemetry = "0.23.0"
 tracing-subscriber = { version = "0.3", default-features = false, features = ["ansi", "env-filter", "fmt", "json"] }
diff --git a/crates/sdk/src/connector/example.rs b/crates/sdk/src/connector/example.rs
index 505634e..0bfe0bc 100644
--- a/crates/sdk/src/connector/example.rs
+++ b/crates/sdk/src/connector/example.rs
@@ -126,7 +126,7 @@ mod tests {
     async fn capabilities_match_ndc_spec_version() -> Result<()> {
         let state =
             crate::default_main::init_server_state(Example::default(), PathBuf::new()).await?;
-        let app = crate::default_main::create_router::<Example>(state, None);
+        let app = crate::default_main::create_router::<Example>(state, None, None);
 
         let client = TestClient::new(app);
         let response = client.get("/capabilities").send().await;
diff --git a/crates/sdk/src/default_main.rs b/crates/sdk/src/default_main.rs
index c5277e2..27b1173 100644
--- a/crates/sdk/src/default_main.rs
+++ b/crates/sdk/src/default_main.rs
@@ -12,7 +12,9 @@ use axum::{
 use axum_extra::extract::WithRejection;
 use clap::{Parser, Subcommand};
 use prometheus::Registry;
-use tower_http::{trace::TraceLayer, validate_request::ValidateRequestHeaderLayer};
+use tower_http::{
+    limit::RequestBodyLimitLayer, trace::TraceLayer, validate_request::ValidateRequestHeaderLayer,
+};
 
 use ndc_models::{
     CapabilitiesResponse, ExplainResponse, MutationRequest, MutationResponse, QueryRequest,
@@ -74,6 +76,8 @@ struct ServeCommand {
     service_token_secret: Option<String>,
     #[arg(long, value_name = "NAME", env = "OTEL_SERVICE_NAME")]
     service_name: Option<String>,
+    #[arg(long, value_name = "MAX_REQUEST_SIZE", env = "HASURA_MAX_REQUEST_SIZE")]
+    max_request_size: Option<usize>,
 }
 
 #[derive(Clone, Parser)]
@@ -242,6 +246,7 @@ where
     let router = create_router::<Setup::Connector>(
         server_state.clone(),
         serve_command.service_token_secret.clone(),
+        serve_command.max_request_size,
     );
 
     let address = net::SocketAddr::new(serve_command.host, serve_command.port);
@@ -296,6 +301,7 @@ pub async fn init_server_state<Setup: ConnectorSetup>(
 pub fn create_router<C>(
     state: ServerState<C>,
     service_token_secret: Option<String>,
+    max_request_size: Option<usize>,
 ) -> axum::Router<()>
 where
     C: Connector + 'static,
@@ -310,6 +316,12 @@ where
         .route("/query/explain", post(post_query_explain::<C>))
         .route("/mutation", post(post_mutation::<C>))
         .route("/mutation/explain", post(post_mutation_explain::<C>))
+        // We want to limit the size of requests to 100MB to prevent various DDoS / SQL overflow
+        // vulnerabilities. We use RequestBodyLimit instead of DefaultBodyLimit to include chunked
+        // requests, too.
+        .layer(RequestBodyLimitLayer::new(
+            max_request_size.unwrap_or(100 * 1024 * 1024),
+        ))
         .layer(ValidateRequestHeaderLayer::custom(auth_handler(
             service_token_secret,
         )))