Merge pull request #139 from brave/increase-nofile-limit

Add switch/mode for increasing nofile limit to handle more concurrent connections within enclave

Cargo.toml

@@ -12,6 +12,7 @@ base64 = "0.21.2"
 clap = { version = "4.3.11", features = ["derive"] }
 curve25519-dalek = "3.2.0"
 ppoprf = { git = "https://github.com/brave/sta-rs", rev = "1937393e9d76c5333af47dd2ae2778b7e63fafd4" }
+rlimit = "0.10"
 serde = "1.0.171"
 serde_json = "1.0.102"
 thiserror = "1.0.43"

src/main.rs

@@ -2,6 +2,7 @@
 
 use axum::{routing::get, routing::post, Router};
 use clap::Parser;
+use rlimit::Resource;
 use std::sync::{Arc, RwLock};
 use time::format_description::well_known::Rfc3339;
 use time::OffsetDateTime;
@@ -39,12 +40,15 @@ pub struct Config {
     /// invocations.
     #[arg(long, value_name = "RFC 3339 timestamp", value_parser = parse_timestamp)]
     epoch_base_time: Option<OffsetDateTime>,
+    /// Increases OS nofile limit to 65535, so the server can handle
+    /// more concurrent connections.
+    #[arg(long, default_value_t = false)]
+    increase_nofile_limit: bool,
 }
 
 /// Parse a timestamp given as a config option
 fn parse_timestamp(stamp: &str) -> Result<OffsetDateTime, &'static str> {
-    OffsetDateTime::parse(stamp, &Rfc3339)
-        .map_err(|_| "Try something like '2023-05-15T04:30:00Z'.")
+    OffsetDateTime::parse(stamp, &Rfc3339).map_err(|_| "Try something like '2023-05-15T04:30:00Z'.")
 }
 
 /// Initialize an axum::Router for our web service
@@ -62,6 +66,20 @@ fn app(oprf_state: OPRFState) -> Router {
         .layer(tower_http::trace::TraceLayer::new_for_http())
 }
 
+fn increase_nofile_limit() {
+    let curr_limits =
+        rlimit::getrlimit(Resource::NOFILE).expect("should be able to get current nofile limit");
+    info!("Current nofile limits = {:?}", curr_limits);
+
+    rlimit::setrlimit(Resource::NOFILE, 65535, 65535).expect("should be able to set nofile limit");
+    let curr_limits = rlimit::getrlimit(Resource::NOFILE)
+        .expect("should be able to get current nofile limit after updating it");
+    info!(
+        "Attempted nofile limit change! Current nofile limits = {:?}",
+        curr_limits
+    );
+}
+
 #[tokio::main]
 async fn main() {
     // Start logging
@@ -74,19 +92,20 @@ async fn main() {
     debug!(?config, "config parsed");
     let addr = config.listen.parse().unwrap();
 
+    if config.increase_nofile_limit {
+        increase_nofile_limit();
+    }
+
     // Oblivious function state
     info!("initializing OPRF state...");
-    let server = state::OPRFServer::new(&config)
-        .expect("Could not initialize PPOPRF state");
+    let server = state::OPRFServer::new(&config).expect("Could not initialize PPOPRF state");
     info!("epoch now {}", server.epoch);
     let oprf_state = Arc::new(RwLock::new(server));
 
     // Spawn a background process to advance the epoch
     info!("Spawning background epoch rotation task...");
     let background_state = oprf_state.clone();
-    tokio::spawn(async move {
-        state::epoch_loop(background_state, &config).await
-    });
+    tokio::spawn(async move { state::epoch_loop(background_state, &config).await });
 
     // Set up routes and middleware
     info!("initializing routes...");

src/tests.rs

@@ -24,10 +24,10 @@ fn test_app() -> crate::Router {
         first_epoch: EPOCH,
         last_epoch: EPOCH * 2,
         epoch_base_time: None,
+        increase_nofile_limit: false,
     };
     // server state
-    let mut server =
-        OPRFServer::new(&config).expect("Could not initialize PPOPRF state");
+    let mut server = OPRFServer::new(&config).expect("Could not initialize PPOPRF state");
     server.next_epoch_time = Some(NEXT_EPOCH_TIME.to_owned());
     let oprf_state = Arc::new(RwLock::new(server));
 
@@ -79,8 +79,8 @@ async fn info() {
     assert_eq!(response.status(), StatusCode::OK);
     let body = hyper::body::to_bytes(response.into_body()).await.unwrap();
     assert!(!body.is_empty());
-    let json: Value = serde_json::from_slice(body.as_ref())
-        .expect("Could not parse response body as json");
+    let json: Value =
+        serde_json::from_slice(body.as_ref()).expect("Could not parse response body as json");
     assert!(json.is_object());
     println!("{:?}", json);
     assert_eq!(json["currentEpoch"], json!(EPOCH));
@@ -173,6 +173,7 @@ async fn epoch_base_time() {
         first_epoch: EPOCH,
         last_epoch: EPOCH * 2,
         epoch_base_time: Some(now - delay),
+        increase_nofile_limit: false,
     };
     // Verify test parameters are compatible with the
     // expected_epoch calculation.
@@ -187,14 +188,11 @@ async fn epoch_base_time() {
         .expect("well-known timestamp format should always succeed");
 
     // server state
-    let server =
-        OPRFServer::new(&config).expect("Could not initialize PPOPRF state");
+    let server = OPRFServer::new(&config).expect("Could not initialize PPOPRF state");
     let oprf_state = Arc::new(RwLock::new(server));
     // background task to manage epoch rotation
     let background_state = oprf_state.clone();
-    tokio::spawn(async move {
-        crate::state::epoch_loop(background_state, &config).await
-    });
+    tokio::spawn(async move { crate::state::epoch_loop(background_state, &config).await });
 
     // Wait for `epoch_loop` to update `next_epoch_time` as a proxy
     // for completing epoch schedule initialization. Use a timeout
@@ -218,8 +216,8 @@ async fn epoch_base_time() {
     assert_eq!(response.status(), StatusCode::OK);
     let body = hyper::body::to_bytes(response.into_body()).await.unwrap();
     assert!(!body.is_empty());
-    let json: Value = serde_json::from_slice(body.as_ref())
-        .expect("Could not parse response body as json");
+    let json: Value =
+        serde_json::from_slice(body.as_ref()).expect("Could not parse response body as json");
     assert!(json.is_object());
     println!("{:?}", json);
     assert_eq!(json["currentEpoch"], json!(expected_epoch));
@@ -232,8 +230,8 @@ async fn epoch_base_time() {
 fn verify_randomness_body(body: axum::body::Bytes, expected_points: usize) {
     // Randomness should return a list of points and an epoch.
     assert!(!body.is_empty());
-    let json: Value = serde_json::from_slice(body.as_ref())
-        .expect("Response body should parse as json");
+    let json: Value =
+        serde_json::from_slice(body.as_ref()).expect("Response body should parse as json");
     // Top-level value should be an object.
     assert!(json.is_object());
     // Epoch should match test_app.