use std::{net::SocketAddr, sync::Arc}; use hudsucker::{ certificate_authority::RcgenAuthority, hyper::{Request, Uri}, rustls::crypto::aws_lc_rs, Body, HttpContext, HttpHandler, Proxy, RequestOrResponse, }; use tokio::{net::TcpListener, sync::oneshot, task::JoinHandle}; use parking_lot::RwLock; use crate::{ cursor::{proxy::UPSTREAM_URL_HEADER, tab::is_tab_path}, store::TabMode, Error, Result, }; use super::ca::LoadedCa; #[derive(Default)] pub struct ProxyRuntime { url: Option, port: Option, stop: Option>, task: Option>, } impl ProxyRuntime { pub fn running(&self) -> bool { self.task.as_ref().is_some_and(|task| !task.is_finished()) } pub fn url(&self) -> Option { self.running().then(|| self.url.clone()).flatten() } pub async fn start( &mut self, backend: SocketAddr, ca: LoadedCa, requested_port: u16, tab_mode: Arc>, ) -> Result<(String, u16)> { if let Some(url) = self.url() { return Ok((url, self.port.unwrap_or_default())); } let listener = bind_proxy_listener(requested_port).await?; let address = listener.local_addr()?; let (stop, done) = oneshot::channel(); let authority = RcgenAuthority::new(ca.issuer, 1_000, aws_lc_rs::default_provider()); let proxy = Proxy::builder() .with_listener(listener) .with_ca(authority) .with_rustls_connector(aws_lc_rs::default_provider()) .with_http_handler(CursorRelay { backend, tab_mode }) .with_graceful_shutdown(async move { let _ = done.await; }) .build() .map_err(|error| Error::Store(format!("build Cursor proxy: {error}")))?; self.stop = Some(stop); self.url = Some(format!("http://{address}")); self.port = Some(address.port()); self.task = Some(tokio::spawn(async move { if let Err(error) = proxy.start().await { tracing::error!(%error, "Cursor proxy stopped unexpectedly"); } })); Ok((self.url.clone().unwrap(), address.port())) } pub async fn stop(&mut self) { if let Some(stop) = self.stop.take() { let _ = stop.send(()); } if let Some(task) = self.task.take() { let _ = tokio::time::timeout(std::time::Duration::from_secs(5), task).await; } self.url = None; self.port = None; } } async fn bind_proxy_listener(requested_port: u16) -> Result { let requested = SocketAddr::from(([127, 0, 0, 1], requested_port)); match TcpListener::bind(requested).await { Ok(listener) => Ok(listener), Err(error) if requested_port != 0 => { tracing::warn!(%requested, %error, "configured proxy port unavailable; selecting a random port"); Ok(TcpListener::bind("127.0.0.1:0").await?) } Err(error) => Err(error.into()), } } #[derive(Clone)] struct CursorRelay { backend: SocketAddr, tab_mode: Arc>, } impl HttpHandler for CursorRelay { async fn handle_request( &mut self, _ctx: &HttpContext, mut request: Request, ) -> RequestOrResponse { let original = request.uri().clone(); let locally_routed = should_route_locally(original.path(), *self.tab_mode.read()); if is_cursor_host(original.host().unwrap_or_default()) && locally_routed { if let Ok(value) = original.to_string().parse() { request.headers_mut().insert(UPSTREAM_URL_HEADER, value); } let path = original .path_and_query() .map(|value| value.as_str()) .unwrap_or("/"); if let Ok(uri) = format!("http://{}{}", self.backend, path).parse::() { *request.uri_mut() = uri; } } request.into() } async fn should_intercept_connect( &mut self, _ctx: &HttpContext, request: &Request, ) -> bool { request .uri() .authority() .is_some_and(|authority| is_cursor_host(authority.host())) } async fn should_intercept_tls( &mut self, _ctx: &HttpContext, hello: hudsucker::rustls::server::ClientHello<'_>, ) -> bool { hello.server_name().is_some_and(is_cursor_host) } } pub fn is_cursor_host(host: &str) -> bool { let host = host.trim_end_matches('.').to_ascii_lowercase(); matches!(host.as_str(), "api2.cursor.sh" | "api3.cursor.sh") || host.ends_with(".cursor.sh") } fn is_local_path(path: &str) -> bool { matches!( path, "/agent.v1.AgentService/RunSSE" | "/aiserver.v1.BidiService/BidiAppend" | "/aiserver.v1.AiService/AvailableModels" | "/agent.v1.AgentService/GetUsableModels" | "/aiserver.v1.AiService/GetUsableModels" | "/aiserver.v1.AuthService/GetEmail" | "/aiserver.v1.DashboardService/GetMe" | "/aiserver.v1.DashboardService/GetTeams" | "/aiserver.v1.DashboardService/GetUserProfile" | "/aiserver.v1.DashboardService/GetCurrentPeriodUsage" | "/aiserver.v1.DashboardService/GetUsageLimitStatusAndActiveGrants" | "/aiserver.v1.AnalyticsService/BootstrapStatsig" | "/auth/full_stripe_profile" ) } fn should_route_locally(path: &str, tab_mode: TabMode) -> bool { is_local_path(path) || (is_tab_path(path) && tab_mode != TabMode::Direct) } #[cfg(test)] mod tests { use super::*; #[tokio::test] async fn proxy_listener_falls_back_when_configured_port_is_busy() { let occupied = TcpListener::bind("127.0.0.1:0").await.unwrap(); let requested_port = occupied.local_addr().unwrap().port(); let listener = bind_proxy_listener(requested_port).await.unwrap(); assert_ne!(listener.local_addr().unwrap().port(), requested_port); } #[test] fn limits_interception_to_cursor_hosts_and_local_paths() { assert!(is_cursor_host("api2.cursor.sh")); assert!(is_cursor_host("repo42.cursor.sh")); assert!(!is_cursor_host("example.com")); assert!(is_local_path("/agent.v1.AgentService/RunSSE")); assert!(is_local_path( "/aiserver.v1.AnalyticsService/BootstrapStatsig" )); assert!(!is_local_path("/unrelated")); assert!(should_route_locally( "/aiserver.v1.AiService/StreamCpp", TabMode::Public )); assert!(should_route_locally( "/aiserver.v1.AiService/StreamCpp", TabMode::Custom )); assert!(!should_route_locally( "/aiserver.v1.AiService/StreamCpp", TabMode::Direct )); } }