Line data Source code
1 : use serde::{Deserialize, Serialize};
2 : use std::fmt;
3 :
4 : use crate::auth::IpPattern;
5 :
6 : use crate::intern::{BranchIdInt, EndpointIdInt, ProjectIdInt};
7 :
8 : /// Generic error response with human-readable description.
9 : /// Note that we can't always present it to user as is.
10 0 : #[derive(Debug, Deserialize)]
11 : pub struct ConsoleError {
12 : pub error: Box<str>,
13 : }
14 :
15 : /// Response which holds client's auth secret, e.g. [`crate::scram::ServerSecret`].
16 : /// Returned by the `/proxy_get_role_secret` API method.
17 30 : #[derive(Deserialize)]
18 : pub struct GetRoleSecret {
19 : pub role_secret: Box<str>,
20 : pub allowed_ips: Option<Vec<IpPattern>>,
21 : pub project_id: Option<ProjectIdInt>,
22 : }
23 :
24 : // Manually implement debug to omit sensitive info.
25 : impl fmt::Debug for GetRoleSecret {
26 0 : fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
27 0 : f.debug_struct("GetRoleSecret").finish_non_exhaustive()
28 0 : }
29 : }
30 :
31 : /// Response which holds compute node's `host:port` pair.
32 : /// Returned by the `/proxy_wake_compute` API method.
33 10 : #[derive(Debug, Deserialize)]
34 : pub struct WakeCompute {
35 : pub address: Box<str>,
36 : pub aux: MetricsAuxInfo,
37 : }
38 :
39 : /// Async response which concludes the link auth flow.
40 : /// Also known as `kickResponse` in the console.
41 12 : #[derive(Debug, Deserialize)]
42 : pub struct KickSession<'a> {
43 : /// Session ID is assigned by the proxy.
44 : pub session_id: &'a str,
45 :
46 : /// Compute node connection params.
47 : #[serde(deserialize_with = "KickSession::parse_db_info")]
48 : pub result: DatabaseInfo,
49 : }
50 :
51 : impl KickSession<'_> {
52 2 : fn parse_db_info<'de, D>(des: D) -> Result<DatabaseInfo, D::Error>
53 2 : where
54 2 : D: serde::Deserializer<'de>,
55 2 : {
56 4 : #[derive(Deserialize)]
57 2 : enum Wrapper {
58 2 : // Currently, console only reports `Success`.
59 2 : // `Failure(String)` used to be here... RIP.
60 2 : Success(DatabaseInfo),
61 2 : }
62 2 :
63 2 : Wrapper::deserialize(des).map(|x| match x {
64 2 : Wrapper::Success(info) => info,
65 2 : })
66 2 : }
67 : }
68 :
69 : /// Compute node connection params.
70 100 : #[derive(Deserialize)]
71 : pub struct DatabaseInfo {
72 : pub host: Box<str>,
73 : pub port: u16,
74 : pub dbname: Box<str>,
75 : pub user: Box<str>,
76 : /// Console always provides a password, but it might
77 : /// be inconvenient for debug with local PG instance.
78 : pub password: Option<Box<str>>,
79 : pub aux: MetricsAuxInfo,
80 : }
81 :
82 : // Manually implement debug to omit sensitive info.
83 : impl fmt::Debug for DatabaseInfo {
84 0 : fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
85 0 : f.debug_struct("DatabaseInfo")
86 0 : .field("host", &self.host)
87 0 : .field("port", &self.port)
88 0 : .field("dbname", &self.dbname)
89 0 : .field("user", &self.user)
90 0 : .finish_non_exhaustive()
91 0 : }
92 : }
93 :
94 : /// Various labels for prometheus metrics.
95 : /// Also known as `ProxyMetricsAuxInfo` in the console.
96 90 : #[derive(Debug, Deserialize, Clone)]
97 : pub struct MetricsAuxInfo {
98 : pub endpoint_id: EndpointIdInt,
99 : pub project_id: ProjectIdInt,
100 : pub branch_id: BranchIdInt,
101 : #[serde(default)]
102 : pub cold_start_info: ColdStartInfo,
103 : }
104 :
105 20 : #[derive(Debug, Default, Serialize, Deserialize, Clone, Copy)]
106 : #[serde(rename_all = "snake_case")]
107 : pub enum ColdStartInfo {
108 : #[default]
109 : Unknown,
110 : /// Compute was already running
111 : Warm,
112 : #[serde(rename = "pool_hit")]
113 : /// Compute was not running but there was an available VM
114 : VmPoolHit,
115 : #[serde(rename = "pool_miss")]
116 : /// Compute was not running and there were no VMs available
117 : VmPoolMiss,
118 :
119 : // not provided by control plane
120 : /// Connection available from HTTP pool
121 : HttpPoolHit,
122 : /// Cached connection info
123 : WarmCached,
124 : }
125 :
126 : impl ColdStartInfo {
127 140 : pub fn as_str(&self) -> &'static str {
128 140 : match self {
129 140 : ColdStartInfo::Unknown => "unknown",
130 0 : ColdStartInfo::Warm => "warm",
131 0 : ColdStartInfo::VmPoolHit => "pool_hit",
132 0 : ColdStartInfo::VmPoolMiss => "pool_miss",
133 0 : ColdStartInfo::HttpPoolHit => "http_pool_hit",
134 0 : ColdStartInfo::WarmCached => "warm_cached",
135 : }
136 140 : }
137 : }
138 :
139 : #[cfg(test)]
140 : mod tests {
141 : use super::*;
142 : use serde_json::json;
143 :
144 10 : fn dummy_aux() -> serde_json::Value {
145 10 : json!({
146 10 : "endpoint_id": "endpoint",
147 10 : "project_id": "project",
148 10 : "branch_id": "branch",
149 10 : "cold_start_info": "unknown",
150 10 : })
151 10 : }
152 :
153 : #[test]
154 2 : fn parse_kick_session() -> anyhow::Result<()> {
155 2 : // This is what the console's kickResponse looks like.
156 2 : let json = json!({
157 2 : "session_id": "deadbeef",
158 2 : "result": {
159 2 : "Success": {
160 2 : "host": "localhost",
161 2 : "port": 5432,
162 2 : "dbname": "postgres",
163 2 : "user": "john_doe",
164 2 : "password": "password",
165 2 : "aux": dummy_aux(),
166 2 : }
167 2 : }
168 2 : });
169 2 : let _: KickSession = serde_json::from_str(&json.to_string())?;
170 :
171 2 : Ok(())
172 2 : }
173 :
174 : #[test]
175 2 : fn parse_db_info() -> anyhow::Result<()> {
176 : // with password
177 2 : let _: DatabaseInfo = serde_json::from_value(json!({
178 2 : "host": "localhost",
179 2 : "port": 5432,
180 2 : "dbname": "postgres",
181 2 : "user": "john_doe",
182 2 : "password": "password",
183 2 : "aux": dummy_aux(),
184 2 : }))?;
185 :
186 : // without password
187 2 : let _: DatabaseInfo = serde_json::from_value(json!({
188 2 : "host": "localhost",
189 2 : "port": 5432,
190 2 : "dbname": "postgres",
191 2 : "user": "john_doe",
192 2 : "aux": dummy_aux(),
193 2 : }))?;
194 :
195 : // new field (forward compatibility)
196 2 : let _: DatabaseInfo = serde_json::from_value(json!({
197 2 : "host": "localhost",
198 2 : "port": 5432,
199 2 : "dbname": "postgres",
200 2 : "user": "john_doe",
201 2 : "project": "hello_world",
202 2 : "N.E.W": "forward compatibility check",
203 2 : "aux": dummy_aux(),
204 2 : }))?;
205 :
206 2 : Ok(())
207 2 : }
208 :
209 : #[test]
210 2 : fn parse_wake_compute() -> anyhow::Result<()> {
211 2 : let json = json!({
212 2 : "address": "0.0.0.0",
213 2 : "aux": dummy_aux(),
214 2 : });
215 2 : let _: WakeCompute = serde_json::from_str(&json.to_string())?;
216 2 : Ok(())
217 2 : }
218 :
219 : #[test]
220 2 : fn parse_get_role_secret() -> anyhow::Result<()> {
221 2 : // Empty `allowed_ips` field.
222 2 : let json = json!({
223 2 : "role_secret": "secret",
224 2 : });
225 2 : let _: GetRoleSecret = serde_json::from_str(&json.to_string())?;
226 2 : let json = json!({
227 2 : "role_secret": "secret",
228 2 : "allowed_ips": ["8.8.8.8"],
229 2 : });
230 2 : let _: GetRoleSecret = serde_json::from_str(&json.to_string())?;
231 2 : let json = json!({
232 2 : "role_secret": "secret",
233 2 : "allowed_ips": ["8.8.8.8"],
234 2 : "project_id": "project",
235 2 : });
236 2 : let _: GetRoleSecret = serde_json::from_str(&json.to_string())?;
237 :
238 2 : Ok(())
239 2 : }
240 : }
|
