Line data Source code
1 : #[cfg(test)]
2 : mod tests;
3 :
4 : pub mod connect_compute;
5 : mod copy_bidirectional;
6 : pub mod handshake;
7 : pub mod passthrough;
8 : pub mod retry;
9 : pub mod wake_compute;
10 : pub use copy_bidirectional::copy_bidirectional_client_compute;
11 : pub use copy_bidirectional::ErrorSource;
12 :
13 : use crate::{
14 : auth,
15 : cancellation::{self, CancellationHandlerMain, CancellationHandlerMainInternal},
16 : compute,
17 : config::{ProxyConfig, TlsConfig},
18 : context::RequestMonitoring,
19 : error::ReportableError,
20 : metrics::{Metrics, NumClientConnectionsGuard},
21 : protocol2::read_proxy_protocol,
22 : proxy::handshake::{handshake, HandshakeData},
23 : rate_limiter::EndpointRateLimiter,
24 : stream::{PqStream, Stream},
25 : EndpointCacheKey,
26 : };
27 : use futures::TryFutureExt;
28 : use itertools::Itertools;
29 : use once_cell::sync::OnceCell;
30 : use pq_proto::{BeMessage as Be, StartupMessageParams};
31 : use regex::Regex;
32 : use smol_str::{format_smolstr, SmolStr};
33 : use std::sync::Arc;
34 : use thiserror::Error;
35 : use tokio::io::{AsyncRead, AsyncWrite, AsyncWriteExt};
36 : use tokio_util::sync::CancellationToken;
37 : use tracing::{error, info, Instrument};
38 :
39 : use self::{
40 : connect_compute::{connect_to_compute, TcpMechanism},
41 : passthrough::ProxyPassthrough,
42 : };
43 :
44 : const ERR_INSECURE_CONNECTION: &str = "connection is insecure (try using `sslmode=require`)";
45 :
46 0 : pub async fn run_until_cancelled<F: std::future::Future>(
47 0 : f: F,
48 0 : cancellation_token: &CancellationToken,
49 0 : ) -> Option<F::Output> {
50 0 : match futures::future::select(
51 0 : std::pin::pin!(f),
52 0 : std::pin::pin!(cancellation_token.cancelled()),
53 0 : )
54 0 : .await
55 : {
56 0 : futures::future::Either::Left((f, _)) => Some(f),
57 0 : futures::future::Either::Right(((), _)) => None,
58 : }
59 0 : }
60 :
61 0 : pub async fn task_main(
62 0 : config: &'static ProxyConfig,
63 0 : listener: tokio::net::TcpListener,
64 0 : cancellation_token: CancellationToken,
65 0 : cancellation_handler: Arc<CancellationHandlerMain>,
66 0 : endpoint_rate_limiter: Arc<EndpointRateLimiter>,
67 0 : ) -> anyhow::Result<()> {
68 : scopeguard::defer! {
69 : info!("proxy has shut down");
70 : }
71 :
72 : // When set for the server socket, the keepalive setting
73 : // will be inherited by all accepted client sockets.
74 0 : socket2::SockRef::from(&listener).set_keepalive(true)?;
75 :
76 0 : let connections = tokio_util::task::task_tracker::TaskTracker::new();
77 :
78 0 : while let Some(accept_result) =
79 0 : run_until_cancelled(listener.accept(), &cancellation_token).await
80 : {
81 0 : let (socket, peer_addr) = accept_result?;
82 :
83 0 : let conn_gauge = Metrics::get()
84 0 : .proxy
85 0 : .client_connections
86 0 : .guard(crate::metrics::Protocol::Tcp);
87 0 :
88 0 : let session_id = uuid::Uuid::new_v4();
89 0 : let cancellation_handler = Arc::clone(&cancellation_handler);
90 0 :
91 0 : tracing::info!(protocol = "tcp", %session_id, "accepted new TCP connection");
92 0 : let endpoint_rate_limiter2 = endpoint_rate_limiter.clone();
93 0 :
94 0 : connections.spawn(async move {
95 0 : let (socket, peer_addr) = match read_proxy_protocol(socket).await {
96 0 : Ok((socket, Some(addr))) => (socket, addr.ip()),
97 0 : Err(e) => {
98 0 : error!("per-client task finished with an error: {e:#}");
99 0 : return;
100 : }
101 0 : Ok((_socket, None)) if config.require_client_ip => {
102 0 : error!("missing required client IP");
103 0 : return;
104 : }
105 0 : Ok((socket, None)) => (socket, peer_addr.ip()),
106 : };
107 :
108 0 : match socket.inner.set_nodelay(true) {
109 0 : Ok(()) => {}
110 0 : Err(e) => {
111 0 : error!("per-client task finished with an error: failed to set socket option: {e:#}");
112 0 : return;
113 : }
114 : };
115 :
116 0 : let mut ctx = RequestMonitoring::new(
117 0 : session_id,
118 0 : peer_addr,
119 0 : crate::metrics::Protocol::Tcp,
120 0 : &config.region,
121 0 : );
122 0 : let span = ctx.span.clone();
123 0 :
124 0 : let startup = Box::pin(
125 0 : handle_client(
126 0 : config,
127 0 : &mut ctx,
128 0 : cancellation_handler,
129 0 : socket,
130 0 : ClientMode::Tcp,
131 0 : endpoint_rate_limiter2,
132 0 : conn_gauge,
133 0 : )
134 0 : .instrument(span.clone()),
135 0 : );
136 0 : let res = startup.await;
137 :
138 0 : match res {
139 0 : Err(e) => {
140 0 : // todo: log and push to ctx the error kind
141 0 : ctx.set_error_kind(e.get_error_kind());
142 : error!(parent: &span, "per-client task finished with an error: {e:#}");
143 : }
144 0 : Ok(None) => {
145 0 : ctx.set_success();
146 0 : }
147 0 : Ok(Some(p)) => {
148 0 : ctx.set_success();
149 0 : ctx.log_connect();
150 0 : match p.proxy_pass().instrument(span.clone()).await {
151 0 : Ok(()) => {}
152 0 : Err(ErrorSource::Client(e)) => {
153 0 : error!(parent: &span, "per-client task finished with an IO error from the client: {e:#}");
154 0 : }
155 0 : Err(ErrorSource::Compute(e)) => {
156 0 : error!(parent: &span, "per-client task finished with an IO error from the compute: {e:#}");
157 0 : }
158 : }
159 : }
160 : }
161 0 : });
162 : }
163 :
164 0 : connections.close();
165 0 : drop(listener);
166 0 :
167 0 : // Drain connections
168 0 : connections.wait().await;
169 :
170 0 : Ok(())
171 0 : }
172 :
173 : pub enum ClientMode {
174 : Tcp,
175 : Websockets { hostname: Option<String> },
176 : }
177 :
178 : /// Abstracts the logic of handling TCP vs WS clients
179 : impl ClientMode {
180 0 : pub fn allow_cleartext(&self) -> bool {
181 0 : match self {
182 0 : ClientMode::Tcp => false,
183 0 : ClientMode::Websockets { .. } => true,
184 : }
185 0 : }
186 :
187 0 : pub fn allow_self_signed_compute(&self, config: &ProxyConfig) -> bool {
188 0 : match self {
189 0 : ClientMode::Tcp => config.allow_self_signed_compute,
190 0 : ClientMode::Websockets { .. } => false,
191 : }
192 0 : }
193 :
194 0 : fn hostname<'a, S>(&'a self, s: &'a Stream<S>) -> Option<&'a str> {
195 0 : match self {
196 0 : ClientMode::Tcp => s.sni_hostname(),
197 0 : ClientMode::Websockets { hostname } => hostname.as_deref(),
198 : }
199 0 : }
200 :
201 0 : fn handshake_tls<'a>(&self, tls: Option<&'a TlsConfig>) -> Option<&'a TlsConfig> {
202 0 : match self {
203 0 : ClientMode::Tcp => tls,
204 : // TLS is None here if using websockets, because the connection is already encrypted.
205 0 : ClientMode::Websockets { .. } => None,
206 : }
207 0 : }
208 : }
209 :
210 0 : #[derive(Debug, Error)]
211 : // almost all errors should be reported to the user, but there's a few cases where we cannot
212 : // 1. Cancellation: we are not allowed to tell the client any cancellation statuses for security reasons
213 : // 2. Handshake: handshake reports errors if it can, otherwise if the handshake fails due to protocol violation,
214 : // we cannot be sure the client even understands our error message
215 : // 3. PrepareClient: The client disconnected, so we can't tell them anyway...
216 : pub enum ClientRequestError {
217 : #[error("{0}")]
218 : Cancellation(#[from] cancellation::CancelError),
219 : #[error("{0}")]
220 : Handshake(#[from] handshake::HandshakeError),
221 : #[error("{0}")]
222 : HandshakeTimeout(#[from] tokio::time::error::Elapsed),
223 : #[error("{0}")]
224 : PrepareClient(#[from] std::io::Error),
225 : #[error("{0}")]
226 : ReportedError(#[from] crate::stream::ReportedError),
227 : }
228 :
229 : impl ReportableError for ClientRequestError {
230 0 : fn get_error_kind(&self) -> crate::error::ErrorKind {
231 0 : match self {
232 0 : ClientRequestError::Cancellation(e) => e.get_error_kind(),
233 0 : ClientRequestError::Handshake(e) => e.get_error_kind(),
234 0 : ClientRequestError::HandshakeTimeout(_) => crate::error::ErrorKind::RateLimit,
235 0 : ClientRequestError::ReportedError(e) => e.get_error_kind(),
236 0 : ClientRequestError::PrepareClient(_) => crate::error::ErrorKind::ClientDisconnect,
237 : }
238 0 : }
239 : }
240 :
241 0 : pub async fn handle_client<S: AsyncRead + AsyncWrite + Unpin>(
242 0 : config: &'static ProxyConfig,
243 0 : ctx: &mut RequestMonitoring,
244 0 : cancellation_handler: Arc<CancellationHandlerMain>,
245 0 : stream: S,
246 0 : mode: ClientMode,
247 0 : endpoint_rate_limiter: Arc<EndpointRateLimiter>,
248 0 : conn_gauge: NumClientConnectionsGuard<'static>,
249 0 : ) -> Result<Option<ProxyPassthrough<CancellationHandlerMainInternal, S>>, ClientRequestError> {
250 0 : info!(
251 : protocol = %ctx.protocol,
252 0 : "handling interactive connection from client"
253 : );
254 :
255 0 : let metrics = &Metrics::get().proxy;
256 0 : let proto = ctx.protocol;
257 0 : let _request_gauge = metrics.connection_requests.guard(proto);
258 0 :
259 0 : let tls = config.tls_config.as_ref();
260 0 :
261 0 : let record_handshake_error = !ctx.has_private_peer_addr();
262 0 : let pause = ctx.latency_timer.pause(crate::metrics::Waiting::Client);
263 0 : let do_handshake = handshake(stream, mode.handshake_tls(tls), record_handshake_error);
264 0 : let (mut stream, params) =
265 0 : match tokio::time::timeout(config.handshake_timeout, do_handshake).await?? {
266 0 : HandshakeData::Startup(stream, params) => (stream, params),
267 0 : HandshakeData::Cancel(cancel_key_data) => {
268 0 : return Ok(cancellation_handler
269 0 : .cancel_session(cancel_key_data, ctx.session_id)
270 0 : .await
271 0 : .map(|()| None)?)
272 : }
273 : };
274 0 : drop(pause);
275 0 :
276 0 : ctx.set_db_options(params.clone());
277 0 :
278 0 : let hostname = mode.hostname(stream.get_ref());
279 0 :
280 0 : let common_names = tls.map(|tls| &tls.common_names);
281 0 :
282 0 : // Extract credentials which we're going to use for auth.
283 0 : let result = config
284 0 : .auth_backend
285 0 : .as_ref()
286 0 : .map(|_| auth::ComputeUserInfoMaybeEndpoint::parse(ctx, ¶ms, hostname, common_names))
287 0 : .transpose();
288 :
289 0 : let user_info = match result {
290 0 : Ok(user_info) => user_info,
291 0 : Err(e) => stream.throw_error(e).await?,
292 : };
293 :
294 0 : let user = user_info.get_user().to_owned();
295 0 : let user_info = match user_info
296 0 : .authenticate(
297 0 : ctx,
298 0 : &mut stream,
299 0 : mode.allow_cleartext(),
300 0 : &config.authentication_config,
301 0 : endpoint_rate_limiter,
302 0 : )
303 0 : .await
304 : {
305 0 : Ok(auth_result) => auth_result,
306 0 : Err(e) => {
307 0 : let db = params.get("database");
308 0 : let app = params.get("application_name");
309 0 : let params_span = tracing::info_span!("", ?user, ?db, ?app);
310 :
311 0 : return stream.throw_error(e).instrument(params_span).await?;
312 : }
313 : };
314 :
315 0 : let mut node = connect_to_compute(
316 0 : ctx,
317 0 : &TcpMechanism {
318 0 : params: ¶ms,
319 0 : locks: &config.connect_compute_locks,
320 0 : },
321 0 : &user_info,
322 0 : mode.allow_self_signed_compute(config),
323 0 : config.wake_compute_retry_config,
324 0 : config.connect_to_compute_retry_config,
325 0 : )
326 0 : .or_else(|e| stream.throw_error(e))
327 0 : .await?;
328 :
329 0 : let session = cancellation_handler.get_session();
330 0 : prepare_client_connection(&node, &session, &mut stream).await?;
331 :
332 : // Before proxy passing, forward to compute whatever data is left in the
333 : // PqStream input buffer. Normally there is none, but our serverless npm
334 : // driver in pipeline mode sends startup, password and first query
335 : // immediately after opening the connection.
336 0 : let (stream, read_buf) = stream.into_inner();
337 0 : node.stream.write_all(&read_buf).await?;
338 :
339 0 : Ok(Some(ProxyPassthrough {
340 0 : client: stream,
341 0 : aux: node.aux.clone(),
342 0 : compute: node,
343 0 : req: _request_gauge,
344 0 : conn: conn_gauge,
345 0 : cancel: session,
346 0 : }))
347 0 : }
348 :
349 : /// Finish client connection initialization: confirm auth success, send params, etc.
350 0 : #[tracing::instrument(skip_all)]
351 : async fn prepare_client_connection<P>(
352 : node: &compute::PostgresConnection,
353 : session: &cancellation::Session<P>,
354 : stream: &mut PqStream<impl AsyncRead + AsyncWrite + Unpin>,
355 : ) -> Result<(), std::io::Error> {
356 : // Register compute's query cancellation token and produce a new, unique one.
357 : // The new token (cancel_key_data) will be sent to the client.
358 : let cancel_key_data = session.enable_query_cancellation(node.cancel_closure.clone());
359 :
360 : // Forward all postgres connection params to the client.
361 : // Right now the implementation is very hacky and inefficent (ideally,
362 : // we don't need an intermediate hashmap), but at least it should be correct.
363 : for (name, value) in &node.params {
364 : // TODO: Theoretically, this could result in a big pile of params...
365 : stream.write_message_noflush(&Be::ParameterStatus {
366 : name: name.as_bytes(),
367 : value: value.as_bytes(),
368 : })?;
369 : }
370 :
371 : stream
372 : .write_message_noflush(&Be::BackendKeyData(cancel_key_data))?
373 : .write_message(&Be::ReadyForQuery)
374 : .await?;
375 :
376 : Ok(())
377 : }
378 :
379 : #[derive(Debug, Clone, PartialEq, Eq, Default)]
380 : pub struct NeonOptions(Vec<(SmolStr, SmolStr)>);
381 :
382 : impl NeonOptions {
383 22 : pub fn parse_params(params: &StartupMessageParams) -> Self {
384 22 : params
385 22 : .options_raw()
386 22 : .map(Self::parse_from_iter)
387 22 : .unwrap_or_default()
388 22 : }
389 14 : pub fn parse_options_raw(options: &str) -> Self {
390 14 : Self::parse_from_iter(StartupMessageParams::parse_options_raw(options))
391 14 : }
392 :
393 4 : pub fn is_ephemeral(&self) -> bool {
394 4 : // Currently, neon endpoint options are all reserved for ephemeral endpoints.
395 4 : !self.0.is_empty()
396 4 : }
397 :
398 26 : fn parse_from_iter<'a>(options: impl Iterator<Item = &'a str>) -> Self {
399 26 : let mut options = options
400 26 : .filter_map(neon_option)
401 26 : .map(|(k, v)| (k.into(), v.into()))
402 26 : .collect_vec();
403 26 : options.sort();
404 26 : Self(options)
405 26 : }
406 :
407 8 : pub fn get_cache_key(&self, prefix: &str) -> EndpointCacheKey {
408 8 : // prefix + format!(" {k}:{v}")
409 8 : // kinda jank because SmolStr is immutable
410 8 : std::iter::once(prefix)
411 8 : .chain(self.0.iter().flat_map(|(k, v)| [" ", &**k, ":", &**v]))
412 8 : .collect::<SmolStr>()
413 8 : .into()
414 8 : }
415 :
416 : /// <https://swagger.io/docs/specification/serialization/> DeepObject format
417 : /// `paramName[prop1]=value1¶mName[prop2]=value2&...`
418 0 : pub fn to_deep_object(&self) -> Vec<(SmolStr, SmolStr)> {
419 0 : self.0
420 0 : .iter()
421 0 : .map(|(k, v)| (format_smolstr!("options[{}]", k), v.clone()))
422 0 : .collect()
423 0 : }
424 : }
425 :
426 64 : pub fn neon_option(bytes: &str) -> Option<(&str, &str)> {
427 64 : static RE: OnceCell<Regex> = OnceCell::new();
428 64 : let re = RE.get_or_init(|| Regex::new(r"^neon_(\w+):(.+)").unwrap());
429 :
430 64 : let cap = re.captures(bytes)?;
431 8 : let (_, [k, v]) = cap.extract();
432 8 : Some((k, v))
433 64 : }
|
