From 1cba761b81cd24b8abb5b33e6826f13ae07d6779 Mon Sep 17 00:00:00 2001
From: Joey Kraut <joey@renegade.fi>
Date: Sun, 21 Apr 2024 11:38:26 -0700
Subject: [PATCH] renegade-dealer: Add size limits to request

---
 renegade-dealer-api/src/lib.rs | 20 +++++++++++++++++
 renegade-dealer/src/main.rs    | 39 ++++++++++++++++++++++++++++------
 2 files changed, 52 insertions(+), 7 deletions(-)

diff --git a/renegade-dealer-api/src/lib.rs b/renegade-dealer-api/src/lib.rs
index caab11c..82286c1 100644
--- a/renegade-dealer-api/src/lib.rs
+++ b/renegade-dealer-api/src/lib.rs
@@ -21,6 +21,15 @@ type Scalar = ark_mpc::algebra::Scalar<Curve>;
 /// A type alias for a scalar share
 type ScalarShare = ark_mpc::algebra::ScalarShare<Curve>;
 
+/// A response to a bad request
+#[derive(Debug, Serialize, Deserialize)]
+pub struct ErrorResponse {
+    /// The error code associated with the response
+    pub code: u32,
+    /// The error message associated with the response
+    pub message: &'static str,
+}
+
 /// A request for offline phase randomness from the dealer
 #[derive(Clone, Debug, Serialize, Deserialize, Eq, PartialEq)]
 pub struct DealerRequest {
@@ -44,6 +53,17 @@ pub struct DealerRequest {
     pub n_triples: u32,
 }
 
+impl DealerRequest {
+    /// Return the total number of requested values
+    pub fn total_values(&self) -> u32 {
+        self.n_random_bits
+            + self.n_random_values
+            + self.n_input_masks
+            + self.n_inverse_pairs
+            + self.n_triples
+    }
+}
+
 /// A response from the Dealer
 #[derive(Clone, Debug, Serialize, Deserialize, Default)]
 pub struct DealerResponse {
diff --git a/renegade-dealer/src/main.rs b/renegade-dealer/src/main.rs
index 94ccc64..8ade0f8 100644
--- a/renegade-dealer/src/main.rs
+++ b/renegade-dealer/src/main.rs
@@ -22,9 +22,17 @@ use clap::Parser;
 use dealer::{
     create_dealer_sender_receiver, create_response_sender_receiver, Dealer, DealerJob, DealerSender,
 };
-use renegade_dealer_api::{DealerRequest, DealerResponse, RequestId};
+use renegade_dealer_api::{DealerRequest, DealerResponse, ErrorResponse, RequestId};
 use warp::Filter;
 
+/// The maximum number of values that may be requested at once by a pair
+const MAX_REQUEST_SIZE: u32 = 1_500_000;
+
+/// An error type indicating a bad request
+#[derive(Debug)]
+struct BadRequestError(&'static str);
+impl warp::reject::Reject for BadRequestError {}
+
 /// Renegade Dealer server configuration
 #[derive(Parser, Debug)]
 #[clap(author, version, about, long_about = None)]
@@ -52,12 +60,16 @@ async fn main() {
         .and_then(move |request_id, body| {
             let dealer_send = dealer_send.clone();
             async move {
-                let resp = handle_req(request_id, body, dealer_send).await;
-                Ok::<_, warp::Rejection>(warp::reply::json(&resp))
+                match handle_req(request_id, body, dealer_send).await {
+                    Ok(resp) => Ok(warp::reply::json(&resp)),
+                    Err(rej) => Err(rej),
+                }
             }
         });
 
-    warp::serve(setup).run(([127, 0, 0, 1], cli.port)).await
+    let routes = setup.recover(handle_rejection);
+
+    warp::serve(routes).run(([127, 0, 0, 1], cli.port)).await
 }
 
 /// Handle an incoming client request
@@ -65,10 +77,23 @@ async fn handle_req(
     request_id: RequestId,
     body: DealerRequest,
     dealer_queue: DealerSender,
-) -> DealerResponse {
-    // Send a request to the dealer
+) -> Result<DealerResponse, warp::Rejection> {
+    if body.total_values() > MAX_REQUEST_SIZE {
+        return Err(warp::reject::custom(BadRequestError("Request size too large")));
+    }
+
     let (send, mut recv) = create_response_sender_receiver();
     dealer_queue.send(DealerJob::new(request_id, body, send)).unwrap();
 
-    recv.recv().await.unwrap()
+    Ok(recv.recv().await.unwrap())
+}
+
+/// Handle a rejection from the dealer
+async fn handle_rejection(err: warp::Rejection) -> Result<impl warp::Reply, warp::Rejection> {
+    if let Some(BadRequestError(msg)) = err.find::<BadRequestError>() {
+        let json = warp::reply::json(&ErrorResponse { message: msg, code: 400 });
+        Ok(warp::reply::with_status(json, warp::http::StatusCode::BAD_REQUEST))
+    } else {
+        Err(err)
+    }
 }