LCOV - code coverage report
Current view: top level - proxy/src/control_plane/client - mock.rs (source / functions) Coverage Total Hit
Test: 07bee600374ccd486c69370d0972d9035964fe68.info Lines: 0.0 % 165 0
Test Date: 2025-02-20 13:11:02 Functions: 0.0 % 26 0

            Line data    Source code
       1              : //! Mock console backend which relies on a user-provided postgres instance.
       2              : 
       3              : use std::str::FromStr;
       4              : use std::sync::Arc;
       5              : 
       6              : use futures::TryFutureExt;
       7              : use thiserror::Error;
       8              : use tokio_postgres::Client;
       9              : use tracing::{error, info, info_span, warn, Instrument};
      10              : 
      11              : use crate::auth::backend::jwt::AuthRule;
      12              : use crate::auth::backend::ComputeUserInfo;
      13              : use crate::auth::IpPattern;
      14              : use crate::cache::Cached;
      15              : use crate::context::RequestContext;
      16              : use crate::control_plane::client::{
      17              :     CachedAllowedIps, CachedAllowedVpcEndpointIds, CachedRoleSecret,
      18              : };
      19              : use crate::control_plane::errors::{
      20              :     ControlPlaneError, GetAuthInfoError, GetEndpointJwksError, WakeComputeError,
      21              : };
      22              : use crate::control_plane::messages::MetricsAuxInfo;
      23              : use crate::control_plane::{AccessBlockerFlags, AuthInfo, AuthSecret, CachedNodeInfo, NodeInfo};
      24              : use crate::error::io_error;
      25              : use crate::intern::RoleNameInt;
      26              : use crate::types::{BranchId, EndpointId, ProjectId, RoleName};
      27              : use crate::url::ApiUrl;
      28              : use crate::{compute, scram};
      29              : 
      30              : #[derive(Debug, Error)]
      31              : enum MockApiError {
      32              :     #[error("Failed to read password: {0}")]
      33              :     PasswordNotSet(tokio_postgres::Error),
      34              : }
      35              : 
      36              : impl From<MockApiError> for ControlPlaneError {
      37            0 :     fn from(e: MockApiError) -> Self {
      38            0 :         io_error(e).into()
      39            0 :     }
      40              : }
      41              : 
      42              : impl From<tokio_postgres::Error> for ControlPlaneError {
      43            0 :     fn from(e: tokio_postgres::Error) -> Self {
      44            0 :         io_error(e).into()
      45            0 :     }
      46              : }
      47              : 
      48              : #[derive(Clone)]
      49              : pub struct MockControlPlane {
      50              :     endpoint: ApiUrl,
      51              :     ip_allowlist_check_enabled: bool,
      52              : }
      53              : 
      54              : impl MockControlPlane {
      55            0 :     pub fn new(endpoint: ApiUrl, ip_allowlist_check_enabled: bool) -> Self {
      56            0 :         Self {
      57            0 :             endpoint,
      58            0 :             ip_allowlist_check_enabled,
      59            0 :         }
      60            0 :     }
      61              : 
      62            0 :     pub(crate) fn url(&self) -> &str {
      63            0 :         self.endpoint.as_str()
      64            0 :     }
      65              : 
      66            0 :     async fn do_get_auth_info(
      67            0 :         &self,
      68            0 :         user_info: &ComputeUserInfo,
      69            0 :     ) -> Result<AuthInfo, GetAuthInfoError> {
      70            0 :         let (secret, allowed_ips) = async {
      71              :             // Perhaps we could persist this connection, but then we'd have to
      72              :             // write more code for reopening it if it got closed, which doesn't
      73              :             // seem worth it.
      74            0 :             let (client, connection) =
      75            0 :                 tokio_postgres::connect(self.endpoint.as_str(), tokio_postgres::NoTls).await?;
      76              : 
      77            0 :             tokio::spawn(connection);
      78              : 
      79            0 :             let secret = if let Some(entry) = get_execute_postgres_query(
      80            0 :                 &client,
      81            0 :                 "select rolpassword from pg_catalog.pg_authid where rolname = $1",
      82            0 :                 &[&&*user_info.user],
      83            0 :                 "rolpassword",
      84            0 :             )
      85            0 :             .await?
      86              :             {
      87            0 :                 info!("got a secret: {entry}"); // safe since it's not a prod scenario
      88            0 :                 let secret = scram::ServerSecret::parse(&entry).map(AuthSecret::Scram);
      89            0 :                 secret.or_else(|| parse_md5(&entry).map(AuthSecret::Md5))
      90              :             } else {
      91            0 :                 warn!("user '{}' does not exist", user_info.user);
      92            0 :                 None
      93              :             };
      94              : 
      95            0 :             let allowed_ips = if self.ip_allowlist_check_enabled {
      96            0 :                 match get_execute_postgres_query(
      97            0 :                     &client,
      98            0 :                     "select allowed_ips from neon_control_plane.endpoints where endpoint_id = $1",
      99            0 :                     &[&user_info.endpoint.as_str()],
     100            0 :                     "allowed_ips",
     101            0 :                 )
     102            0 :                 .await?
     103              :                 {
     104            0 :                     Some(s) => {
     105            0 :                         info!("got allowed_ips: {s}");
     106            0 :                         s.split(',')
     107            0 :                             .map(|s| {
     108            0 :                                 IpPattern::from_str(s).expect("mocked ip pattern should be correct")
     109            0 :                             })
     110            0 :                             .collect()
     111              :                     }
     112            0 :                     None => vec![],
     113              :                 }
     114              :             } else {
     115            0 :                 vec![]
     116              :             };
     117              : 
     118            0 :             Ok((secret, allowed_ips))
     119            0 :         }
     120            0 :         .inspect_err(|e: &GetAuthInfoError| tracing::error!("{e}"))
     121            0 :         .instrument(info_span!("postgres", url = self.endpoint.as_str()))
     122            0 :         .await?;
     123            0 :         Ok(AuthInfo {
     124            0 :             secret,
     125            0 :             allowed_ips,
     126            0 :             allowed_vpc_endpoint_ids: vec![],
     127            0 :             project_id: None,
     128            0 :             account_id: None,
     129            0 :             access_blocker_flags: AccessBlockerFlags::default(),
     130            0 :         })
     131            0 :     }
     132              : 
     133            0 :     async fn do_get_endpoint_jwks(
     134            0 :         &self,
     135            0 :         endpoint: EndpointId,
     136            0 :     ) -> Result<Vec<AuthRule>, GetEndpointJwksError> {
     137            0 :         let (client, connection) =
     138            0 :             tokio_postgres::connect(self.endpoint.as_str(), tokio_postgres::NoTls).await?;
     139              : 
     140            0 :         let connection = tokio::spawn(connection);
     141              : 
     142            0 :         let res = client.query(
     143            0 :                 "select id, jwks_url, audience, role_names from neon_control_plane.endpoint_jwks where endpoint_id = $1",
     144            0 :                 &[&endpoint.as_str()],
     145            0 :             )
     146            0 :             .await?;
     147              : 
     148            0 :         let mut rows = vec![];
     149            0 :         for row in res {
     150            0 :             rows.push(AuthRule {
     151            0 :                 id: row.get("id"),
     152            0 :                 jwks_url: url::Url::parse(row.get("jwks_url"))?,
     153            0 :                 audience: row.get("audience"),
     154            0 :                 role_names: row
     155            0 :                     .get::<_, Vec<String>>("role_names")
     156            0 :                     .into_iter()
     157            0 :                     .map(RoleName::from)
     158            0 :                     .map(|s| RoleNameInt::from(&s))
     159            0 :                     .collect(),
     160            0 :             });
     161            0 :         }
     162              : 
     163            0 :         drop(client);
     164            0 :         connection.await??;
     165              : 
     166            0 :         Ok(rows)
     167            0 :     }
     168              : 
     169            0 :     async fn do_wake_compute(&self) -> Result<NodeInfo, WakeComputeError> {
     170            0 :         let mut config = compute::ConnCfg::new(
     171            0 :             self.endpoint.host_str().unwrap_or("localhost").to_owned(),
     172            0 :             self.endpoint.port().unwrap_or(5432),
     173            0 :         );
     174            0 :         config.ssl_mode(postgres_client::config::SslMode::Disable);
     175            0 : 
     176            0 :         let node = NodeInfo {
     177            0 :             config,
     178            0 :             aux: MetricsAuxInfo {
     179            0 :                 endpoint_id: (&EndpointId::from("endpoint")).into(),
     180            0 :                 project_id: (&ProjectId::from("project")).into(),
     181            0 :                 branch_id: (&BranchId::from("branch")).into(),
     182            0 :                 cold_start_info: crate::control_plane::messages::ColdStartInfo::Warm,
     183            0 :             },
     184            0 :         };
     185            0 : 
     186            0 :         Ok(node)
     187            0 :     }
     188              : }
     189              : 
     190            0 : async fn get_execute_postgres_query(
     191            0 :     client: &Client,
     192            0 :     query: &str,
     193            0 :     params: &[&(dyn tokio_postgres::types::ToSql + Sync)],
     194            0 :     idx: &str,
     195            0 : ) -> Result<Option<String>, GetAuthInfoError> {
     196            0 :     let rows = client.query(query, params).await?;
     197              : 
     198              :     // We can get at most one row, because `rolname` is unique.
     199            0 :     let Some(row) = rows.first() else {
     200              :         // This means that the user doesn't exist, so there can be no secret.
     201              :         // However, this is still a *valid* outcome which is very similar
     202              :         // to getting `404 Not found` from the Neon console.
     203            0 :         return Ok(None);
     204              :     };
     205              : 
     206            0 :     let entry = row.try_get(idx).map_err(MockApiError::PasswordNotSet)?;
     207            0 :     Ok(Some(entry))
     208            0 : }
     209              : 
     210              : impl super::ControlPlaneApi for MockControlPlane {
     211              :     #[tracing::instrument(skip_all)]
     212              :     async fn get_role_secret(
     213              :         &self,
     214              :         _ctx: &RequestContext,
     215              :         user_info: &ComputeUserInfo,
     216              :     ) -> Result<CachedRoleSecret, GetAuthInfoError> {
     217              :         Ok(CachedRoleSecret::new_uncached(
     218              :             self.do_get_auth_info(user_info).await?.secret,
     219              :         ))
     220              :     }
     221              : 
     222            0 :     async fn get_allowed_ips(
     223            0 :         &self,
     224            0 :         _ctx: &RequestContext,
     225            0 :         user_info: &ComputeUserInfo,
     226            0 :     ) -> Result<CachedAllowedIps, GetAuthInfoError> {
     227            0 :         Ok(Cached::new_uncached(Arc::new(
     228            0 :             self.do_get_auth_info(user_info).await?.allowed_ips,
     229              :         )))
     230            0 :     }
     231              : 
     232            0 :     async fn get_allowed_vpc_endpoint_ids(
     233            0 :         &self,
     234            0 :         _ctx: &RequestContext,
     235            0 :         user_info: &ComputeUserInfo,
     236            0 :     ) -> Result<CachedAllowedVpcEndpointIds, super::errors::GetAuthInfoError> {
     237            0 :         Ok(Cached::new_uncached(Arc::new(
     238            0 :             self.do_get_auth_info(user_info)
     239            0 :                 .await?
     240              :                 .allowed_vpc_endpoint_ids,
     241              :         )))
     242            0 :     }
     243              : 
     244            0 :     async fn get_block_public_or_vpc_access(
     245            0 :         &self,
     246            0 :         _ctx: &RequestContext,
     247            0 :         user_info: &ComputeUserInfo,
     248            0 :     ) -> Result<super::CachedAccessBlockerFlags, super::errors::GetAuthInfoError> {
     249            0 :         Ok(Cached::new_uncached(
     250            0 :             self.do_get_auth_info(user_info).await?.access_blocker_flags,
     251              :         ))
     252            0 :     }
     253              : 
     254            0 :     async fn get_endpoint_jwks(
     255            0 :         &self,
     256            0 :         _ctx: &RequestContext,
     257            0 :         endpoint: EndpointId,
     258            0 :     ) -> Result<Vec<AuthRule>, GetEndpointJwksError> {
     259            0 :         self.do_get_endpoint_jwks(endpoint).await
     260            0 :     }
     261              : 
     262              :     #[tracing::instrument(skip_all)]
     263              :     async fn wake_compute(
     264              :         &self,
     265              :         _ctx: &RequestContext,
     266              :         _user_info: &ComputeUserInfo,
     267              :     ) -> Result<CachedNodeInfo, WakeComputeError> {
     268              :         self.do_wake_compute().map_ok(Cached::new_uncached).await
     269              :     }
     270              : }
     271              : 
     272            0 : fn parse_md5(input: &str) -> Option<[u8; 16]> {
     273            0 :     let text = input.strip_prefix("md5")?;
     274              : 
     275            0 :     let mut bytes = [0u8; 16];
     276            0 :     hex::decode_to_slice(text, &mut bytes).ok()?;
     277              : 
     278            0 :     Some(bytes)
     279            0 : }
        

Generated by: LCOV version 2.1-beta