Merge pull request #173 from NethermindEth/feature/rate-limiter

Rate limiter feature
NethermindEth · Oct 24, 2023 · 8b6ca3b · 8b6ca3b
2 parents 2745fb6 + cae89a1
commit 8b6ca3b
Show file tree

Hide file tree

Showing 13 changed files with 254 additions and 72 deletions.
diff --git a/api/Cargo.lock b/api/Cargo.lock
diff --git a/api/src/handlers/cairo_version.rs b/api/src/handlers/cairo_version.rs
@@ -1,5 +1,6 @@
 use crate::handlers::process::{do_process_command, fetch_process_result};
 use crate::handlers::types::{ApiCommand, ApiCommandResult};
+use crate::rate_limiter::RateLimited;
 use crate::types::{ApiError, Result};
 use crate::utils::lib::DEFAULT_CAIRO_DIR;
 use crate::worker::WorkerEngine;
@@ -10,15 +11,18 @@ use tracing::{error, info, instrument};
 // Read the version from the cairo Cargo.toml file.
 #[instrument]
 #[get("/cairo_version")]
-pub async fn cairo_version() -> String {
+pub async fn cairo_version(_rate_limited: RateLimited) -> String {
     info!("/cairo_version");
     do_cairo_version().unwrap_or_else(|e| format!("Failed to get cairo version: {:?}", e))
 }
 
 // Read the version from the cairo Cargo.toml file.
 #[instrument]
 #[get("/cairo_version_async")]
-pub async fn cairo_version_async(engine: &State<WorkerEngine>) -> String {
+pub async fn cairo_version_async(
+    engine: &State<WorkerEngine>,
+    _rate_limited: RateLimited,
+) -> String {
     info!("/cairo_version_async");
     do_process_command(ApiCommand::CairoVersion, engine)
 }
@@ -51,9 +55,9 @@ pub fn do_cairo_version() -> Result<String> {
             .arg("--version")
             .stdout(Stdio::piped())
             .spawn()
-            .map_err(|e| ApiError::FailedToExecuteCommand(e))?
+            .map_err(ApiError::FailedToExecuteCommand)?
             .wait_with_output()
-            .map_err(|e| ApiError::FailedToReadOutput(e))?
+            .map_err(ApiError::FailedToReadOutput)?
             .stdout,
     ) {
         Ok(version) => Ok(version),

diff --git a/api/src/handlers/cairo_versions.rs b/api/src/handlers/cairo_versions.rs
@@ -16,16 +16,13 @@ pub async fn cairo_versions() -> String {
 pub async fn do_cairo_versions() -> Result<String> {
     let path = Path::new(CAIRO_COMPILERS_DIR);
 
-    let mut dir = read_dir(path)
-        .await
-        .map_err(|e| ApiError::FailedToReadDir(e))?;
+    let mut dir = read_dir(path).await.map_err(ApiError::FailedToReadDir)?;
     let mut result = vec![];
 
     while let Ok(Some(entry)) = dir.next_entry().await {
         let entry = entry;
         let path = entry.path();
         if path.is_dir() {
-            println!("{:?}", entry.file_name());
             result.push(entry.file_name().to_string_lossy().to_string());
         }
     }

diff --git a/api/src/handlers/compile_casm.rs b/api/src/handlers/compile_casm.rs
@@ -1,5 +1,6 @@
 use crate::handlers::process::{do_process_command, fetch_process_result};
 use crate::handlers::types::{ApiCommand, ApiCommandResult, CompileResponse};
+use crate::rate_limiter::RateLimited;
 use crate::types::{ApiError, Result};
 use crate::utils::lib::{get_file_ext, get_file_path, CAIRO_COMPILERS_DIR, CASM_ROOT};
 use crate::worker::WorkerEngine;
@@ -14,7 +15,11 @@ use tracing::info;
 use tracing::instrument;
 
 #[get("/compile-to-casm/<version>/<remix_file_path..>")]
-pub async fn compile_to_casm(version: String, remix_file_path: PathBuf) -> Json<CompileResponse> {
+pub async fn compile_to_casm(
+    version: String,
+    remix_file_path: PathBuf,
+    _rate_limited: RateLimited,
+) -> Json<CompileResponse> {
     info!("/compile-to-casm/{:?}", remix_file_path);
     do_compile_to_casm(version.clone(), remix_file_path)
         .await
@@ -34,6 +39,7 @@ pub async fn compile_to_casm_async(
     version: String,
     remix_file_path: PathBuf,
     engine: &State<WorkerEngine>,
+    _rate_limited: RateLimited,
 ) -> String {
     info!("/compile-to-casm-async/{:?}", remix_file_path);
     do_process_command(
@@ -120,28 +126,28 @@ pub async fn do_compile_to_casm(
         .arg(&casm_path)
         .stderr(Stdio::piped())
         .spawn()
-        .map_err(|e| ApiError::FailedToExecuteCommand(e))?;
+        .map_err(ApiError::FailedToExecuteCommand)?;
 
     debug!("LOG: ran command:{:?}", compile);
 
     let output = result
         .wait_with_output()
-        .map_err(|e| ApiError::FailedToReadOutput(e))?;
+        .map_err(ApiError::FailedToReadOutput)?;
 
     let file_content = fs::read_to_string(
         NamedFile::open(&casm_path)
             .await
-            .map_err(|e| ApiError::FailedToReadFile(e))?
+            .map_err(ApiError::FailedToReadFile)?
             .path()
             .to_str()
             .ok_or(ApiError::FailedToParseString)?
             .to_string(),
     )
     .await
-    .map_err(|e| ApiError::FailedToReadFile(e))?;
+    .map_err(ApiError::FailedToReadFile)?;
 
     let message = String::from_utf8(output.stderr)
-        .map_err(|e| ApiError::UTF8Error(e))?
+        .map_err(ApiError::UTF8Error)?
         .replace(
             &file_path
                 .to_str()

diff --git a/api/src/handlers/compile_sierra.rs b/api/src/handlers/compile_sierra.rs
@@ -1,5 +1,6 @@
 use crate::handlers::process::{do_process_command, fetch_process_result};
 use crate::handlers::types::{ApiCommand, ApiCommandResult, CompileResponse};
+use crate::rate_limiter::RateLimited;
 use crate::types::{ApiError, Result};
 use crate::utils::lib::{get_file_ext, get_file_path, CAIRO_COMPILERS_DIR, SIERRA_ROOT};
 use crate::worker::WorkerEngine;
@@ -14,7 +15,11 @@ use tracing::{debug, instrument};
 
 #[instrument]
 #[get("/compile-to-sierra/<version>/<remix_file_path..>")]
-pub async fn compile_to_sierra(version: String, remix_file_path: PathBuf) -> Json<CompileResponse> {
+pub async fn compile_to_sierra(
+    version: String,
+    remix_file_path: PathBuf,
+    _rate_limited: RateLimited,
+) -> Json<CompileResponse> {
     info!("/compile-to-sierra");
 
     let res = do_compile_to_sierra(version.clone(), remix_file_path).await;
@@ -36,6 +41,7 @@ pub async fn compile_to_siera_async(
     version: String,
     remix_file_path: PathBuf,
     engine: &State<WorkerEngine>,
+    _rate_limited: RateLimited,
 ) -> String {
     info!("/compile-to-sierra-async");
     do_process_command(
@@ -124,27 +130,27 @@ pub async fn do_compile_to_sierra(
         .stderr(Stdio::piped())
         .stdout(Stdio::piped())
         .spawn()
-        .map_err(|e| ApiError::FailedToExecuteCommand(e))?;
+        .map_err(ApiError::FailedToExecuteCommand)?;
 
     debug!("LOG: ran command:{:?}", compile);
 
     let output = result
         .wait_with_output()
-        .map_err(|e| ApiError::FailedToReadOutput(e))?;
+        .map_err(ApiError::FailedToReadOutput)?;
 
     let file_content = fs::read_to_string(
         NamedFile::open(&sierra_path)
             .await
-            .map_err(|e| ApiError::FailedToReadFile(e))?
+            .map_err(ApiError::FailedToReadFile)?
             .path()
             .to_str()
             .ok_or(ApiError::FailedToParseString)?,
     )
     .await
-    .map_err(|e| ApiError::FailedToReadFile(e))?;
+    .map_err(ApiError::FailedToReadFile)?;
 
     let message = String::from_utf8(output.stderr)
-        .map_err(|e| ApiError::UTF8Error(e))?
+        .map_err(ApiError::UTF8Error)?
         .replace(
             &file_path
                 .to_str()

diff --git a/api/src/handlers/mod.rs b/api/src/handlers/mod.rs
@@ -70,7 +70,7 @@ fn get_files_recursive(base_path: &Path) -> Result<Vec<FileContentMap>> {
     if base_path.is_dir() {
         for entry in base_path
             .read_dir()
-            .map_err(|e| ApiError::FailedToReadDir(e))?
+            .map_err(ApiError::FailedToReadDir)?
             .flatten()
         {
             let path = entry.path();

diff --git a/api/src/handlers/save_code.rs b/api/src/handlers/save_code.rs
@@ -44,7 +44,7 @@ pub async fn do_save_code(file: Data<'_>, remix_file_path: PathBuf) -> Result<St
         .open(128_i32.gibibytes())
         .into_file(&file_path)
         .await
-        .map_err(|e| ApiError::FailedToSaveFile(e))?;
+        .map_err(ApiError::FailedToSaveFile)?;
 
     Ok(file_path
         .to_str()

diff --git a/api/src/handlers/scarb_compile.rs b/api/src/handlers/scarb_compile.rs
@@ -1,6 +1,7 @@
 use crate::handlers::get_files_recursive;
 use crate::handlers::process::{do_process_command, fetch_process_result};
 use crate::handlers::types::{ApiCommand, ApiCommandResult, ScarbCompileResponse};
+use crate::rate_limiter::RateLimited;
 use crate::types::{ApiError, Result};
 use crate::utils::lib::get_file_path;
 use crate::worker::WorkerEngine;
@@ -13,7 +14,10 @@ use tracing::{info, instrument};
 
 #[instrument]
 #[get("/compile-scarb/<remix_file_path..>")]
-pub async fn scarb_compile(remix_file_path: PathBuf) -> Json<ScarbCompileResponse> {
+pub async fn scarb_compile(
+    remix_file_path: PathBuf,
+    _rate_limited: RateLimited,
+) -> Json<ScarbCompileResponse> {
     info!("/compile-scarb/{:?}", remix_file_path);
     do_scarb_compile(remix_file_path).await.unwrap_or_else(|e| {
         Json(ScarbCompileResponse {
@@ -26,7 +30,11 @@ pub async fn scarb_compile(remix_file_path: PathBuf) -> Json<ScarbCompileRespons
 
 #[instrument]
 #[get("/compile-scarb-async/<remix_file_path..>")]
-pub async fn scarb_compile_async(remix_file_path: PathBuf, engine: &State<WorkerEngine>) -> String {
+pub async fn scarb_compile_async(
+    remix_file_path: PathBuf,
+    engine: &State<WorkerEngine>,
+    _rate_limited: RateLimited,
+) -> String {
     info!("/compile-scarb-async/{:?}", remix_file_path);
     do_process_command(ApiCommand::ScarbCompile { remix_file_path }, engine)
 }
@@ -60,18 +68,18 @@ pub async fn do_scarb_compile(remix_file_path: PathBuf) -> Result<Json<ScarbComp
         .stdout(Stdio::piped())
         .stderr(Stdio::piped())
         .spawn()
-        .map_err(|e| ApiError::FailedToExecuteCommand(e))?;
+        .map_err(ApiError::FailedToExecuteCommand)?;
 
     debug!("LOG: ran command:{:?}", compile);
 
     let output = result
         .wait_with_output()
-        .map_err(|e| ApiError::FailedToReadOutput(e))?;
+        .map_err(ApiError::FailedToReadOutput)?;
 
     let file_content_map_array = get_files_recursive(&file_path.join("target/dev"))?;
 
     let message = String::from_utf8(output.stdout)
-        .map_err(|e| ApiError::UTF8Error(e))?
+        .map_err(ApiError::UTF8Error)?
         .replace(
             &file_path
                 .to_str()
@@ -80,7 +88,7 @@ pub async fn do_scarb_compile(remix_file_path: PathBuf) -> Result<Json<ScarbComp
             &remix_file_path,
         )
         + &String::from_utf8(output.stderr)
-            .map_err(|e| ApiError::UTF8Error(e))?
+            .map_err(ApiError::UTF8Error)?
             .replace(
                 &file_path
                     .to_str()

diff --git a/api/src/main.rs b/api/src/main.rs
@@ -3,12 +3,14 @@ extern crate rocket;
 
 pub mod cors;
 pub mod handlers;
+pub mod rate_limiter;
 pub mod tracing_log;
 pub mod types;
 pub mod utils;
 pub mod worker;
 
 use crate::cors::CORS;
+use crate::rate_limiter::RateLimiter;
 use crate::tracing_log::init_logger;
 use crate::worker::WorkerEngine;
 use handlers::cairo_version::{cairo_version, cairo_version_async, get_cairo_version_result};
@@ -49,26 +51,30 @@ async fn rocket() -> _ {
 
     info!("Starting Rocket webserver...");
 
-    rocket::build().manage(engine).attach(CORS).mount(
-        "/",
-        routes![
-            compile_to_sierra,
-            compile_to_siera_async,
-            get_siera_compile_result,
-            compile_to_casm,
-            compile_to_casm_async,
-            compile_to_casm_result,
-            scarb_compile,
-            scarb_compile_async,
-            get_scarb_compile_result,
-            save_code,
-            cairo_versions,
-            cairo_version,
-            cairo_version_async,
-            get_cairo_version_result,
-            get_process_status,
-            health,
-            who_is_this,
-        ],
-    )
+    rocket::build()
+        .manage(engine)
+        .attach(CORS)
+        .manage(RateLimiter::new())
+        .mount(
+            "/",
+            routes![
+                compile_to_sierra,
+                compile_to_siera_async,
+                get_siera_compile_result,
+                compile_to_casm,
+                compile_to_casm_async,
+                compile_to_casm_result,
+                scarb_compile,
+                scarb_compile_async,
+                get_scarb_compile_result,
+                save_code,
+                cairo_versions,
+                cairo_version,
+                cairo_version_async,
+                get_cairo_version_result,
+                get_process_status,
+                health,
+                who_is_this,
+            ],
+        )
 }