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