diff --git a/src/error.rs b/src/error.rs index 9a839dd..ff58f4e 100644 --- a/src/error.rs +++ b/src/error.rs @@ -228,6 +228,15 @@ pub enum ProxyError { #[error("Session terminated")] RouteSwitched, + #[error("Traffic budget wait cancelled")] + TrafficBudgetWaitCancelled, + + #[error("Traffic budget wait deadline exceeded")] + TrafficBudgetWaitDeadlineExceeded, + + #[error("ME client writer cancelled")] + MiddleClientWriterCancelled, + // ============= Config Errors ============= #[error("Config error: {0}")] Config(String), diff --git a/src/proxy/middle_relay.rs b/src/proxy/middle_relay.rs index 6b0329c..8fadf9b 100644 --- a/src/proxy/middle_relay.rs +++ b/src/proxy/middle_relay.rs @@ -697,7 +697,7 @@ async fn wait_for_traffic_budget( let wait_started_at = Instant::now(); if deadline.is_some_and(|deadline| wait_started_at >= deadline) { - return Err(ProxyError::Proxy("traffic budget wait deadline exceeded".into())); + return Err(ProxyError::TrafficBudgetWaitDeadlineExceeded); } tokio::time::sleep(next_refill_delay()).await; let wait_ms = wait_started_at @@ -741,13 +741,13 @@ async fn wait_for_traffic_budget_or_cancel( let wait_started_at = Instant::now(); if deadline.is_some_and(|deadline| wait_started_at >= deadline) { stats.increment_flow_wait_middle_rate_limit_cancelled_total(); - return Err(ProxyError::Proxy("traffic budget wait deadline exceeded".into())); + return Err(ProxyError::TrafficBudgetWaitDeadlineExceeded); } tokio::select! { _ = tokio::time::sleep(next_refill_delay()) => {} _ = cancel.cancelled() => { stats.increment_flow_wait_middle_rate_limit_cancelled_total(); - return Err(ProxyError::Proxy("traffic budget wait cancelled".into())); + return Err(ProxyError::TrafficBudgetWaitCancelled); } } let wait_ms = wait_started_at @@ -2763,7 +2763,7 @@ where { tokio::select! { result = client_writer.write_all(bytes) => result.map_err(ProxyError::Io), - _ = cancel.cancelled() => Err(ProxyError::Proxy("ME client writer cancelled".into())), + _ = cancel.cancelled() => Err(ProxyError::MiddleClientWriterCancelled), } } @@ -2776,7 +2776,7 @@ where { tokio::select! { result = client_writer.flush() => result.map_err(ProxyError::Io), - _ = cancel.cancelled() => Err(ProxyError::Proxy("ME client writer cancelled".into())), + _ = cancel.cancelled() => Err(ProxyError::MiddleClientWriterCancelled), } } diff --git a/src/proxy/tests/middle_relay_atomic_quota_invariant_tests.rs b/src/proxy/tests/middle_relay_atomic_quota_invariant_tests.rs index e1b3511..5cf6ac5 100644 --- a/src/proxy/tests/middle_relay_atomic_quota_invariant_tests.rs +++ b/src/proxy/tests/middle_relay_atomic_quota_invariant_tests.rs @@ -251,7 +251,7 @@ async fn me_writer_data_write_obeys_flow_cancellation() { .await; assert!( - matches!(result, Err(ProxyError::Proxy(ref message)) if message == "ME client writer cancelled"), + matches!(result, Err(ProxyError::MiddleClientWriterCancelled)), "cancelled middle writer must return a bounded cancellation error" ); assert_eq!(