feat(logger): rate limit based on peer address

cargo-shuttle/src/lib.rs

@@ -698,9 +698,22 @@ impl Shuttle {
 
             while let Some(Ok(msg)) = stream.next().await {
                 if let tokio_tungstenite::tungstenite::Message::Text(line) = msg {
-                    let log_item: shuttle_common::LogItem = serde_json::from_str(&line)
-                        .context("Failed parsing logs. Is your cargo-shuttle outdated?")?;
-                    println!("{log_item}")
+                    match serde_json::from_str::<shuttle_common::LogItem>(&line) {
+                        Ok(log_item) => {
+                            println!("{log_item}")
+                        }
+                        Err(err) => {
+                            debug!(error = %err, "failed to parse message into log item");
+
+                            let message = if let Ok(err) = serde_json::from_str::<ApiError>(&line) {
+                                err.to_string()
+                            } else {
+                                "failed to parse logs, is your cargo-shuttle outdated?".to_string()
+                            };
+
+                            bail!(message);
+                        }
+                    }
                 }
             }
         } else {

common/src/models/error.rs

@@ -59,6 +59,7 @@ pub enum ErrorKind {
     NotReady,
     ServiceUnavailable,
     DeleteProjectFailed,
+    RateLimited(String),
 }
 
 impl From<ErrorKind> for ApiError {
@@ -121,6 +122,12 @@ impl From<ErrorKind> for ApiError {
             ErrorKind::Forbidden => (StatusCode::FORBIDDEN, "forbidden"),
             ErrorKind::NotReady => (StatusCode::INTERNAL_SERVER_ERROR, "service not ready"),
             ErrorKind::DeleteProjectFailed => (StatusCode::INTERNAL_SERVER_ERROR, "deleting project failed"),
+            ErrorKind::RateLimited(message) => {
+                return Self {
+                    message,
+                    status_code: StatusCode::TOO_MANY_REQUESTS.as_u16(),
+                }
+            },
         };
         Self {
             message: error_message.to_string(),

deployer/src/handlers/error.rs

@@ -26,6 +26,8 @@ pub enum Error {
     Internal(#[from] anyhow::Error),
     #[error("Missing header: {0}")]
     MissingHeader(String),
+    #[error("{0}. Retry the request in a few minutes")]
+    RateLimited(String),
 }
 
 impl Serialize for Error {
@@ -47,6 +49,7 @@ impl IntoResponse for Error {
 
         let code = match self {
             Error::NotFound(_) => StatusCode::NOT_FOUND,
+            Error::RateLimited(_) => StatusCode::TOO_MANY_REQUESTS,
             _ => StatusCode::INTERNAL_SERVER_ERROR,
         };
 

deployer/src/handlers/mod.rs

@@ -33,7 +33,7 @@ use shuttle_common::{
     claims::{Claim, Scope},
     models::{
         deployment::{DeploymentRequest, CREATE_SERVICE_BODY_LIMIT, GIT_STRINGS_MAX_LENGTH},
-        error::axum::CustomErrorPath,
+        error::{axum::CustomErrorPath, ApiError, ErrorKind},
         project::ProjectName,
     },
     request_span, LogItem,
@@ -653,16 +653,27 @@ pub async fn get_logs(
     logs_request.extensions_mut().insert(claim);
 
     let mut client = deployment_manager.logs_fetcher().clone();
-    if let Ok(logs) = client.get_logs(logs_request).await {
-        Ok(Json(
+
+    match client.get_logs(logs_request).await {
+        Ok(logs) => Ok(Json(
             logs.into_inner()
                 .log_items
                 .into_iter()
                 .map(|l| l.to_log_item_with_id(deployment_id))
                 .collect(),
-        ))
-    } else {
-        Err(Error::NotFound("deployment not found".to_string()))
+        )),
+        Err(error) => {
+            if error.code() == tonic::Code::Unavailable
+                && error.metadata().get("x-ratelimit-limit").is_some()
+            {
+                Err(Error::RateLimited(
+                    "your application is producing too many logs. Interactions with the shuttle logger service will be rate limited"
+                        .to_string(),
+                ))
+            } else {
+                Err(anyhow!("failed to retrieve logs for deployment").into())
+            }
+        }
     }
 }
 
@@ -712,9 +723,24 @@ async fn logs_websocket_handler(
                 "failed to get backlog of logs"
             );
 
-            let _ = s
-                .send(ws::Message::Text("failed to get logs".to_string()))
-                .await;
+            if error.code() == tonic::Code::Unavailable
+                && error.metadata().get("x-ratelimit-limit").is_some()
+            {
+                let message = serde_json::to_string(
+                    &ApiError::from(
+                        ErrorKind::RateLimited(
+                            "your application is producing too many logs. Interactions with the shuttle logger service will be rate limited"
+                            .to_string()
+                        )
+                    ))
+                    .expect("to convert error to json");
+
+                let _ = s.send(ws::Message::Text(message)).await;
+            } else {
+                let _ = s
+                    .send(ws::Message::Text("failed to get logs".to_string()))
+                    .await;
+            }
             let _ = s.close().await;
             return;
         }

logger/Cargo.toml

@@ -12,6 +12,7 @@ shuttle-proto = { workspace = true }
 async-trait = { workspace = true }
 chrono = { workspace = true }
 clap = { workspace = true }
+http = { workspace = true }
 prost-types = { workspace = true }
 serde_json = { workspace = true }
 sqlx = { workspace = true, features = [
@@ -24,6 +25,8 @@ thiserror = { workspace = true }
 tokio = { workspace = true, features = ["rt-multi-thread"] }
 tokio-stream = { workspace = true }
 tonic = { workspace = true }
+tower = { workspace = true }
+tower_governor = { version= "0.1.0", features = ["tracing"] }
 tracing = { workspace = true }
 tracing-subscriber = { workspace = true, features = ["default"] }
 
@@ -35,3 +38,5 @@ serde_json = { workspace = true }
 shuttle-common-tests = { workspace = true }
 uuid = { workspace = true }
 ctor = { workspace = true }
+futures = { workspace = true }
+tower = { workspace = true, features = ["util"] }

logger/src/lib.rs

@@ -15,6 +15,7 @@ use tracing::{debug, error, field, Span};
 
 pub mod args;
 mod dal;
+pub mod rate_limiting;
 
 pub use dal::Postgres;