Line data Source code
1 : use std::sync::Arc;
2 :
3 : use futures::{FutureExt, TryFutureExt};
4 : use tokio::io::{AsyncRead, AsyncWrite};
5 : use tokio_util::sync::CancellationToken;
6 : use tracing::{Instrument, debug, error, info};
7 :
8 : use crate::auth::backend::ConsoleRedirectBackend;
9 : use crate::cancellation::CancellationHandler;
10 : use crate::config::{ProxyConfig, ProxyProtocolV2};
11 : use crate::context::RequestContext;
12 : use crate::error::ReportableError;
13 : use crate::metrics::{Metrics, NumClientConnectionsGuard};
14 : use crate::pglb::ClientRequestError;
15 : use crate::pglb::handshake::{HandshakeData, handshake};
16 : use crate::pglb::passthrough::ProxyPassthrough;
17 : use crate::protocol2::{ConnectHeader, ConnectionInfo, read_proxy_protocol};
18 : use crate::proxy::connect_compute::{TcpMechanism, connect_to_compute};
19 : use crate::proxy::{ErrorSource, finish_client_init};
20 : use crate::util::run_until_cancelled;
21 :
22 0 : pub async fn task_main(
23 0 : config: &'static ProxyConfig,
24 0 : backend: &'static ConsoleRedirectBackend,
25 0 : listener: tokio::net::TcpListener,
26 0 : cancellation_token: CancellationToken,
27 0 : cancellation_handler: Arc<CancellationHandler>,
28 0 : ) -> anyhow::Result<()> {
29 0 : scopeguard::defer! {
30 : info!("proxy has shut down");
31 : }
32 :
33 : // When set for the server socket, the keepalive setting
34 : // will be inherited by all accepted client sockets.
35 0 : socket2::SockRef::from(&listener).set_keepalive(true)?;
36 :
37 0 : let connections = tokio_util::task::task_tracker::TaskTracker::new();
38 0 : let cancellations = tokio_util::task::task_tracker::TaskTracker::new();
39 :
40 0 : while let Some(accept_result) =
41 0 : run_until_cancelled(listener.accept(), &cancellation_token).await
42 : {
43 0 : let (socket, peer_addr) = accept_result?;
44 :
45 0 : let conn_gauge = Metrics::get()
46 0 : .proxy
47 0 : .client_connections
48 0 : .guard(crate::metrics::Protocol::Tcp);
49 :
50 0 : let session_id = uuid::Uuid::new_v4();
51 0 : let cancellation_handler = Arc::clone(&cancellation_handler);
52 0 : let cancellations = cancellations.clone();
53 :
54 0 : debug!(protocol = "tcp", %session_id, "accepted new TCP connection");
55 :
56 0 : connections.spawn(async move {
57 0 : let (socket, conn_info) = match config.proxy_protocol_v2 {
58 : ProxyProtocolV2::Required => {
59 0 : match read_proxy_protocol(socket).await {
60 0 : Err(e) => {
61 0 : error!("per-client task finished with an error: {e:#}");
62 0 : return;
63 : }
64 : // our load balancers will not send any more data. let's just exit immediately
65 0 : Ok((_socket, ConnectHeader::Local)) => {
66 0 : debug!("healthcheck received");
67 0 : return;
68 : }
69 0 : Ok((socket, ConnectHeader::Proxy(info))) => (socket, info),
70 : }
71 : }
72 : // ignore the header - it cannot be confused for a postgres or http connection so will
73 : // error later.
74 0 : ProxyProtocolV2::Rejected => (
75 0 : socket,
76 0 : ConnectionInfo {
77 0 : addr: peer_addr,
78 0 : extra: None,
79 0 : },
80 0 : ),
81 : };
82 :
83 0 : match socket.set_nodelay(true) {
84 0 : Ok(()) => {}
85 0 : Err(e) => {
86 0 : error!(
87 0 : "per-client task finished with an error: failed to set socket option: {e:#}"
88 : );
89 0 : return;
90 : }
91 : }
92 :
93 0 : let ctx = RequestContext::new(session_id, conn_info, crate::metrics::Protocol::Tcp);
94 :
95 0 : let res = handle_client(
96 0 : config,
97 0 : backend,
98 0 : &ctx,
99 0 : cancellation_handler,
100 0 : socket,
101 0 : conn_gauge,
102 0 : cancellations,
103 0 : )
104 0 : .instrument(ctx.span())
105 0 : .boxed()
106 0 : .await;
107 :
108 0 : match res {
109 0 : Err(e) => {
110 0 : ctx.set_error_kind(e.get_error_kind());
111 0 : error!(parent: &ctx.span(), "per-client task finished with an error: {e:#}");
112 : }
113 0 : Ok(None) => {
114 0 : ctx.set_success();
115 0 : }
116 0 : Ok(Some(p)) => {
117 0 : ctx.set_success();
118 0 : let _disconnect = ctx.log_connect();
119 0 : match p.proxy_pass().await {
120 0 : Ok(()) => {}
121 0 : Err(ErrorSource::Client(e)) => {
122 0 : error!(
123 : ?session_id,
124 0 : "per-client task finished with an IO error from the client: {e:#}"
125 : );
126 : }
127 0 : Err(ErrorSource::Compute(e)) => {
128 0 : error!(
129 : ?session_id,
130 0 : "per-client task finished with an IO error from the compute: {e:#}"
131 : );
132 : }
133 : }
134 : }
135 : }
136 0 : });
137 : }
138 :
139 0 : connections.close();
140 0 : cancellations.close();
141 0 : drop(listener);
142 :
143 : // Drain connections
144 0 : connections.wait().await;
145 0 : cancellations.wait().await;
146 :
147 0 : Ok(())
148 0 : }
149 :
150 : #[allow(clippy::too_many_arguments)]
151 0 : pub(crate) async fn handle_client<S: AsyncRead + AsyncWrite + Unpin + Send>(
152 0 : config: &'static ProxyConfig,
153 0 : backend: &'static ConsoleRedirectBackend,
154 0 : ctx: &RequestContext,
155 0 : cancellation_handler: Arc<CancellationHandler>,
156 0 : stream: S,
157 0 : conn_gauge: NumClientConnectionsGuard<'static>,
158 0 : cancellations: tokio_util::task::task_tracker::TaskTracker,
159 0 : ) -> Result<Option<ProxyPassthrough<S>>, ClientRequestError> {
160 0 : debug!(
161 0 : protocol = %ctx.protocol(),
162 0 : "handling interactive connection from client"
163 : );
164 :
165 0 : let metrics = &Metrics::get().proxy;
166 0 : let proto = ctx.protocol();
167 0 : let request_gauge = metrics.connection_requests.guard(proto);
168 :
169 0 : let tls = config.tls_config.load();
170 0 : let tls = tls.as_deref();
171 :
172 0 : let record_handshake_error = !ctx.has_private_peer_addr();
173 0 : let pause = ctx.latency_timer_pause(crate::metrics::Waiting::Client);
174 0 : let do_handshake = handshake(ctx, stream, tls, record_handshake_error);
175 :
176 0 : let (mut stream, params) = match tokio::time::timeout(config.handshake_timeout, do_handshake)
177 0 : .await??
178 : {
179 0 : HandshakeData::Startup(stream, params) => (stream, params),
180 0 : HandshakeData::Cancel(cancel_key_data) => {
181 : // spawn a task to cancel the session, but don't wait for it
182 0 : cancellations.spawn({
183 0 : let cancellation_handler_clone = Arc::clone(&cancellation_handler);
184 0 : let ctx = ctx.clone();
185 0 : let cancel_span = tracing::span!(parent: None, tracing::Level::INFO, "cancel_session", session_id = ?ctx.session_id());
186 0 : cancel_span.follows_from(tracing::Span::current());
187 0 : async move {
188 0 : cancellation_handler_clone
189 0 : .cancel_session(
190 0 : cancel_key_data,
191 0 : ctx,
192 0 : config.authentication_config.ip_allowlist_check_enabled,
193 0 : config.authentication_config.is_vpc_acccess_proxy,
194 0 : backend.get_api(),
195 0 : )
196 0 : .await
197 0 : .inspect_err(|e | debug!(error = ?e, "cancel_session failed")).ok();
198 0 : }.instrument(cancel_span)
199 : });
200 :
201 0 : return Ok(None);
202 : }
203 : };
204 0 : drop(pause);
205 :
206 0 : ctx.set_db_options(params.clone());
207 :
208 0 : let (node_info, mut auth_info, user_info) = match backend
209 0 : .authenticate(ctx, &config.authentication_config, &mut stream)
210 0 : .await
211 : {
212 0 : Ok(auth_result) => auth_result,
213 0 : Err(e) => Err(stream.throw_error(e, Some(ctx)).await)?,
214 : };
215 0 : auth_info.set_startup_params(¶ms, true);
216 :
217 0 : let mut node = connect_to_compute(
218 0 : ctx,
219 0 : &TcpMechanism {
220 0 : locks: &config.connect_compute_locks,
221 0 : },
222 0 : &node_info,
223 0 : config.wake_compute_retry_config,
224 0 : &config.connect_to_compute,
225 : )
226 0 : .or_else(|e| async { Err(stream.throw_error(e, Some(ctx)).await) })
227 0 : .await?;
228 :
229 0 : let pg_settings = auth_info
230 0 : .authenticate(ctx, &mut node, &user_info)
231 0 : .or_else(|e| async { Err(stream.throw_error(e, Some(ctx)).await) })
232 0 : .await?;
233 :
234 0 : let session = cancellation_handler.get_key();
235 :
236 0 : finish_client_init(&pg_settings, *session.key(), &mut stream);
237 0 : let stream = stream.flush_and_into_inner().await?;
238 :
239 0 : let session_id = ctx.session_id();
240 0 : let (cancel_on_shutdown, cancel) = tokio::sync::oneshot::channel();
241 0 : tokio::spawn(async move {
242 0 : session
243 0 : .maintain_cancel_key(
244 0 : session_id,
245 0 : cancel,
246 0 : &pg_settings.cancel_closure,
247 0 : &config.connect_to_compute,
248 0 : )
249 0 : .await;
250 0 : });
251 :
252 0 : Ok(Some(ProxyPassthrough {
253 0 : client: stream,
254 0 : compute: node.stream,
255 0 :
256 0 : aux: node.aux,
257 0 : private_link_id: None,
258 0 :
259 0 : _cancel_on_shutdown: cancel_on_shutdown,
260 0 :
261 0 : _req: request_gauge,
262 0 : _conn: conn_gauge,
263 0 : _db_conn: node.guage,
264 0 : }))
265 0 : }
|
