LCOV - code coverage report
Current view: top level - proxy/src/control_plane/client - cplane_proxy_v1.rs (source / functions) Coverage Total Hit
Test: 15f04989d2faf4ce76cecb56042184aca56ebae6.info Lines: 8.9 % 237 21
Test Date: 2025-07-14 11:50:36 Functions: 16.0 % 25 4

            Line data    Source code
       1              : //! Production console backend.
       2              : 
       3              : use std::net::IpAddr;
       4              : use std::str::FromStr;
       5              : use std::sync::Arc;
       6              : use std::time::Duration;
       7              : 
       8              : use ::http::HeaderName;
       9              : use ::http::header::AUTHORIZATION;
      10              : use bytes::Bytes;
      11              : use futures::TryFutureExt;
      12              : use hyper::StatusCode;
      13              : use postgres_client::config::SslMode;
      14              : use tokio::time::Instant;
      15              : use tracing::{Instrument, debug, info, info_span, warn};
      16              : 
      17              : use super::super::messages::{ControlPlaneErrorMessage, GetEndpointAccessControl, WakeCompute};
      18              : use crate::auth::backend::ComputeUserInfo;
      19              : use crate::auth::backend::jwt::AuthRule;
      20              : use crate::context::RequestContext;
      21              : use crate::control_plane::caches::ApiCaches;
      22              : use crate::control_plane::errors::{
      23              :     ControlPlaneError, GetAuthInfoError, GetEndpointJwksError, WakeComputeError,
      24              : };
      25              : use crate::control_plane::locks::ApiLocks;
      26              : use crate::control_plane::messages::{ColdStartInfo, EndpointJwksResponse, Reason};
      27              : use crate::control_plane::{
      28              :     AccessBlockerFlags, AuthInfo, AuthSecret, CachedNodeInfo, EndpointAccessControl, NodeInfo,
      29              :     RoleAccessControl,
      30              : };
      31              : use crate::metrics::Metrics;
      32              : use crate::rate_limiter::WakeComputeRateLimiter;
      33              : use crate::types::{EndpointCacheKey, EndpointId, RoleName};
      34              : use crate::{compute, http, scram};
      35              : 
      36              : pub(crate) const X_REQUEST_ID: HeaderName = HeaderName::from_static("x-request-id");
      37              : 
      38              : #[derive(Clone)]
      39              : pub struct NeonControlPlaneClient {
      40              :     endpoint: http::Endpoint,
      41              :     pub caches: &'static ApiCaches,
      42              :     pub(crate) locks: &'static ApiLocks<EndpointCacheKey>,
      43              :     pub(crate) wake_compute_endpoint_rate_limiter: Arc<WakeComputeRateLimiter>,
      44              :     // put in a shared ref so we don't copy secrets all over in memory
      45              :     jwt: Arc<str>,
      46              : }
      47              : 
      48              : impl NeonControlPlaneClient {
      49              :     /// Construct an API object containing the auth parameters.
      50            0 :     pub fn new(
      51            0 :         endpoint: http::Endpoint,
      52            0 :         jwt: Arc<str>,
      53            0 :         caches: &'static ApiCaches,
      54            0 :         locks: &'static ApiLocks<EndpointCacheKey>,
      55            0 :         wake_compute_endpoint_rate_limiter: Arc<WakeComputeRateLimiter>,
      56            0 :     ) -> Self {
      57            0 :         Self {
      58            0 :             endpoint,
      59            0 :             caches,
      60            0 :             locks,
      61            0 :             wake_compute_endpoint_rate_limiter,
      62            0 :             jwt,
      63            0 :         }
      64            0 :     }
      65              : 
      66            0 :     pub(crate) fn url(&self) -> &str {
      67            0 :         self.endpoint.url().as_str()
      68            0 :     }
      69              : 
      70            0 :     async fn do_get_auth_req(
      71            0 :         &self,
      72            0 :         ctx: &RequestContext,
      73            0 :         endpoint: &EndpointId,
      74            0 :         role: &RoleName,
      75            0 :     ) -> Result<AuthInfo, GetAuthInfoError> {
      76            0 :         async {
      77            0 :             let response = {
      78            0 :                 let request = self
      79            0 :                     .endpoint
      80            0 :                     .get_path("get_endpoint_access_control")
      81            0 :                     .header(X_REQUEST_ID, ctx.session_id().to_string())
      82            0 :                     .header(AUTHORIZATION, format!("Bearer {}", &self.jwt))
      83            0 :                     .query(&[("session_id", ctx.session_id())])
      84            0 :                     .query(&[
      85            0 :                         ("application_name", ctx.console_application_name().as_str()),
      86            0 :                         ("endpointish", endpoint.as_str()),
      87            0 :                         ("role", role.as_str()),
      88            0 :                     ])
      89            0 :                     .build()?;
      90              : 
      91            0 :                 debug!(url = request.url().as_str(), "sending http request");
      92            0 :                 let start = Instant::now();
      93            0 :                 let _pause = ctx.latency_timer_pause_at(start, crate::metrics::Waiting::Cplane);
      94            0 :                 let response = self.endpoint.execute(request).await?;
      95              : 
      96            0 :                 info!(duration = ?start.elapsed(), "received http response");
      97              : 
      98            0 :                 response
      99              :             };
     100              : 
     101            0 :             let body = match parse_body::<GetEndpointAccessControl>(
     102            0 :                 response.status(),
     103            0 :                 response.bytes().await?,
     104              :             ) {
     105            0 :                 Ok(body) => body,
     106              :                 // Error 404 is special: it's ok not to have a secret.
     107              :                 // TODO(anna): retry
     108            0 :                 Err(e) => {
     109            0 :                     return if e.get_reason().is_not_found() {
     110              :                         // TODO: refactor this because it's weird
     111              :                         // this is a failure to authenticate but we return Ok.
     112            0 :                         Ok(AuthInfo::default())
     113              :                     } else {
     114            0 :                         Err(e.into())
     115              :                     };
     116              :                 }
     117              :             };
     118              : 
     119            0 :             let secret = if body.role_secret.is_empty() {
     120            0 :                 None
     121              :             } else {
     122            0 :                 let secret = scram::ServerSecret::parse(&body.role_secret)
     123            0 :                     .map(AuthSecret::Scram)
     124            0 :                     .ok_or(GetAuthInfoError::BadSecret)?;
     125            0 :                 Some(secret)
     126              :             };
     127            0 :             let allowed_ips = body.allowed_ips.unwrap_or_default();
     128            0 :             Metrics::get()
     129            0 :                 .proxy
     130            0 :                 .allowed_ips_number
     131            0 :                 .observe(allowed_ips.len() as f64);
     132            0 :             let allowed_vpc_endpoint_ids = body.allowed_vpc_endpoint_ids.unwrap_or_default();
     133            0 :             Metrics::get()
     134            0 :                 .proxy
     135            0 :                 .allowed_vpc_endpoint_ids
     136            0 :                 .observe(allowed_vpc_endpoint_ids.len() as f64);
     137            0 :             let block_public_connections = body.block_public_connections.unwrap_or_default();
     138            0 :             let block_vpc_connections = body.block_vpc_connections.unwrap_or_default();
     139            0 :             Ok(AuthInfo {
     140            0 :                 secret,
     141            0 :                 allowed_ips,
     142            0 :                 allowed_vpc_endpoint_ids,
     143            0 :                 project_id: body.project_id,
     144            0 :                 account_id: body.account_id,
     145            0 :                 access_blocker_flags: AccessBlockerFlags {
     146            0 :                     public_access_blocked: block_public_connections,
     147            0 :                     vpc_access_blocked: block_vpc_connections,
     148            0 :                 },
     149            0 :                 rate_limits: body.rate_limits,
     150            0 :             })
     151            0 :         }
     152            0 :         .inspect_err(|e| tracing::debug!(error = ?e))
     153            0 :         .instrument(info_span!("do_get_auth_info"))
     154            0 :         .await
     155            0 :     }
     156              : 
     157            0 :     async fn do_get_endpoint_jwks(
     158            0 :         &self,
     159            0 :         ctx: &RequestContext,
     160            0 :         endpoint: &EndpointId,
     161            0 :     ) -> Result<Vec<AuthRule>, GetEndpointJwksError> {
     162            0 :         let request_id = ctx.session_id().to_string();
     163            0 :         async {
     164            0 :             let request = self
     165            0 :                 .endpoint
     166            0 :                 .get_with_url(|url| {
     167            0 :                     url.path_segments_mut()
     168            0 :                         .push("endpoints")
     169            0 :                         .push(endpoint.as_str())
     170            0 :                         .push("jwks");
     171            0 :                 })
     172            0 :                 .header(X_REQUEST_ID, &request_id)
     173            0 :                 .header(AUTHORIZATION, format!("Bearer {}", &self.jwt))
     174            0 :                 .query(&[("session_id", ctx.session_id())])
     175            0 :                 .build()
     176            0 :                 .map_err(GetEndpointJwksError::RequestBuild)?;
     177              : 
     178            0 :             debug!(url = request.url().as_str(), "sending http request");
     179            0 :             let start = Instant::now();
     180            0 :             let pause = ctx.latency_timer_pause(crate::metrics::Waiting::Cplane);
     181            0 :             let response = self
     182            0 :                 .endpoint
     183            0 :                 .execute(request)
     184            0 :                 .await
     185            0 :                 .map_err(GetEndpointJwksError::RequestExecute)?;
     186            0 :             drop(pause);
     187            0 :             info!(duration = ?start.elapsed(), "received http response");
     188              : 
     189            0 :             let body = parse_body::<EndpointJwksResponse>(
     190            0 :                 response.status(),
     191            0 :                 response.bytes().await.map_err(ControlPlaneError::from)?,
     192            0 :             )?;
     193              : 
     194            0 :             let rules = body
     195            0 :                 .jwks
     196            0 :                 .into_iter()
     197            0 :                 .map(|jwks| AuthRule {
     198            0 :                     id: jwks.id,
     199            0 :                     jwks_url: jwks.jwks_url,
     200            0 :                     audience: jwks.jwt_audience,
     201            0 :                     role_names: jwks.role_names,
     202            0 :                 })
     203            0 :                 .collect();
     204              : 
     205            0 :             Ok(rules)
     206            0 :         }
     207            0 :         .inspect_err(|e| tracing::debug!(error = ?e))
     208            0 :         .instrument(info_span!("do_get_endpoint_jwks"))
     209            0 :         .await
     210            0 :     }
     211              : 
     212            0 :     async fn do_wake_compute(
     213            0 :         &self,
     214            0 :         ctx: &RequestContext,
     215            0 :         user_info: &ComputeUserInfo,
     216            0 :     ) -> Result<NodeInfo, WakeComputeError> {
     217            0 :         let request_id = ctx.session_id().to_string();
     218            0 :         let application_name = ctx.console_application_name();
     219            0 :         async {
     220            0 :             let mut request_builder = self
     221            0 :                 .endpoint
     222            0 :                 .get_path("wake_compute")
     223            0 :                 .header("X-Request-ID", &request_id)
     224            0 :                 .header("Authorization", format!("Bearer {}", &self.jwt))
     225            0 :                 .query(&[("session_id", ctx.session_id())])
     226            0 :                 .query(&[
     227            0 :                     ("application_name", application_name.as_str()),
     228            0 :                     ("endpointish", user_info.endpoint.as_str()),
     229            0 :                 ]);
     230              : 
     231            0 :             let options = user_info.options.to_deep_object();
     232            0 :             if !options.is_empty() {
     233            0 :                 request_builder = request_builder.query(&options);
     234            0 :             }
     235              : 
     236            0 :             let request = request_builder.build()?;
     237              : 
     238            0 :             debug!(url = request.url().as_str(), "sending http request");
     239            0 :             let start = Instant::now();
     240            0 :             let pause = ctx.latency_timer_pause(crate::metrics::Waiting::Cplane);
     241            0 :             let response = self.endpoint.execute(request).await?;
     242            0 :             drop(pause);
     243            0 :             info!(duration = ?start.elapsed(), "received http response");
     244            0 :             let body = parse_body::<WakeCompute>(response.status(), response.bytes().await?)?;
     245              : 
     246            0 :             let Some((host, port)) = parse_host_port(&body.address) else {
     247            0 :                 return Err(WakeComputeError::BadComputeAddress(body.address));
     248              :             };
     249              : 
     250            0 :             let host_addr = IpAddr::from_str(host).ok();
     251              : 
     252            0 :             let ssl_mode = match &body.server_name {
     253            0 :                 Some(_) => SslMode::Require,
     254            0 :                 None => SslMode::Disable,
     255              :             };
     256            0 :             let host = match body.server_name {
     257            0 :                 Some(host) => host.into(),
     258            0 :                 None => host.into(),
     259              :             };
     260              : 
     261            0 :             let node = NodeInfo {
     262            0 :                 conn_info: compute::ConnectInfo {
     263            0 :                     host_addr,
     264            0 :                     host,
     265            0 :                     port,
     266            0 :                     ssl_mode,
     267            0 :                 },
     268            0 :                 aux: body.aux,
     269            0 :             };
     270              : 
     271            0 :             Ok(node)
     272            0 :         }
     273            0 :         .inspect_err(|e| tracing::debug!(error = ?e))
     274            0 :         .instrument(info_span!("do_wake_compute"))
     275            0 :         .await
     276            0 :     }
     277              : }
     278              : 
     279              : impl super::ControlPlaneApi for NeonControlPlaneClient {
     280              :     #[tracing::instrument(skip_all)]
     281              :     async fn get_role_access_control(
     282              :         &self,
     283              :         ctx: &RequestContext,
     284              :         endpoint: &EndpointId,
     285              :         role: &RoleName,
     286              :     ) -> Result<RoleAccessControl, crate::control_plane::errors::GetAuthInfoError> {
     287              :         let normalized_ep = &endpoint.normalize();
     288              :         if let Some(secret) = self
     289              :             .caches
     290              :             .project_info
     291              :             .get_role_secret(normalized_ep, role)
     292              :         {
     293              :             return Ok(secret);
     294              :         }
     295              : 
     296              :         let auth_info = self.do_get_auth_req(ctx, endpoint, role).await?;
     297              : 
     298              :         let control = EndpointAccessControl {
     299              :             allowed_ips: Arc::new(auth_info.allowed_ips),
     300              :             allowed_vpce: Arc::new(auth_info.allowed_vpc_endpoint_ids),
     301              :             flags: auth_info.access_blocker_flags,
     302              :             rate_limits: auth_info.rate_limits,
     303              :         };
     304              :         let role_control = RoleAccessControl {
     305              :             secret: auth_info.secret,
     306              :         };
     307              : 
     308              :         if let Some(project_id) = auth_info.project_id {
     309              :             let normalized_ep_int = normalized_ep.into();
     310              : 
     311              :             self.caches.project_info.insert_endpoint_access(
     312              :                 auth_info.account_id,
     313              :                 project_id,
     314              :                 normalized_ep_int,
     315              :                 role.into(),
     316              :                 control,
     317              :                 role_control.clone(),
     318              :             );
     319              :             ctx.set_project_id(project_id);
     320              :         }
     321              : 
     322              :         Ok(role_control)
     323              :     }
     324              : 
     325              :     #[tracing::instrument(skip_all)]
     326              :     async fn get_endpoint_access_control(
     327              :         &self,
     328              :         ctx: &RequestContext,
     329              :         endpoint: &EndpointId,
     330              :         role: &RoleName,
     331              :     ) -> Result<EndpointAccessControl, GetAuthInfoError> {
     332              :         let normalized_ep = &endpoint.normalize();
     333              :         if let Some(control) = self.caches.project_info.get_endpoint_access(normalized_ep) {
     334              :             return Ok(control);
     335              :         }
     336              : 
     337              :         let auth_info = self.do_get_auth_req(ctx, endpoint, role).await?;
     338              : 
     339              :         let control = EndpointAccessControl {
     340              :             allowed_ips: Arc::new(auth_info.allowed_ips),
     341              :             allowed_vpce: Arc::new(auth_info.allowed_vpc_endpoint_ids),
     342              :             flags: auth_info.access_blocker_flags,
     343              :             rate_limits: auth_info.rate_limits,
     344              :         };
     345              :         let role_control = RoleAccessControl {
     346              :             secret: auth_info.secret,
     347              :         };
     348              : 
     349              :         if let Some(project_id) = auth_info.project_id {
     350              :             let normalized_ep_int = normalized_ep.into();
     351              : 
     352              :             self.caches.project_info.insert_endpoint_access(
     353              :                 auth_info.account_id,
     354              :                 project_id,
     355              :                 normalized_ep_int,
     356              :                 role.into(),
     357              :                 control.clone(),
     358              :                 role_control,
     359              :             );
     360              :             ctx.set_project_id(project_id);
     361              :         }
     362              : 
     363              :         Ok(control)
     364              :     }
     365              : 
     366              :     #[tracing::instrument(skip_all)]
     367              :     async fn get_endpoint_jwks(
     368              :         &self,
     369              :         ctx: &RequestContext,
     370              :         endpoint: &EndpointId,
     371              :     ) -> Result<Vec<AuthRule>, GetEndpointJwksError> {
     372              :         self.do_get_endpoint_jwks(ctx, endpoint).await
     373              :     }
     374              : 
     375              :     #[tracing::instrument(skip_all)]
     376              :     async fn wake_compute(
     377              :         &self,
     378              :         ctx: &RequestContext,
     379              :         user_info: &ComputeUserInfo,
     380              :     ) -> Result<CachedNodeInfo, WakeComputeError> {
     381              :         let key = user_info.endpoint_cache_key();
     382              : 
     383              :         macro_rules! check_cache {
     384              :             () => {
     385              :                 if let Some(cached) = self.caches.node_info.get(&key) {
     386              :                     let (cached, info) = cached.take_value();
     387            0 :                     let info = info.map_err(|c| {
     388            0 :                         info!(key = &*key, "found cached wake_compute error");
     389            0 :                         WakeComputeError::ControlPlane(ControlPlaneError::Message(Box::new(*c)))
     390            0 :                     })?;
     391              : 
     392              :                     debug!(key = &*key, "found cached compute node info");
     393              :                     ctx.set_project(info.aux.clone());
     394              :                     return Ok(cached.map(|()| info));
     395              :                 }
     396              :             };
     397              :         }
     398              : 
     399              :         // Every time we do a wakeup http request, the compute node will stay up
     400              :         // for some time (highly depends on the console's scale-to-zero policy);
     401              :         // The connection info remains the same during that period of time,
     402              :         // which means that we might cache it to reduce the load and latency.
     403              :         check_cache!();
     404              : 
     405              :         let permit = self.locks.get_permit(&key).await?;
     406              : 
     407              :         // after getting back a permit - it's possible the cache was filled
     408              :         // double check
     409              :         if permit.should_check_cache() {
     410              :             // TODO: if there is something in the cache, mark the permit as success.
     411              :             check_cache!();
     412              :         }
     413              : 
     414              :         // check rate limit
     415              :         if !self
     416              :             .wake_compute_endpoint_rate_limiter
     417              :             .check(user_info.endpoint.normalize_intern(), 1)
     418              :         {
     419              :             return Err(WakeComputeError::TooManyConnections);
     420              :         }
     421              : 
     422              :         let node = permit.release_result(self.do_wake_compute(ctx, user_info).await);
     423              :         match node {
     424              :             Ok(node) => {
     425              :                 ctx.set_project(node.aux.clone());
     426              :                 debug!(key = &*key, "created a cache entry for woken compute node");
     427              : 
     428              :                 let mut stored_node = node.clone();
     429              :                 // store the cached node as 'warm_cached'
     430              :                 stored_node.aux.cold_start_info = ColdStartInfo::WarmCached;
     431              : 
     432              :                 let (_, cached) = self.caches.node_info.insert_unit(key, Ok(stored_node));
     433              : 
     434              :                 Ok(cached.map(|()| node))
     435              :             }
     436              :             Err(err) => match err {
     437              :                 WakeComputeError::ControlPlane(ControlPlaneError::Message(err)) => {
     438              :                     let Some(status) = &err.status else {
     439              :                         return Err(WakeComputeError::ControlPlane(ControlPlaneError::Message(
     440              :                             err,
     441              :                         )));
     442              :                     };
     443              : 
     444              :                     let reason = status
     445              :                         .details
     446              :                         .error_info
     447              :                         .map_or(Reason::Unknown, |x| x.reason);
     448              : 
     449              :                     // if we can retry this error, do not cache it.
     450              :                     if reason.can_retry() {
     451              :                         return Err(WakeComputeError::ControlPlane(ControlPlaneError::Message(
     452              :                             err,
     453              :                         )));
     454              :                     }
     455              : 
     456              :                     // at this point, we should only have quota errors.
     457              :                     debug!(
     458              :                         key = &*key,
     459              :                         "created a cache entry for the wake compute error"
     460              :                     );
     461              : 
     462              :                     self.caches.node_info.insert_ttl(
     463              :                         key,
     464              :                         Err(err.clone()),
     465              :                         Duration::from_secs(30),
     466              :                     );
     467              : 
     468              :                     Err(WakeComputeError::ControlPlane(ControlPlaneError::Message(
     469              :                         err,
     470              :                     )))
     471              :                 }
     472              :                 err => return Err(err),
     473              :             },
     474              :         }
     475              :     }
     476              : }
     477              : 
     478              : /// Parse http response body, taking status code into account.
     479            0 : fn parse_body<T: for<'a> serde::Deserialize<'a>>(
     480            0 :     status: StatusCode,
     481            0 :     body: Bytes,
     482            0 : ) -> Result<T, ControlPlaneError> {
     483            0 :     if status.is_success() {
     484              :         // We shouldn't log raw body because it may contain secrets.
     485            0 :         info!("request succeeded, processing the body");
     486            0 :         return Ok(serde_json::from_slice(&body).map_err(std::io::Error::other)?);
     487            0 :     }
     488              : 
     489              :     // Log plaintext to be able to detect, whether there are some cases not covered by the error struct.
     490            0 :     info!("response_error plaintext: {:?}", body);
     491              : 
     492              :     // Don't throw an error here because it's not as important
     493              :     // as the fact that the request itself has failed.
     494            0 :     let mut body = serde_json::from_slice(&body).unwrap_or_else(|e| {
     495            0 :         warn!("failed to parse error body: {e}");
     496            0 :         Box::new(ControlPlaneErrorMessage {
     497            0 :             error: "reason unclear (malformed error message)".into(),
     498            0 :             http_status_code: status,
     499            0 :             status: None,
     500            0 :         })
     501            0 :     });
     502            0 :     body.http_status_code = status;
     503              : 
     504            0 :     warn!("console responded with an error ({status}): {body:?}");
     505            0 :     Err(ControlPlaneError::Message(body))
     506            0 : }
     507              : 
     508            3 : fn parse_host_port(input: &str) -> Option<(&str, u16)> {
     509            3 :     let (host, port) = input.rsplit_once(':')?;
     510            3 :     let ipv6_brackets: &[_] = &['[', ']'];
     511            3 :     Some((host.trim_matches(ipv6_brackets), port.parse().ok()?))
     512            3 : }
     513              : 
     514              : #[cfg(test)]
     515              : mod tests {
     516              :     use super::*;
     517              : 
     518              :     #[test]
     519            1 :     fn test_parse_host_port_v4() {
     520            1 :         let (host, port) = parse_host_port("127.0.0.1:5432").expect("failed to parse");
     521            1 :         assert_eq!(host, "127.0.0.1");
     522            1 :         assert_eq!(port, 5432);
     523            1 :     }
     524              : 
     525              :     #[test]
     526            1 :     fn test_parse_host_port_v6() {
     527            1 :         let (host, port) = parse_host_port("[2001:db8::1]:5432").expect("failed to parse");
     528            1 :         assert_eq!(host, "2001:db8::1");
     529            1 :         assert_eq!(port, 5432);
     530            1 :     }
     531              : 
     532              :     #[test]
     533            1 :     fn test_parse_host_port_url() {
     534            1 :         let (host, port) = parse_host_port("compute-foo-bar-1234.default.svc.cluster.local:5432")
     535            1 :             .expect("failed to parse");
     536            1 :         assert_eq!(host, "compute-foo-bar-1234.default.svc.cluster.local");
     537            1 :         assert_eq!(port, 5432);
     538            1 :     }
     539              : }
        

Generated by: LCOV version 2.1-beta