LCOV - code coverage report
Current view: top level - proxy/src/control_plane - messages.rs (source / functions) Coverage Total Hit
Test: 1b0a6a0c05cee5a7de360813c8034804e105ce1c.info Lines: 65.5 % 203 133
Test Date: 2025-03-12 00:01:28 Functions: 15.2 % 132 20

            Line data    Source code
       1              : use std::fmt::{self, Display};
       2              : 
       3              : use measured::FixedCardinalityLabel;
       4              : use serde::{Deserialize, Serialize};
       5              : use smol_str::SmolStr;
       6              : 
       7              : use crate::auth::IpPattern;
       8              : use crate::intern::{AccountIdInt, BranchIdInt, EndpointIdInt, ProjectIdInt, RoleNameInt};
       9              : use crate::proxy::retry::CouldRetry;
      10              : 
      11              : /// Generic error response with human-readable description.
      12              : /// Note that we can't always present it to user as is.
      13            0 : #[derive(Debug, Deserialize, Clone)]
      14              : pub(crate) struct ControlPlaneErrorMessage {
      15              :     pub(crate) error: Box<str>,
      16              :     #[serde(skip)]
      17              :     pub(crate) http_status_code: http::StatusCode,
      18              :     pub(crate) status: Option<Status>,
      19              : }
      20              : 
      21              : impl ControlPlaneErrorMessage {
      22            3 :     pub(crate) fn get_reason(&self) -> Reason {
      23            3 :         self.status
      24            3 :             .as_ref()
      25            3 :             .and_then(|s| s.details.error_info.as_ref())
      26            3 :             .map_or(Reason::Unknown, |e| e.reason)
      27            3 :     }
      28              : 
      29            0 :     pub(crate) fn get_user_facing_message(&self) -> String {
      30              :         use super::errors::REQUEST_FAILED;
      31            0 :         self.status
      32            0 :             .as_ref()
      33            0 :             .and_then(|s| s.details.user_facing_message.as_ref())
      34            0 :             .map_or_else(|| {
      35            0 :                 // Ask @neondatabase/control-plane for review before adding more.
      36            0 :                 match self.http_status_code {
      37              :                     http::StatusCode::NOT_FOUND => {
      38              :                         // Status 404: failed to get a project-related resource.
      39            0 :                         format!("{REQUEST_FAILED}: endpoint cannot be found")
      40              :                     }
      41              :                     http::StatusCode::NOT_ACCEPTABLE => {
      42              :                         // Status 406: endpoint is disabled (we don't allow connections).
      43            0 :                         format!("{REQUEST_FAILED}: endpoint is disabled")
      44              :                     }
      45              :                     http::StatusCode::LOCKED | http::StatusCode::UNPROCESSABLE_ENTITY => {
      46              :                         // Status 423: project might be in maintenance mode (or bad state), or quotas exceeded.
      47            0 :                         format!("{REQUEST_FAILED}: endpoint is temporarily unavailable. Check your quotas and/or contact our support.")
      48              :                     }
      49            0 :                     _ => REQUEST_FAILED.to_owned(),
      50              :                 }
      51            0 :             }, |m| m.message.clone().into())
      52            0 :     }
      53              : }
      54              : 
      55              : impl Display for ControlPlaneErrorMessage {
      56            0 :     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
      57            0 :         let msg: &str = self
      58            0 :             .status
      59            0 :             .as_ref()
      60            0 :             .and_then(|s| s.details.user_facing_message.as_ref())
      61            0 :             .map_or_else(|| self.error.as_ref(), |m| m.message.as_ref());
      62            0 :         write!(f, "{msg}")
      63            0 :     }
      64              : }
      65              : 
      66              : impl CouldRetry for ControlPlaneErrorMessage {
      67            6 :     fn could_retry(&self) -> bool {
      68              :         // If the error message does not have a status,
      69              :         // the error is unknown and probably should not retry automatically
      70            6 :         let Some(status) = &self.status else {
      71            2 :             return false;
      72              :         };
      73              : 
      74              :         // retry if the retry info is set.
      75            4 :         if status.details.retry_info.is_some() {
      76            4 :             return true;
      77            0 :         }
      78            0 : 
      79            0 :         // if no retry info set, attempt to use the error code to guess the retry state.
      80            0 :         let reason = status
      81            0 :             .details
      82            0 :             .error_info
      83            0 :             .map_or(Reason::Unknown, |e| e.reason);
      84            0 : 
      85            0 :         reason.can_retry()
      86            6 :     }
      87              : }
      88              : 
      89            0 : #[derive(Debug, Deserialize, Clone)]
      90              : #[allow(dead_code)]
      91              : pub(crate) struct Status {
      92              :     pub(crate) code: Box<str>,
      93              :     pub(crate) message: Box<str>,
      94              :     pub(crate) details: Details,
      95              : }
      96              : 
      97            0 : #[derive(Debug, Deserialize, Clone)]
      98              : pub(crate) struct Details {
      99              :     pub(crate) error_info: Option<ErrorInfo>,
     100              :     pub(crate) retry_info: Option<RetryInfo>,
     101              :     pub(crate) user_facing_message: Option<UserFacingMessage>,
     102              : }
     103              : 
     104            0 : #[derive(Copy, Clone, Debug, Deserialize)]
     105              : pub(crate) struct ErrorInfo {
     106              :     pub(crate) reason: Reason,
     107              :     // Schema could also have `metadata` field, but it's not structured. Skip it for now.
     108              : }
     109              : 
     110            0 : #[derive(Clone, Copy, Debug, Deserialize, Default)]
     111              : pub(crate) enum Reason {
     112              :     /// RoleProtected indicates that the role is protected and the attempted operation is not permitted on protected roles.
     113              :     #[serde(rename = "ROLE_PROTECTED")]
     114              :     RoleProtected,
     115              :     /// ResourceNotFound indicates that a resource (project, endpoint, branch, etc.) wasn't found,
     116              :     /// usually due to the provided ID not being correct or because the subject doesn't have enough permissions to
     117              :     /// access the requested resource.
     118              :     /// Prefer a more specific reason if possible, e.g., ProjectNotFound, EndpointNotFound, etc.
     119              :     #[serde(rename = "RESOURCE_NOT_FOUND")]
     120              :     ResourceNotFound,
     121              :     /// ProjectNotFound indicates that the project wasn't found, usually due to the provided ID not being correct,
     122              :     /// or that the subject doesn't have enough permissions to access the requested project.
     123              :     #[serde(rename = "PROJECT_NOT_FOUND")]
     124              :     ProjectNotFound,
     125              :     /// EndpointNotFound indicates that the endpoint wasn't found, usually due to the provided ID not being correct,
     126              :     /// or that the subject doesn't have enough permissions to access the requested endpoint.
     127              :     #[serde(rename = "ENDPOINT_NOT_FOUND")]
     128              :     EndpointNotFound,
     129              :     /// BranchNotFound indicates that the branch wasn't found, usually due to the provided ID not being correct,
     130              :     /// or that the subject doesn't have enough permissions to access the requested branch.
     131              :     #[serde(rename = "BRANCH_NOT_FOUND")]
     132              :     BranchNotFound,
     133              :     /// RateLimitExceeded indicates that the rate limit for the operation has been exceeded.
     134              :     #[serde(rename = "RATE_LIMIT_EXCEEDED")]
     135              :     RateLimitExceeded,
     136              :     /// NonDefaultBranchComputeTimeExceeded indicates that the compute time quota of non-default branches has been
     137              :     /// exceeded.
     138              :     #[serde(rename = "NON_PRIMARY_BRANCH_COMPUTE_TIME_EXCEEDED")]
     139              :     NonDefaultBranchComputeTimeExceeded,
     140              :     /// ActiveTimeQuotaExceeded indicates that the active time quota was exceeded.
     141              :     #[serde(rename = "ACTIVE_TIME_QUOTA_EXCEEDED")]
     142              :     ActiveTimeQuotaExceeded,
     143              :     /// ComputeTimeQuotaExceeded indicates that the compute time quota was exceeded.
     144              :     #[serde(rename = "COMPUTE_TIME_QUOTA_EXCEEDED")]
     145              :     ComputeTimeQuotaExceeded,
     146              :     /// WrittenDataQuotaExceeded indicates that the written data quota was exceeded.
     147              :     #[serde(rename = "WRITTEN_DATA_QUOTA_EXCEEDED")]
     148              :     WrittenDataQuotaExceeded,
     149              :     /// DataTransferQuotaExceeded indicates that the data transfer quota was exceeded.
     150              :     #[serde(rename = "DATA_TRANSFER_QUOTA_EXCEEDED")]
     151              :     DataTransferQuotaExceeded,
     152              :     /// LogicalSizeQuotaExceeded indicates that the logical size quota was exceeded.
     153              :     #[serde(rename = "LOGICAL_SIZE_QUOTA_EXCEEDED")]
     154              :     LogicalSizeQuotaExceeded,
     155              :     /// RunningOperations indicates that the project already has some running operations
     156              :     /// and scheduling of new ones is prohibited.
     157              :     #[serde(rename = "RUNNING_OPERATIONS")]
     158              :     RunningOperations,
     159              :     /// ConcurrencyLimitReached indicates that the concurrency limit for an action was reached.
     160              :     #[serde(rename = "CONCURRENCY_LIMIT_REACHED")]
     161              :     ConcurrencyLimitReached,
     162              :     /// LockAlreadyTaken indicates that the we attempted to take a lock that was already taken.
     163              :     #[serde(rename = "LOCK_ALREADY_TAKEN")]
     164              :     LockAlreadyTaken,
     165              :     /// ActiveEndpointsLimitExceeded indicates that the limit of concurrently active endpoints was exceeded.
     166              :     #[serde(rename = "ACTIVE_ENDPOINTS_LIMIT_EXCEEDED")]
     167              :     ActiveEndpointsLimitExceeded,
     168              :     #[default]
     169              :     #[serde(other)]
     170              :     Unknown,
     171              : }
     172              : 
     173              : impl Reason {
     174            0 :     pub(crate) fn is_not_found(self) -> bool {
     175            0 :         matches!(
     176            0 :             self,
     177              :             Reason::ResourceNotFound
     178              :                 | Reason::ProjectNotFound
     179              :                 | Reason::EndpointNotFound
     180              :                 | Reason::BranchNotFound
     181              :         )
     182            0 :     }
     183              : 
     184            0 :     pub(crate) fn can_retry(self) -> bool {
     185            0 :         match self {
     186              :             // do not retry role protected errors
     187              :             // not a transitive error
     188            0 :             Reason::RoleProtected => false,
     189              :             // on retry, it will still not be found
     190              :             Reason::ResourceNotFound
     191              :             | Reason::ProjectNotFound
     192              :             | Reason::EndpointNotFound
     193            0 :             | Reason::BranchNotFound => false,
     194              :             // we were asked to go away
     195              :             Reason::RateLimitExceeded
     196              :             | Reason::NonDefaultBranchComputeTimeExceeded
     197              :             | Reason::ActiveTimeQuotaExceeded
     198              :             | Reason::ComputeTimeQuotaExceeded
     199              :             | Reason::WrittenDataQuotaExceeded
     200              :             | Reason::DataTransferQuotaExceeded
     201              :             | Reason::LogicalSizeQuotaExceeded
     202            0 :             | Reason::ActiveEndpointsLimitExceeded => false,
     203              :             // transitive error. control plane is currently busy
     204              :             // but might be ready soon
     205              :             Reason::RunningOperations
     206              :             | Reason::ConcurrencyLimitReached
     207            0 :             | Reason::LockAlreadyTaken => true,
     208              :             // unknown error. better not retry it.
     209            0 :             Reason::Unknown => false,
     210              :         }
     211            0 :     }
     212              : }
     213              : 
     214            0 : #[derive(Copy, Clone, Debug, Deserialize)]
     215              : #[allow(dead_code)]
     216              : pub(crate) struct RetryInfo {
     217              :     pub(crate) retry_delay_ms: u64,
     218              : }
     219              : 
     220            0 : #[derive(Debug, Deserialize, Clone)]
     221              : pub(crate) struct UserFacingMessage {
     222              :     pub(crate) message: Box<str>,
     223              : }
     224              : 
     225              : /// Response which holds client's auth secret, e.g. [`crate::scram::ServerSecret`].
     226              : /// Returned by the `/get_endpoint_access_control` API method.
     227           12 : #[derive(Deserialize)]
     228              : pub(crate) struct GetEndpointAccessControl {
     229              :     pub(crate) role_secret: Box<str>,
     230              :     pub(crate) allowed_ips: Option<Vec<IpPattern>>,
     231              :     pub(crate) allowed_vpc_endpoint_ids: Option<Vec<String>>,
     232              :     pub(crate) project_id: Option<ProjectIdInt>,
     233              :     pub(crate) account_id: Option<AccountIdInt>,
     234              :     pub(crate) block_public_connections: Option<bool>,
     235              :     pub(crate) block_vpc_connections: Option<bool>,
     236              : }
     237              : 
     238              : /// Response which holds compute node's `host:port` pair.
     239              : /// Returned by the `/proxy_wake_compute` API method.
     240            2 : #[derive(Debug, Deserialize)]
     241              : pub(crate) struct WakeCompute {
     242              :     pub(crate) address: Box<str>,
     243              :     pub(crate) server_name: Option<String>,
     244              :     pub(crate) aux: MetricsAuxInfo,
     245              : }
     246              : 
     247              : /// Async response which concludes the console redirect auth flow.
     248              : /// Also known as `kickResponse` in the console.
     249            2 : #[derive(Debug, Deserialize)]
     250              : pub(crate) struct KickSession<'a> {
     251              :     /// Session ID is assigned by the proxy.
     252              :     pub(crate) session_id: &'a str,
     253              : 
     254              :     /// Compute node connection params.
     255              :     #[serde(deserialize_with = "KickSession::parse_db_info")]
     256              :     pub(crate) result: DatabaseInfo,
     257              : }
     258              : 
     259              : impl KickSession<'_> {
     260            1 :     fn parse_db_info<'de, D>(des: D) -> Result<DatabaseInfo, D::Error>
     261            1 :     where
     262            1 :         D: serde::Deserializer<'de>,
     263            1 :     {
     264            1 :         #[derive(Deserialize)]
     265              :         enum Wrapper {
     266              :             // Currently, console only reports `Success`.
     267              :             // `Failure(String)` used to be here... RIP.
     268              :             Success(DatabaseInfo),
     269              :         }
     270              : 
     271            1 :         Wrapper::deserialize(des).map(|x| match x {
     272            1 :             Wrapper::Success(info) => info,
     273            1 :         })
     274            1 :     }
     275              : }
     276              : 
     277              : /// Compute node connection params.
     278           31 : #[derive(Deserialize)]
     279              : pub(crate) struct DatabaseInfo {
     280              :     pub(crate) host: Box<str>,
     281              :     pub(crate) port: u16,
     282              :     pub(crate) dbname: Box<str>,
     283              :     pub(crate) user: Box<str>,
     284              :     /// Console always provides a password, but it might
     285              :     /// be inconvenient for debug with local PG instance.
     286              :     pub(crate) password: Option<Box<str>>,
     287              :     pub(crate) aux: MetricsAuxInfo,
     288              :     #[serde(default)]
     289              :     pub(crate) allowed_ips: Option<Vec<IpPattern>>,
     290              :     #[serde(default)]
     291              :     pub(crate) allowed_vpc_endpoint_ids: Option<Vec<String>>,
     292              :     #[serde(default)]
     293              :     pub(crate) public_access_allowed: Option<bool>,
     294              : }
     295              : 
     296              : // Manually implement debug to omit sensitive info.
     297              : impl fmt::Debug for DatabaseInfo {
     298            0 :     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
     299            0 :         f.debug_struct("DatabaseInfo")
     300            0 :             .field("host", &self.host)
     301            0 :             .field("port", &self.port)
     302            0 :             .field("dbname", &self.dbname)
     303            0 :             .field("user", &self.user)
     304            0 :             .field("allowed_ips", &self.allowed_ips)
     305            0 :             .field("allowed_vpc_endpoint_ids", &self.allowed_vpc_endpoint_ids)
     306            0 :             .finish_non_exhaustive()
     307            0 :     }
     308              : }
     309              : 
     310              : /// Various labels for prometheus metrics.
     311              : /// Also known as `ProxyMetricsAuxInfo` in the console.
     312           30 : #[derive(Debug, Deserialize, Clone)]
     313              : pub(crate) struct MetricsAuxInfo {
     314              :     pub(crate) endpoint_id: EndpointIdInt,
     315              :     pub(crate) project_id: ProjectIdInt,
     316              :     pub(crate) branch_id: BranchIdInt,
     317              :     // note: we don't use interned strings for compute IDs.
     318              :     // they churn too quickly and we have no way to clean up interned strings.
     319              :     pub(crate) compute_id: SmolStr,
     320              :     #[serde(default)]
     321              :     pub(crate) cold_start_info: ColdStartInfo,
     322              : }
     323              : 
     324            6 : #[derive(Debug, Default, Serialize, Deserialize, Clone, Copy, FixedCardinalityLabel)]
     325              : #[serde(rename_all = "snake_case")]
     326              : pub enum ColdStartInfo {
     327              :     #[default]
     328              :     Unknown,
     329              :     /// Compute was already running
     330              :     Warm,
     331              :     #[serde(rename = "pool_hit")]
     332              :     #[label(rename = "pool_hit")]
     333              :     /// Compute was not running but there was an available VM
     334              :     VmPoolHit,
     335              :     #[serde(rename = "pool_miss")]
     336              :     #[label(rename = "pool_miss")]
     337              :     /// Compute was not running and there were no VMs available
     338              :     VmPoolMiss,
     339              : 
     340              :     // not provided by control plane
     341              :     /// Connection available from HTTP pool
     342              :     HttpPoolHit,
     343              :     /// Cached connection info
     344              :     WarmCached,
     345              : }
     346              : 
     347              : impl ColdStartInfo {
     348            0 :     pub(crate) fn as_str(self) -> &'static str {
     349            0 :         match self {
     350            0 :             ColdStartInfo::Unknown => "unknown",
     351            0 :             ColdStartInfo::Warm => "warm",
     352            0 :             ColdStartInfo::VmPoolHit => "pool_hit",
     353            0 :             ColdStartInfo::VmPoolMiss => "pool_miss",
     354            0 :             ColdStartInfo::HttpPoolHit => "http_pool_hit",
     355            0 :             ColdStartInfo::WarmCached => "warm_cached",
     356              :         }
     357            0 :     }
     358              : }
     359              : 
     360            0 : #[derive(Debug, Deserialize, Clone)]
     361              : pub struct EndpointJwksResponse {
     362              :     pub jwks: Vec<JwksSettings>,
     363              : }
     364              : 
     365            0 : #[derive(Debug, Deserialize, Clone)]
     366              : pub struct JwksSettings {
     367              :     pub id: String,
     368              :     pub jwks_url: url::Url,
     369              :     #[serde(rename = "provider_name")]
     370              :     pub _provider_name: String,
     371              :     pub jwt_audience: Option<String>,
     372              :     pub role_names: Vec<RoleNameInt>,
     373              : }
     374              : 
     375              : #[cfg(test)]
     376              : mod tests {
     377              :     use serde_json::json;
     378              : 
     379              :     use super::*;
     380              : 
     381            6 :     fn dummy_aux() -> serde_json::Value {
     382            6 :         json!({
     383            6 :             "endpoint_id": "endpoint",
     384            6 :             "project_id": "project",
     385            6 :             "branch_id": "branch",
     386            6 :             "compute_id": "compute",
     387            6 :             "cold_start_info": "unknown",
     388            6 :         })
     389            6 :     }
     390              : 
     391              :     #[test]
     392            1 :     fn parse_kick_session() -> anyhow::Result<()> {
     393            1 :         // This is what the console's kickResponse looks like.
     394            1 :         let json = json!({
     395            1 :             "session_id": "deadbeef",
     396            1 :             "result": {
     397            1 :                 "Success": {
     398            1 :                     "host": "localhost",
     399            1 :                     "port": 5432,
     400            1 :                     "dbname": "postgres",
     401            1 :                     "user": "john_doe",
     402            1 :                     "password": "password",
     403            1 :                     "aux": dummy_aux(),
     404            1 :                 }
     405            1 :             }
     406            1 :         });
     407            1 :         serde_json::from_str::<KickSession<'_>>(&json.to_string())?;
     408              : 
     409            1 :         Ok(())
     410            1 :     }
     411              : 
     412              :     #[test]
     413            1 :     fn parse_db_info() -> anyhow::Result<()> {
     414            1 :         // with password
     415            1 :         serde_json::from_value::<DatabaseInfo>(json!({
     416            1 :             "host": "localhost",
     417            1 :             "port": 5432,
     418            1 :             "dbname": "postgres",
     419            1 :             "user": "john_doe",
     420            1 :             "password": "password",
     421            1 :             "aux": dummy_aux(),
     422            1 :         }))?;
     423              : 
     424              :         // without password
     425            1 :         serde_json::from_value::<DatabaseInfo>(json!({
     426            1 :             "host": "localhost",
     427            1 :             "port": 5432,
     428            1 :             "dbname": "postgres",
     429            1 :             "user": "john_doe",
     430            1 :             "aux": dummy_aux(),
     431            1 :         }))?;
     432              : 
     433              :         // new field (forward compatibility)
     434            1 :         serde_json::from_value::<DatabaseInfo>(json!({
     435            1 :             "host": "localhost",
     436            1 :             "port": 5432,
     437            1 :             "dbname": "postgres",
     438            1 :             "user": "john_doe",
     439            1 :             "project": "hello_world",
     440            1 :             "N.E.W": "forward compatibility check",
     441            1 :             "aux": dummy_aux(),
     442            1 :         }))?;
     443              : 
     444              :         // with allowed_ips
     445            1 :         let dbinfo = serde_json::from_value::<DatabaseInfo>(json!({
     446            1 :             "host": "localhost",
     447            1 :             "port": 5432,
     448            1 :             "dbname": "postgres",
     449            1 :             "user": "john_doe",
     450            1 :             "password": "password",
     451            1 :             "aux": dummy_aux(),
     452            1 :             "allowed_ips": ["127.0.0.1"],
     453            1 :         }))?;
     454              : 
     455            1 :         assert_eq!(
     456            1 :             dbinfo.allowed_ips,
     457            1 :             Some(vec![IpPattern::Single("127.0.0.1".parse()?)])
     458              :         );
     459              : 
     460            1 :         Ok(())
     461            1 :     }
     462              : 
     463              :     #[test]
     464            1 :     fn parse_wake_compute() -> anyhow::Result<()> {
     465            1 :         let json = json!({
     466            1 :             "address": "0.0.0.0",
     467            1 :             "aux": dummy_aux(),
     468            1 :         });
     469            1 :         serde_json::from_str::<WakeCompute>(&json.to_string())?;
     470            1 :         Ok(())
     471            1 :     }
     472              : 
     473              :     #[test]
     474            1 :     fn parse_get_role_secret() -> anyhow::Result<()> {
     475            1 :         // Empty `allowed_ips` and `allowed_vpc_endpoint_ids` field.
     476            1 :         let json = json!({
     477            1 :             "role_secret": "secret",
     478            1 :         });
     479            1 :         serde_json::from_str::<GetEndpointAccessControl>(&json.to_string())?;
     480            1 :         let json = json!({
     481            1 :             "role_secret": "secret",
     482            1 :             "allowed_ips": ["8.8.8.8"],
     483            1 :         });
     484            1 :         serde_json::from_str::<GetEndpointAccessControl>(&json.to_string())?;
     485            1 :         let json = json!({
     486            1 :             "role_secret": "secret",
     487            1 :             "allowed_vpc_endpoint_ids": ["vpce-0abcd1234567890ef"],
     488            1 :         });
     489            1 :         serde_json::from_str::<GetEndpointAccessControl>(&json.to_string())?;
     490            1 :         let json = json!({
     491            1 :             "role_secret": "secret",
     492            1 :             "allowed_ips": ["8.8.8.8"],
     493            1 :             "allowed_vpc_endpoint_ids": ["vpce-0abcd1234567890ef"],
     494            1 :         });
     495            1 :         serde_json::from_str::<GetEndpointAccessControl>(&json.to_string())?;
     496            1 :         let json = json!({
     497            1 :             "role_secret": "secret",
     498            1 :             "allowed_ips": ["8.8.8.8"],
     499            1 :             "allowed_vpc_endpoint_ids": ["vpce-0abcd1234567890ef"],
     500            1 :             "project_id": "project",
     501            1 :         });
     502            1 :         serde_json::from_str::<GetEndpointAccessControl>(&json.to_string())?;
     503              : 
     504            1 :         Ok(())
     505            1 :     }
     506              : }
        

Generated by: LCOV version 2.1-beta