TLA Line data Source code
1 : mod classic;
2 : mod hacks;
3 : mod link;
4 :
5 : pub use link::LinkAuthError;
6 :
7 : use crate::{
8 : auth::{self, ClientCredentials},
9 : console::{
10 : self,
11 : provider::{CachedNodeInfo, ConsoleReqExtra},
12 : Api,
13 : },
14 : stream, url,
15 : };
16 : use futures::TryFutureExt;
17 : use std::borrow::Cow;
18 : use tokio::io::{AsyncRead, AsyncWrite};
19 : use tracing::info;
20 :
21 : /// A product of successful authentication.
22 : pub struct AuthSuccess<T> {
23 : /// Did we send [`pq_proto::BeMessage::AuthenticationOk`] to client?
24 : pub reported_auth_ok: bool,
25 : /// Something to be considered a positive result.
26 : pub value: T,
27 : }
28 :
29 : impl<T> AuthSuccess<T> {
30 : /// Very similar to [`std::option::Option::map`].
31 : /// Maps [`AuthSuccess<T>`] to [`AuthSuccess<R>`] by applying
32 : /// a function to a contained value.
33 CBC 3 : pub fn map<R>(self, f: impl FnOnce(T) -> R) -> AuthSuccess<R> {
34 3 : AuthSuccess {
35 3 : reported_auth_ok: self.reported_auth_ok,
36 3 : value: f(self.value),
37 3 : }
38 3 : }
39 : }
40 :
41 : /// This type serves two purposes:
42 : ///
43 : /// * When `T` is `()`, it's just a regular auth backend selector
44 : /// which we use in [`crate::config::ProxyConfig`].
45 : ///
46 : /// * However, when we substitute `T` with [`ClientCredentials`],
47 : /// this helps us provide the credentials only to those auth
48 : /// backends which require them for the authentication process.
49 : pub enum BackendType<'a, T> {
50 : /// Current Cloud API (V2).
51 : Console(Cow<'a, console::provider::neon::Api>, T),
52 : /// Local mock of Cloud API (V2).
53 : Postgres(Cow<'a, console::provider::mock::Api>, T),
54 : /// Authentication via a web browser.
55 : Link(Cow<'a, url::ApiUrl>),
56 : /// Test backend.
57 : Test(&'a dyn TestBackend),
58 : }
59 :
60 : pub trait TestBackend: Send + Sync + 'static {
61 : fn wake_compute(&self) -> Result<CachedNodeInfo, console::errors::WakeComputeError>;
62 : }
63 :
64 : impl std::fmt::Display for BackendType<'_, ()> {
65 16 : fn fmt(&self, fmt: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
66 16 : use BackendType::*;
67 16 : match self {
68 UBC 0 : Console(endpoint, _) => fmt.debug_tuple("Console").field(&endpoint.url()).finish(),
69 CBC 13 : Postgres(endpoint, _) => fmt.debug_tuple("Postgres").field(&endpoint.url()).finish(),
70 3 : Link(url) => fmt.debug_tuple("Link").field(&url.as_str()).finish(),
71 UBC 0 : Test(_) => fmt.debug_tuple("Test").finish(),
72 : }
73 CBC 16 : }
74 : }
75 :
76 : impl<T> BackendType<'_, T> {
77 : /// Very similar to [`std::option::Option::as_ref`].
78 : /// This helps us pass structured config to async tasks.
79 60 : pub fn as_ref(&self) -> BackendType<'_, &T> {
80 60 : use BackendType::*;
81 60 : match self {
82 UBC 0 : Console(c, x) => Console(Cow::Borrowed(c), x),
83 CBC 57 : Postgres(c, x) => Postgres(Cow::Borrowed(c), x),
84 3 : Link(c) => Link(Cow::Borrowed(c)),
85 UBC 0 : Test(x) => Test(*x),
86 : }
87 CBC 60 : }
88 : }
89 :
90 : impl<'a, T> BackendType<'a, T> {
91 : /// Very similar to [`std::option::Option::map`].
92 : /// Maps [`BackendType<T>`] to [`BackendType<R>`] by applying
93 : /// a function to a contained value.
94 60 : pub fn map<R>(self, f: impl FnOnce(T) -> R) -> BackendType<'a, R> {
95 60 : use BackendType::*;
96 60 : match self {
97 UBC 0 : Console(c, x) => Console(c, f(x)),
98 CBC 57 : Postgres(c, x) => Postgres(c, f(x)),
99 3 : Link(c) => Link(c),
100 UBC 0 : Test(x) => Test(x),
101 : }
102 CBC 60 : }
103 : }
104 :
105 : impl<'a, T, E> BackendType<'a, Result<T, E>> {
106 : /// Very similar to [`std::option::Option::transpose`].
107 : /// This is most useful for error handling.
108 60 : pub fn transpose(self) -> Result<BackendType<'a, T>, E> {
109 60 : use BackendType::*;
110 60 : match self {
111 UBC 0 : Console(c, x) => x.map(|x| Console(c, x)),
112 CBC 57 : Postgres(c, x) => x.map(|x| Postgres(c, x)),
113 3 : Link(c) => Ok(Link(c)),
114 UBC 0 : Test(x) => Ok(Test(x)),
115 : }
116 CBC 60 : }
117 : }
118 :
119 : /// True to its name, this function encapsulates our current auth trade-offs.
120 : /// Here, we choose the appropriate auth flow based on circumstances.
121 30 : async fn auth_quirks(
122 30 : api: &impl console::Api,
123 30 : extra: &ConsoleReqExtra<'_>,
124 30 : creds: &mut ClientCredentials<'_>,
125 30 : client: &mut stream::PqStream<impl AsyncRead + AsyncWrite + Unpin>,
126 30 : allow_cleartext: bool,
127 30 : ) -> auth::Result<AuthSuccess<CachedNodeInfo>> {
128 30 : // If there's no project so far, that entails that client doesn't
129 30 : // support SNI or other means of passing the endpoint (project) name.
130 30 : // We now expect to see a very specific payload in the place of password.
131 30 : if creds.project.is_none() {
132 : // Password will be checked by the compute node later.
133 3 : return hacks::password_hack(api, extra, creds, client).await;
134 27 : }
135 27 :
136 27 : // Password hack should set the project name.
137 27 : // TODO: make `creds.project` more type-safe.
138 27 : assert!(creds.project.is_some());
139 :
140 : // Perform cleartext auth if we're allowed to do that.
141 : // Currently, we use it for websocket connections (latency).
142 27 : if allow_cleartext {
143 : // Password will be checked by the compute node later.
144 UBC 0 : return hacks::cleartext_hack(api, extra, creds, client).await;
145 CBC 27 : }
146 27 :
147 27 : // Finally, proceed with the main auth flow (SCRAM-based).
148 189 : classic::authenticate(api, extra, creds, client).await
149 30 : }
150 :
151 : impl BackendType<'_, ClientCredentials<'_>> {
152 : /// Get compute endpoint name from the credentials.
153 33 : pub fn get_endpoint(&self) -> Option<String> {
154 33 : use BackendType::*;
155 33 :
156 33 : match self {
157 UBC 0 : Console(_, creds) => creds.project.clone(),
158 CBC 30 : Postgres(_, creds) => creds.project.clone(),
159 3 : Link(_) => Some("link".to_owned()),
160 UBC 0 : Test(_) => Some("test".to_owned()),
161 : }
162 CBC 33 : }
163 :
164 : /// Get username from the credentials.
165 4 : pub fn get_user(&self) -> &str {
166 4 : use BackendType::*;
167 4 :
168 4 : match self {
169 UBC 0 : Console(_, creds) => creds.user,
170 CBC 4 : Postgres(_, creds) => creds.user,
171 UBC 0 : Link(_) => "link",
172 0 : Test(_) => "test",
173 : }
174 CBC 4 : }
175 :
176 : /// Authenticate the client via the requested backend, possibly using credentials.
177 99 : #[tracing::instrument(fields(allow_cleartext = allow_cleartext), skip_all)]
178 : pub async fn authenticate(
179 : &mut self,
180 : extra: &ConsoleReqExtra<'_>,
181 : client: &mut stream::PqStream<impl AsyncRead + AsyncWrite + Unpin>,
182 : allow_cleartext: bool,
183 : ) -> auth::Result<AuthSuccess<CachedNodeInfo>> {
184 : use BackendType::*;
185 :
186 : let res = match self {
187 : Console(api, creds) => {
188 UBC 0 : info!(
189 0 : user = creds.user,
190 0 : project = creds.project(),
191 0 : "performing authentication using the console"
192 0 : );
193 :
194 : let api = api.as_ref();
195 : auth_quirks(api, extra, creds, client, allow_cleartext).await?
196 : }
197 : Postgres(api, creds) => {
198 CBC 30 : info!(
199 30 : user = creds.user,
200 30 : project = creds.project(),
201 30 : "performing authentication using a local postgres instance"
202 30 : );
203 :
204 : let api = api.as_ref();
205 : auth_quirks(api, extra, creds, client, allow_cleartext).await?
206 : }
207 : // NOTE: this auth backend doesn't use client credentials.
208 : Link(url) => {
209 3 : info!("performing link authentication");
210 :
211 : link::authenticate(url, client)
212 : .await?
213 : .map(CachedNodeInfo::new_uncached)
214 : }
215 : Test(_) => {
216 : unreachable!("this function should never be called in the test backend")
217 : }
218 : };
219 :
220 29 : info!("user successfully authenticated");
221 : Ok(res)
222 : }
223 :
224 : /// When applicable, wake the compute node, gaining its connection info in the process.
225 : /// The link auth flow doesn't support this, so we return [`None`] in that case.
226 27 : pub async fn wake_compute(
227 27 : &self,
228 27 : extra: &ConsoleReqExtra<'_>,
229 27 : ) -> Result<Option<CachedNodeInfo>, console::errors::WakeComputeError> {
230 27 : use BackendType::*;
231 27 :
232 27 : match self {
233 UBC 0 : Console(api, creds) => api.wake_compute(extra, creds).map_ok(Some).await,
234 CBC 27 : Postgres(api, creds) => api.wake_compute(extra, creds).map_ok(Some).await,
235 UBC 0 : Link(_) => Ok(None),
236 0 : Test(x) => x.wake_compute().map(Some),
237 : }
238 CBC 27 : }
239 : }
|