Add ability to configure HTTP endpoint limits (#4670)

src/cnf/mod.rs

@@ -28,6 +28,34 @@ pub const PKG_NAME: &str = "surrealdb";
 /// The public endpoint for the administration interface
 pub const APP_ENDPOINT: &str = "https://surrealdb.com/app";
 
+/// The maximum HTTP body size of the HTTP /ml endpoints (defaults to 4 GiB)
+pub static HTTP_MAX_ML_BODY_SIZE: Lazy<usize> =
+	lazy_env_parse!("SURREAL_HTTP_MAX_ML_BODY_SIZE", usize, 4 << 30);
+
+/// The maximum HTTP body size of the HTTP /sql endpoint (defaults to 1 MiB)
+pub static HTTP_MAX_SQL_BODY_SIZE: Lazy<usize> =
+	lazy_env_parse!("SURREAL_HTTP_MAX_SQL_BODY_SIZE", usize, 1 << 20);
+
+/// The maximum HTTP body size of the HTTP /rpc endpoint (defaults to 4 MiB)
+pub static HTTP_MAX_RPC_BODY_SIZE: Lazy<usize> =
+	lazy_env_parse!("SURREAL_HTTP_MAX_RPC_BODY_SIZE", usize, 4 << 20);
+
+/// The maximum HTTP body size of the HTTP /key endpoints (defaults to 16 KiB)
+pub static HTTP_MAX_KEY_BODY_SIZE: Lazy<usize> =
+	lazy_env_parse!("SURREAL_HTTP_MAX_KEY_BODY_SIZE", usize, 16 << 10);
+
+/// The maximum HTTP body size of the HTTP /signup endpoint (defaults to 1 KiB)
+pub static HTTP_MAX_SIGNUP_BODY_SIZE: Lazy<usize> =
+	lazy_env_parse!("SURREAL_HTTP_MAX_SIGNUP_BODY_SIZE", usize, 1 << 10);
+
+/// The maximum HTTP body size of the HTTP /key endpoints (defaults to 1 KiB)
+pub static HTTP_MAX_SIGNIN_BODY_SIZE: Lazy<usize> =
+	lazy_env_parse!("SURREAL_HTTP_MAX_SIGNIN_BODY_SIZE", usize, 1 << 10);
+
+/// The maximum HTTP body size of the HTTP /key endpoints (defaults to 4 GiB)
+pub static HTTP_MAX_IMPORT_BODY_SIZE: Lazy<usize> =
+	lazy_env_parse!("SURREAL_HTTP_MAX_IMPORT_BODY_SIZE", usize, 4 << 30);
+
 /// Specifies the frequency with which ping messages should be sent to the client
 pub const WEBSOCKET_PING_FREQUENCY: Duration = Duration::from_secs(5);
 

src/net/import.rs

@@ -1,5 +1,6 @@
 use super::headers::Accept;
 use super::AppState;
+use crate::cnf::HTTP_MAX_IMPORT_BODY_SIZE;
 use crate::err::Error;
 use crate::net::input::bytes_to_utf8;
 use crate::net::output;
@@ -15,8 +16,6 @@ use surrealdb::iam::Action::Edit;
 use surrealdb::iam::ResourceKind::Any;
 use tower_http::limit::RequestBodyLimitLayer;
 
-const MAX: usize = 1024 * 1024 * 1024 * 4; // 4 GiB
-
 pub(super) fn router<S>() -> Router<S>
 where
 	S: Clone + Send + Sync + 'static,
@@ -24,7 +23,7 @@ where
 	Router::new()
 		.route("/import", post(handler))
 		.route_layer(DefaultBodyLimit::disable())
-		.layer(RequestBodyLimitLayer::new(MAX))
+		.layer(RequestBodyLimitLayer::new(*HTTP_MAX_IMPORT_BODY_SIZE))
 }
 
 async fn handler(

src/net/key.rs

@@ -1,3 +1,6 @@
+use super::headers::Accept;
+use super::AppState;
+use crate::cnf::HTTP_MAX_KEY_BODY_SIZE;
 use crate::err::Error;
 use crate::net::input::bytes_to_utf8;
 use crate::net::output;
@@ -5,7 +8,8 @@ use crate::net::params::Params;
 use axum::extract::{DefaultBodyLimit, Path};
 use axum::response::IntoResponse;
 use axum::routing::options;
-use axum::{Extension, Router};
+use axum::Extension;
+use axum::Router;
 use axum_extra::extract::Query;
 use axum_extra::TypedHeader;
 use bytes::Bytes;
@@ -16,11 +20,6 @@ use surrealdb::iam::check::check_ns_db;
 use surrealdb::sql::Value;
 use tower_http::limit::RequestBodyLimitLayer;
 
-use super::headers::Accept;
-use super::AppState;
-
-const MAX: usize = 1024 * 16; // 16 KiB
-
 #[derive(Default, Deserialize, Debug, Clone)]
 struct QueryOptions {
 	pub limit: Option<i64>,
@@ -43,7 +42,7 @@ where
 				.delete(delete_all),
 		)
 		.route_layer(DefaultBodyLimit::disable())
-		.layer(RequestBodyLimitLayer::new(MAX))
+		.layer(RequestBodyLimitLayer::new(*HTTP_MAX_KEY_BODY_SIZE))
 		.merge(
 			Router::new()
 				.route(
@@ -56,7 +55,7 @@ where
 						.delete(delete_one),
 				)
 				.route_layer(DefaultBodyLimit::disable())
-				.layer(RequestBodyLimitLayer::new(MAX)),
+				.layer(RequestBodyLimitLayer::new(*HTTP_MAX_KEY_BODY_SIZE)),
 		)
 }
 

