Limit request body sizes

spruceid · Oct 25, 2021 · cb95825 · cb95825
1 parent c13582d
commit cb95825
Showing 1 changed file with 55 additions and 9 deletions.
diff --git a/http/src/lib.rs b/http/src/lib.rs
@@ -22,14 +22,18 @@ pub mod error;
 use accept::HttpAccept;
 pub use error::Error;
 
-use hyper::body::Buf;
 use hyper::header::{ACCEPT, CONTENT_TYPE, LOCATION};
 use hyper::{Body, Response};
 use hyper::{Method, Request, StatusCode};
 use serde::{Deserialize, Serialize};
 use serde_json::json;
 use tower_service::Service;
 
+/// Maximum size (bytes) for request payloads.
+///
+/// Larger payloads will result in a HTTP 413 error.
+pub const MAX_BODY_LENGTH: usize = 2097152; // 2MB
+
 #[derive(Debug, Serialize, Deserialize, Clone)]
 #[serde(untagged)]
 pub enum PresentationOrJWT {
@@ -104,6 +108,28 @@ pub async fn pick_key<'a>(
     keys.get(&public_key)
 }
 
+async fn get_body_bytes_limited(req: Request<Body>) -> Result<Vec<u8>, (StatusCode, String)> {
+    use hyper::body::HttpBody;
+    let mut body = req.into_body();
+    let mut bytes = Vec::new();
+    while let Some(buf) = body.data().await {
+        let buf = match buf {
+            Ok(buf) => buf,
+            Err(err) => {
+                return Err((StatusCode::INTERNAL_SERVER_ERROR, err.to_string()));
+            }
+        };
+        if bytes.len() + buf.len() > MAX_BODY_LENGTH {
+            return Err((
+                StatusCode::PAYLOAD_TOO_LARGE,
+                "Request too large".to_string(),
+            ));
+        }
+        bytes.append(&mut buf.to_vec());
+    }
+    Ok(bytes)
+}
+
 impl DIDKitHTTPSvc {
     pub fn new(keys: KeyMap, resolver_options: ResolverOptions) -> Self {
         Self {
@@ -204,8 +230,13 @@ impl DIDKitHTTPSvc {
         let keys = self.keys.clone();
         let resolver_options = self.resolver_options.clone();
         Box::pin(async move {
-            let body_reader = hyper::body::aggregate(req).await?.reader();
-            let issue_req: IssueCredentialRequest = match serde_json::from_reader(body_reader) {
+            let bytes = match get_body_bytes_limited(req).await {
+                Ok(bytes) => bytes,
+                Err((status_code, message)) => {
+                    return Self::response(status_code, message).await;
+                }
+            };
+            let issue_req: IssueCredentialRequest = match serde_json::from_slice(&bytes) {
                 Ok(reader) => reader,
                 Err(err) => {
                     return Self::response(StatusCode::BAD_REQUEST, err.to_string()).await;
@@ -277,8 +308,13 @@ impl DIDKitHTTPSvc {
         };
         let resolver_options = self.resolver_options.clone();
         Box::pin(async move {
-            let body_reader = hyper::body::aggregate(req).await?.reader();
-            let verify_req: VerifyCredentialRequest = match serde_json::from_reader(body_reader) {
+            let bytes = match get_body_bytes_limited(req).await {
+                Ok(bytes) => bytes,
+                Err((status_code, message)) => {
+                    return Self::response(status_code, message).await;
+                }
+            };
+            let verify_req: VerifyCredentialRequest = match serde_json::from_slice(&bytes) {
                 Ok(reader) => reader,
                 Err(err) => {
                     return Self::response(StatusCode::BAD_REQUEST, err.to_string()).await;
@@ -333,8 +369,13 @@ impl DIDKitHTTPSvc {
         let keys = self.keys.clone();
         let resolver_options = self.resolver_options.clone();
         Box::pin(async move {
-            let body_reader = hyper::body::aggregate(req).await?.reader();
-            let issue_req: ProvePresentationRequest = match serde_json::from_reader(body_reader) {
+            let bytes = match get_body_bytes_limited(req).await {
+                Ok(bytes) => bytes,
+                Err((status_code, message)) => {
+                    return Self::response(status_code, message).await;
+                }
+            };
+            let issue_req: ProvePresentationRequest = match serde_json::from_slice(&bytes) {
                 Ok(reader) => reader,
                 Err(err) => {
                     return Self::response(StatusCode::BAD_REQUEST, err.to_string()).await;
@@ -408,8 +449,13 @@ impl DIDKitHTTPSvc {
         };
         let resolver_options = self.resolver_options.clone();
         Box::pin(async move {
-            let body_reader = hyper::body::aggregate(req).await?.reader();
-            let verify_req: VerifyPresentationRequest = match serde_json::from_reader(body_reader) {
+            let bytes = match get_body_bytes_limited(req).await {
+                Ok(bytes) => bytes,
+                Err((status_code, message)) => {
+                    return Self::response(status_code, message).await;
+                }
+            };
+            let verify_req: VerifyPresentationRequest = match serde_json::from_slice(&bytes) {
                 Ok(reader) => reader,
                 Err(err) => {
                     return Self::response(StatusCode::BAD_REQUEST, err.to_string()).await;