From 2b27e5456a8f7f54a4e01de22a5ba5fd2a8f1415 Mon Sep 17 00:00:00 2001 From: Or Ricon Date: Mon, 13 Jun 2022 11:53:10 +0300 Subject: [PATCH] limit request body size --- Cargo.lock | 5 +++-- Cargo.toml | 1 + src/main.rs | 21 ++++++++++++++++++++- 3 files changed, 24 insertions(+), 3 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 56d9a66..9c954cd 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -797,9 +797,9 @@ dependencies = [ [[package]] name = "http-body" -version = "0.4.4" +version = "0.4.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1ff4f84919677303da5f147645dbea6b1881f368d03ac84e1dc09031ebd7b2c6" +checksum = "d5f38f16d184e36f2408a55281cd658ecbd3ca05cce6d6510a176eca393e26d1" dependencies = [ "bytes", "http", @@ -961,6 +961,7 @@ dependencies = [ "futures", "garcon", "hex", + "http-body", "hyper", "hyper-rustls", "hyper-tls", diff --git a/Cargo.toml b/Cargo.toml index 4d6a183..8155cfb 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -26,6 +26,7 @@ flate2 = "1.0.0" futures = "0.3.21" garcon = { version = "0.2", features = ["async"] } hex = "0.4" +http-body = "0.4.5" hyper = { version = "0.14", features = ["full"] } hyper-rustls = { version = "0.23", features = [ "webpki-roots" ] } hyper-tls = "0.5" diff --git a/src/main.rs b/src/main.rs index 090f2a7..62393be 100644 --- a/src/main.rs +++ b/src/main.rs @@ -2,6 +2,7 @@ use axum::{handler::Handler, routing::get, Extension, Router}; use clap::{crate_authors, crate_version, Parser}; use flate2::read::{DeflateDecoder, GzDecoder}; use futures::{future::OptionFuture, try_join, FutureExt}; +use http_body::{LengthLimitError, Limited}; use hyper::{ body, body::Bytes, @@ -64,6 +65,11 @@ const MAX_LOG_CERT_B64_SIZE: usize = 2000; const MAX_CHUNK_SIZE_TO_DECOMPRESS: usize = 1024; const MAX_CHUNKS_TO_DECOMPRESS: u64 = 10_240; +const KB: usize = 1024; +const MB: usize = 1024 * KB; + +const REQUEST_BODY_SIZE_LIMIT: usize = 10 * MB; + /// Resolve overrides for [`reqwest::ClientBuilder::resolve()`] /// `ic0.app=[::1]:9090` pub(crate) struct OptResolve { @@ -301,7 +307,20 @@ async fn forward_request( }) .collect::>(); - let entire_body = body::to_bytes(body).await?.to_vec(); + // Limit request body size + let body = Limited::new(body, REQUEST_BODY_SIZE_LIMIT); + let entire_body = match hyper::body::to_bytes(body).await { + Ok(data) => data, + Err(err) => { + if err.downcast_ref::().is_some() { + return Ok(Response::builder() + .status(StatusCode::PAYLOAD_TOO_LARGE) + .body(Body::from("Request size exceeds limit"))?); + } + return Err(err); + } + } + .to_vec(); slog::trace!(logger, "<<"); if logger.is_trace_enabled() {