src/net/ml.rs

@@ -1,28 +1,38 @@
 //! This file defines the endpoints for the ML API for importing and exporting SurrealML models.
 use super::AppState;
+use crate::cnf::HTTP_MAX_ML_BODY_SIZE;
 use crate::err::Error;
+#[cfg(feature = "ml")]
 use crate::net::output;
 use axum::body::Body;
 use axum::extract::{DefaultBodyLimit, Path};
 use axum::response::IntoResponse;
+#[cfg(feature = "ml")]
 use axum::response::Response;
 use axum::routing::{get, post};
 use axum::Extension;
 use axum::Router;
+#[cfg(feature = "ml")]
 use bytes::Bytes;
+#[cfg(feature = "ml")]
 use futures_util::StreamExt;
+#[cfg(feature = "ml")]
 use http::StatusCode;
 use surrealdb::dbs::Session;
+#[cfg(feature = "ml")]
 use surrealdb::iam::check::check_ns_db;
+#[cfg(feature = "ml")]
 use surrealdb::iam::Action::{Edit, View};
+#[cfg(feature = "ml")]
 use surrealdb::iam::ResourceKind::Model;
+#[cfg(feature = "ml")]
 use surrealdb::kvs::{LockType::Optimistic, TransactionType::Read};
+#[cfg(feature = "ml")]
 use surrealdb::ml::storage::surml_file::SurMlFile;
+#[cfg(feature = "ml")]
 use surrealdb::sql::statements::{DefineModelStatement, DefineStatement};
 use tower_http::limit::RequestBodyLimitLayer;
 
-const MAX: usize = 1024 * 1024 * 1024 * 4; // 4 GiB
-
 /// The router definition for the ML API endpoints.
 pub(super) fn router<S>() -> Router<S>
 where
@@ -32,10 +42,11 @@ where
 		.route("/ml/import", post(import))
 		.route("/ml/export/:name/:version", get(export))
 		.route_layer(DefaultBodyLimit::disable())
-		.layer(RequestBodyLimitLayer::new(MAX))
+		.layer(RequestBodyLimitLayer::new(*HTTP_MAX_ML_BODY_SIZE))
 }
 
 /// This endpoint allows the user to import a model into the database.
+#[cfg(feature = "ml")]
 async fn import(
 	Extension(state): Extension<AppState>,
 	Extension(session): Extension<Session>,
@@ -89,6 +100,7 @@ async fn import(
 }
 
 /// This endpoint allows the user to export a model from the database.
+#[cfg(feature = "ml")]
 async fn export(
 	Extension(state): Extension<AppState>,
 	Extension(session): Extension<Session>,
@@ -120,3 +132,23 @@ async fn export(
 	// Return the streamed body
 	Ok(Response::builder().status(StatusCode::OK).body(body).unwrap())
 }
+
+/// This endpoint allows the user to import a model into the database.
+#[cfg(not(feature = "ml"))]
+async fn import(
+	Extension(_): Extension<AppState>,
+	Extension(_): Extension<Session>,
+	_: Body,
+) -> Result<(), impl IntoResponse> {
+	Err(Error::Request)
+}
+
+/// This endpoint allows the user to export a model from the database.
+#[cfg(not(feature = "ml"))]
+async fn export(
+	Extension(_): Extension<AppState>,
+	Extension(_): Extension<Session>,
+	Path((_, _)): Path<(String, String)>,
+) -> Result<(), impl IntoResponse> {
+	Err(Error::Request)
+}

src/net/mod.rs

@@ -8,6 +8,7 @@ mod health;
 mod import;
 mod input;
 mod key;
+mod ml;
 pub(crate) mod output;
 mod params;
 mod rpc;
@@ -19,9 +20,6 @@ mod sync;
 mod tracer;
 mod version;
 
-#[cfg(feature = "ml")]
-mod ml;
-
 use crate::cli::CF;
 use crate::cnf::{self, GRAPHQL_ENABLE};
 use crate::err::Error;
@@ -173,7 +171,8 @@ pub async fn init(ds: Arc<Datastore>, ct: CancellationToken) -> Result<(), Error
 		.merge(sql::router())
 		.merge(signin::router())
 		.merge(signup::router())
-		.merge(key::router());
+		.merge(key::router())
+		.merge(ml::router());
 
 	let axum_app = if *GRAPHQL_ENABLE {
 		#[cfg(surrealdb_unstable)]
@@ -190,9 +189,6 @@ pub async fn init(ds: Arc<Datastore>, ct: CancellationToken) -> Result<(), Error
 		axum_app
 	};
 
-	#[cfg(feature = "ml")]
-	let axum_app = axum_app.merge(ml::router());
-
 	let axum_app = axum_app.layer(service);
 
 	// Get a new server handler

src/net/rpc.rs

@@ -4,15 +4,16 @@ use std::sync::Arc;
 
 use super::headers::SurrealId;
 use crate::cnf;
+use crate::cnf::HTTP_MAX_RPC_BODY_SIZE;
 use crate::err::Error;
 use crate::rpc::connection::Connection;
 use crate::rpc::format::HttpFormat;
 use crate::rpc::post_context::PostRpcContext;
 use crate::rpc::response::IntoRpcResponse;
 use crate::rpc::RpcState;
+use axum::extract::DefaultBodyLimit;
 use axum::extract::State;
-use axum::routing::get;
-use axum::routing::post;
+use axum::routing::options;
 use axum::{
 	extract::ws::{WebSocket, WebSocketUpgrade},
 	response::IntoResponse,
@@ -28,6 +29,7 @@ use surrealdb::kvs::Datastore;
 use surrealdb::rpc::format::Format;
 use surrealdb::rpc::format::PROTOCOLS;
 use surrealdb::rpc::method::Method;
+use tower_http::limit::RequestBodyLimitLayer;
 use tower_http::request_id::RequestId;
 use uuid::Uuid;
 
@@ -38,7 +40,10 @@ use super::AppState;
 use surrealdb::rpc::rpc_context::RpcContext;
 
 pub(super) fn router() -> Router<Arc<RpcState>> {
-	Router::new().route("/rpc", get(get_handler)).route("/rpc", post(post_handler))
+	Router::new()
+		.route("/rpc", options(|| async {}).get(get_handler).post(post_handler))
+		.route_layer(DefaultBodyLimit::disable())
+		.layer(RequestBodyLimitLayer::new(*HTTP_MAX_RPC_BODY_SIZE))
 }
 
 async fn get_handler(

src/net/signin.rs

@@ -1,3 +1,6 @@
+use super::headers::Accept;
+use super::AppState;
+use crate::cnf::HTTP_MAX_SIGNUP_BODY_SIZE;
 use crate::err::Error;
 use crate::net::input::bytes_to_utf8;
 use crate::net::output;
@@ -13,11 +16,6 @@ use surrealdb::dbs::Session;
 use surrealdb::sql::Value;
 use tower_http::limit::RequestBodyLimitLayer;
 
-use super::headers::Accept;
-use super::AppState;
-
-const MAX: usize = 1024; // 1 KiB
-
 #[derive(Serialize)]
 struct Success {
 	code: u16,
@@ -42,7 +40,7 @@ where
 	Router::new()
 		.route("/signin", options(|| async {}).post(handler))
 		.route_layer(DefaultBodyLimit::disable())
-		.layer(RequestBodyLimitLayer::new(MAX))
+		.layer(RequestBodyLimitLayer::new(*HTTP_MAX_SIGNUP_BODY_SIZE))
 }
 
 async fn handler(

src/net/signup.rs

@@ -1,10 +1,14 @@
+use super::headers::Accept;
+use super::AppState;
+use crate::cnf::HTTP_MAX_SIGNIN_BODY_SIZE;
 use crate::err::Error;
 use crate::net::input::bytes_to_utf8;
 use crate::net::output;
 use axum::extract::DefaultBodyLimit;
 use axum::response::IntoResponse;
 use axum::routing::options;
-use axum::{Extension, Router};
+use axum::Extension;
+use axum::Router;
 use axum_extra::TypedHeader;
 use bytes::Bytes;
 use serde::Serialize;
@@ -12,11 +16,6 @@ use surrealdb::dbs::Session;
 use surrealdb::sql::Value;
 use tower_http::limit::RequestBodyLimitLayer;
 
-use super::headers::Accept;
-use super::AppState;
-
-const MAX: usize = 1024; // 1 KiB
-
 #[derive(Serialize)]
 struct Success {
 	code: u16,
@@ -41,7 +40,7 @@ where
 	Router::new()
 		.route("/signup", options(|| async {}).post(handler))
 		.route_layer(DefaultBodyLimit::disable())
-		.layer(RequestBodyLimitLayer::new(MAX))
+		.layer(RequestBodyLimitLayer::new(*HTTP_MAX_SIGNIN_BODY_SIZE))
 }
 
 async fn handler(

src/net/sql.rs

@@ -1,3 +1,6 @@
+use super::headers::Accept;
+use super::AppState;
+use crate::cnf::HTTP_MAX_SQL_BODY_SIZE;
 use crate::err::Error;
 use crate::net::input::bytes_to_utf8;
 use crate::net::output;
@@ -17,11 +20,6 @@ use futures::{SinkExt, StreamExt};
 use surrealdb::dbs::Session;
 use tower_http::limit::RequestBodyLimitLayer;
 
-use super::headers::Accept;
-use super::AppState;
-
-const MAX: usize = 1024 * 1024; // 1 MiB
-
 pub(super) fn router<S>() -> Router<S>
 where
 	S: Clone + Send + Sync + 'static,
@@ -29,7 +27,7 @@ where
 	Router::new()
 		.route("/sql", options(|| async {}).get(ws_handler).post(post_handler))
 		.route_layer(DefaultBodyLimit::disable())
-		.layer(RequestBodyLimitLayer::new(MAX))
+		.layer(RequestBodyLimitLayer::new(*HTTP_MAX_SQL_BODY_SIZE))
 }
 
 async fn post_handler(