qail_pg/driver/connection/
startup.rs

1//! Startup handshake — authentication, parameter negotiation, prepared stmt mgmt.
2
3use super::helpers::{generate_gss_token, md5_password_message, select_scram_mechanism};
4use super::types::{GSS_SESSION_COUNTER, PgConnection, StartupAuthFlow};
5use crate::driver::stream::PgStream;
6use crate::driver::{AuthSettings, EnterpriseAuthMechanism, GssTokenProvider, PgError, PgResult};
7use crate::protocol::{BackendMessage, FrontendMessage, ScramClient, TransactionStatus};
8use sha2::{Digest, Sha256};
9use std::sync::atomic::Ordering;
10
11impl PgConnection {
12    /// Handle startup sequence (auth + params).
13    pub(super) async fn handle_startup(
14        &mut self,
15        user: &str,
16        password: Option<&str>,
17        auth_settings: AuthSettings,
18        gss_token_provider: Option<GssTokenProvider>,
19    ) -> PgResult<()> {
20        let mut scram_client: Option<ScramClient> = None;
21        let mut startup_auth_flow: Option<StartupAuthFlow> = None;
22        let mut saw_auth_ok = false;
23        let gss_session_id = GSS_SESSION_COUNTER.fetch_add(1, Ordering::Relaxed);
24        let mut gss_roundtrips: u32 = 0;
25        const MAX_GSS_ROUNDTRIPS: u32 = 32;
26
27        loop {
28            let msg = self.recv().await?;
29            if saw_auth_ok
30                && matches!(
31                    &msg,
32                    BackendMessage::AuthenticationOk
33                        | BackendMessage::AuthenticationKerberosV5
34                        | BackendMessage::AuthenticationGSS
35                        | BackendMessage::AuthenticationSCMCredential
36                        | BackendMessage::AuthenticationGSSContinue(_)
37                        | BackendMessage::AuthenticationSSPI
38                        | BackendMessage::AuthenticationCleartextPassword
39                        | BackendMessage::AuthenticationMD5Password(_)
40                        | BackendMessage::AuthenticationSASL(_)
41                        | BackendMessage::AuthenticationSASLContinue(_)
42                        | BackendMessage::AuthenticationSASLFinal(_)
43                )
44            {
45                return Err(PgError::Protocol(
46                    "Received authentication challenge after AuthenticationOk".to_string(),
47                ));
48            }
49            match msg {
50                BackendMessage::AuthenticationOk => {
51                    if let Some(StartupAuthFlow::Scram {
52                        server_final_seen: false,
53                    }) = startup_auth_flow
54                    {
55                        return Err(PgError::Protocol(
56                            "Received AuthenticationOk before AuthenticationSASLFinal".to_string(),
57                        ));
58                    }
59                    saw_auth_ok = true;
60                }
61                BackendMessage::AuthenticationKerberosV5 => {
62                    if let Some(flow) = startup_auth_flow {
63                        return Err(PgError::Protocol(format!(
64                            "Received AuthenticationKerberosV5 while {} authentication is in progress",
65                            flow.label()
66                        )));
67                    }
68                    startup_auth_flow = Some(StartupAuthFlow::EnterpriseGss {
69                        mechanism: EnterpriseAuthMechanism::KerberosV5,
70                    });
71
72                    if !auth_settings.allow_kerberos_v5 {
73                        return Err(PgError::Auth(
74                            "Server requested Kerberos V5 authentication, but Kerberos V5 is disabled by AuthSettings".to_string(),
75                        ));
76                    }
77
78                    if gss_token_provider.is_none() {
79                        return Err(PgError::Auth(
80                            "Kerberos V5 authentication requested but no GSS token provider is configured. Set ConnectOptions.gss_token_provider.".to_string(),
81                        ));
82                    }
83
84                    let token = generate_gss_token(
85                        gss_session_id,
86                        EnterpriseAuthMechanism::KerberosV5,
87                        None,
88                        gss_token_provider.as_ref(),
89                    )
90                    .map_err(|e| {
91                        PgError::Auth(format!("Kerberos V5 token generation failed: {}", e))
92                    })?;
93
94                    self.send(FrontendMessage::GSSResponse(token)).await?;
95                }
96                BackendMessage::AuthenticationGSS => {
97                    if let Some(flow) = startup_auth_flow {
98                        return Err(PgError::Protocol(format!(
99                            "Received AuthenticationGSS while {} authentication is in progress",
100                            flow.label()
101                        )));
102                    }
103                    startup_auth_flow = Some(StartupAuthFlow::EnterpriseGss {
104                        mechanism: EnterpriseAuthMechanism::GssApi,
105                    });
106
107                    if !auth_settings.allow_gssapi {
108                        return Err(PgError::Auth(
109                            "Server requested GSSAPI authentication, but GSSAPI is disabled by AuthSettings".to_string(),
110                        ));
111                    }
112
113                    if gss_token_provider.is_none() {
114                        return Err(PgError::Auth(
115                            "GSSAPI authentication requested but no GSS token provider is configured. Set ConnectOptions.gss_token_provider.".to_string(),
116                        ));
117                    }
118
119                    let token = generate_gss_token(
120                        gss_session_id,
121                        EnterpriseAuthMechanism::GssApi,
122                        None,
123                        gss_token_provider.as_ref(),
124                    )
125                    .map_err(|e| {
126                        PgError::Auth(format!("GSSAPI initial token generation failed: {}", e))
127                    })?;
128
129                    self.send(FrontendMessage::GSSResponse(token)).await?;
130                }
131                BackendMessage::AuthenticationSCMCredential => {
132                    if let Some(flow) = startup_auth_flow {
133                        return Err(PgError::Protocol(format!(
134                            "Received AuthenticationSCMCredential while {} authentication is in progress",
135                            flow.label()
136                        )));
137                    }
138                    return Err(PgError::Auth(
139                        "Server requested SCM credential authentication (auth code 6). This driver currently does not support Unix-socket credential passing; use SCRAM, GSS/SSPI, or password auth for this connection."
140                            .to_string(),
141                    ));
142                }
143                BackendMessage::AuthenticationSSPI => {
144                    if let Some(flow) = startup_auth_flow {
145                        return Err(PgError::Protocol(format!(
146                            "Received AuthenticationSSPI while {} authentication is in progress",
147                            flow.label()
148                        )));
149                    }
150                    startup_auth_flow = Some(StartupAuthFlow::EnterpriseGss {
151                        mechanism: EnterpriseAuthMechanism::Sspi,
152                    });
153
154                    if !auth_settings.allow_sspi {
155                        return Err(PgError::Auth(
156                            "Server requested SSPI authentication, but SSPI is disabled by AuthSettings".to_string(),
157                        ));
158                    }
159
160                    if gss_token_provider.is_none() {
161                        return Err(PgError::Auth(
162                            "SSPI authentication requested but no GSS token provider is configured. Set ConnectOptions.gss_token_provider.".to_string(),
163                        ));
164                    }
165
166                    let token = generate_gss_token(
167                        gss_session_id,
168                        EnterpriseAuthMechanism::Sspi,
169                        None,
170                        gss_token_provider.as_ref(),
171                    )
172                    .map_err(|e| {
173                        PgError::Auth(format!("SSPI initial token generation failed: {}", e))
174                    })?;
175
176                    self.send(FrontendMessage::GSSResponse(token)).await?;
177                }
178                BackendMessage::AuthenticationGSSContinue(server_token) => {
179                    gss_roundtrips += 1;
180                    if gss_roundtrips > MAX_GSS_ROUNDTRIPS {
181                        return Err(PgError::Auth(format!(
182                            "GSS handshake exceeded {} roundtrips — aborting",
183                            MAX_GSS_ROUNDTRIPS
184                        )));
185                    }
186
187                    let mechanism = match startup_auth_flow {
188                        Some(StartupAuthFlow::EnterpriseGss { mechanism }) => mechanism,
189                        Some(flow) => {
190                            return Err(PgError::Protocol(format!(
191                                "Received AuthenticationGSSContinue while {} authentication is in progress",
192                                flow.label()
193                            )));
194                        }
195                        None => {
196                            return Err(PgError::Auth(
197                                "Received GSSContinue without AuthenticationGSS/SSPI/KerberosV5 init"
198                                    .to_string(),
199                            ));
200                        }
201                    };
202
203                    if gss_token_provider.is_none() {
204                        return Err(PgError::Auth(
205                            "Received GSSContinue but no GSS token provider is configured. Set ConnectOptions.gss_token_provider.".to_string(),
206                        ));
207                    }
208
209                    let token = generate_gss_token(
210                        gss_session_id,
211                        mechanism,
212                        Some(&server_token),
213                        gss_token_provider.as_ref(),
214                    )
215                    .map_err(|e| {
216                        PgError::Auth(format!("GSS continue token generation failed: {}", e))
217                    })?;
218
219                    // Only send the response if there is actually a token to
220                    // send.  When gss_init_sec_context returns GSS_S_COMPLETE
221                    // on the final round, the token may be empty.  Sending an
222                    // empty GSSResponse ('p') after the server already
223                    // considers auth complete trips the "invalid frontend
224                    // message type 112" FATAL in PostgreSQL.
225                    if !token.is_empty() {
226                        self.send(FrontendMessage::GSSResponse(token)).await?;
227                    }
228                }
229                BackendMessage::AuthenticationCleartextPassword => {
230                    if let Some(flow) = startup_auth_flow {
231                        return Err(PgError::Protocol(format!(
232                            "Received AuthenticationCleartextPassword while {} authentication is in progress",
233                            flow.label()
234                        )));
235                    }
236                    startup_auth_flow = Some(StartupAuthFlow::CleartextPassword);
237
238                    if !auth_settings.allow_cleartext_password {
239                        return Err(PgError::Auth(
240                            "Server requested cleartext authentication, but cleartext is disabled by AuthSettings"
241                                .to_string(),
242                        ));
243                    }
244                    let password = password.ok_or_else(|| {
245                        PgError::Auth("Password required for cleartext authentication".to_string())
246                    })?;
247                    self.send(FrontendMessage::PasswordMessage(password.to_string()))
248                        .await?;
249                }
250                BackendMessage::AuthenticationMD5Password(salt) => {
251                    if let Some(flow) = startup_auth_flow {
252                        return Err(PgError::Protocol(format!(
253                            "Received AuthenticationMD5Password while {} authentication is in progress",
254                            flow.label()
255                        )));
256                    }
257                    startup_auth_flow = Some(StartupAuthFlow::Md5Password);
258
259                    if !auth_settings.allow_md5_password {
260                        return Err(PgError::Auth(
261                            "Server requested MD5 authentication, but MD5 is disabled by AuthSettings"
262                                .to_string(),
263                        ));
264                    }
265                    let password = password.ok_or_else(|| {
266                        PgError::Auth("Password required for MD5 authentication".to_string())
267                    })?;
268                    let md5_password = md5_password_message(user, password, salt);
269                    self.send(FrontendMessage::PasswordMessage(md5_password))
270                        .await?;
271                }
272                BackendMessage::AuthenticationSASL(mechanisms) => {
273                    if let Some(flow) = startup_auth_flow {
274                        return Err(PgError::Protocol(format!(
275                            "Received AuthenticationSASL while {} authentication is in progress",
276                            flow.label()
277                        )));
278                    }
279                    startup_auth_flow = Some(StartupAuthFlow::Scram {
280                        server_final_seen: false,
281                    });
282
283                    if !auth_settings.allow_scram_sha_256 {
284                        return Err(PgError::Auth(
285                            "Server requested SCRAM authentication, but SCRAM is disabled by AuthSettings"
286                                .to_string(),
287                        ));
288                    }
289                    let password = password.ok_or_else(|| {
290                        PgError::Auth("Password required for SCRAM authentication".to_string())
291                    })?;
292
293                    let tls_binding = self.tls_server_end_point_channel_binding();
294                    let (mechanism, channel_binding_data) = select_scram_mechanism(
295                        &mechanisms,
296                        tls_binding,
297                        auth_settings.channel_binding,
298                    )
299                    .map_err(PgError::Auth)?;
300
301                    let client = if let Some(binding_data) = channel_binding_data {
302                        ScramClient::new_with_tls_server_end_point(user, password, binding_data)
303                    } else {
304                        ScramClient::new(user, password)
305                    };
306                    let first_message = client.client_first_message();
307
308                    self.send(FrontendMessage::SASLInitialResponse {
309                        mechanism,
310                        data: first_message,
311                    })
312                    .await?;
313
314                    scram_client = Some(client);
315                }
316                BackendMessage::AuthenticationSASLContinue(server_data) => {
317                    match startup_auth_flow {
318                        Some(StartupAuthFlow::Scram {
319                            server_final_seen: false,
320                        }) => {}
321                        Some(StartupAuthFlow::Scram {
322                            server_final_seen: true,
323                        }) => {
324                            return Err(PgError::Protocol(
325                                "Received AuthenticationSASLContinue after AuthenticationSASLFinal"
326                                    .to_string(),
327                            ));
328                        }
329                        Some(flow) => {
330                            return Err(PgError::Protocol(format!(
331                                "Received AuthenticationSASLContinue while {} authentication is in progress",
332                                flow.label()
333                            )));
334                        }
335                        None => {
336                            return Err(PgError::Auth(
337                                "Received SASL Continue without SASL init".to_string(),
338                            ));
339                        }
340                    }
341
342                    let client = scram_client.as_mut().ok_or_else(|| {
343                        PgError::Auth("Received SASL Continue without SASL init".to_string())
344                    })?;
345
346                    let final_message = client
347                        .process_server_first(&server_data)
348                        .map_err(|e| PgError::Auth(format!("SCRAM error: {}", e)))?;
349
350                    self.send(FrontendMessage::SASLResponse(final_message))
351                        .await?;
352                }
353                BackendMessage::AuthenticationSASLFinal(server_signature) => {
354                    match startup_auth_flow {
355                        Some(StartupAuthFlow::Scram {
356                            server_final_seen: false,
357                        }) => {
358                            startup_auth_flow = Some(StartupAuthFlow::Scram {
359                                server_final_seen: true,
360                            });
361                        }
362                        Some(StartupAuthFlow::Scram {
363                            server_final_seen: true,
364                        }) => {
365                            return Err(PgError::Protocol(
366                                "Received duplicate AuthenticationSASLFinal".to_string(),
367                            ));
368                        }
369                        Some(flow) => {
370                            return Err(PgError::Protocol(format!(
371                                "Received AuthenticationSASLFinal while {} authentication is in progress",
372                                flow.label()
373                            )));
374                        }
375                        None => {
376                            return Err(PgError::Auth(
377                                "Received SASL Final without SASL init".to_string(),
378                            ));
379                        }
380                    }
381
382                    let client = scram_client.as_ref().ok_or_else(|| {
383                        PgError::Auth("Received SASL Final without SASL init".to_string())
384                    })?;
385                    client
386                        .verify_server_final(&server_signature)
387                        .map_err(|e| PgError::Auth(format!("Server verification failed: {}", e)))?;
388                }
389                BackendMessage::ParameterStatus { .. } => {
390                    if !saw_auth_ok {
391                        return Err(PgError::Protocol(
392                            "Received ParameterStatus before AuthenticationOk".to_string(),
393                        ));
394                    }
395                }
396                BackendMessage::NegotiateProtocolVersion {
397                    newest_minor_supported,
398                    unrecognized_protocol_options,
399                } => {
400                    if saw_auth_ok {
401                        return Err(PgError::Protocol(
402                            "Received NegotiateProtocolVersion after AuthenticationOk".to_string(),
403                        ));
404                    }
405                    let negotiated = if let Ok(minor) = u16::try_from(newest_minor_supported) {
406                        minor
407                    } else {
408                        let packed = u32::try_from(newest_minor_supported).map_err(|_| {
409                            PgError::Protocol(format!(
410                                "Invalid NegotiateProtocolVersion newest_minor_supported: {}",
411                                newest_minor_supported
412                            ))
413                        })?;
414                        let major = (packed >> 16) as u16;
415                        let minor = (packed & 0xFFFF) as u16;
416                        if major != 3 {
417                            return Err(PgError::Protocol(format!(
418                                "Invalid NegotiateProtocolVersion newest_minor_supported: {}",
419                                newest_minor_supported
420                            )));
421                        }
422                        minor
423                    };
424                    if negotiated > self.requested_protocol_minor {
425                        return Err(PgError::Protocol(format!(
426                            "Server negotiated protocol minor {} above requested {}",
427                            negotiated, self.requested_protocol_minor
428                        )));
429                    }
430                    self.negotiated_protocol_minor = negotiated;
431                    if !unrecognized_protocol_options.is_empty() {
432                        tracing::debug!(
433                            negotiated_minor = negotiated,
434                            unrecognized_count = unrecognized_protocol_options.len(),
435                            "startup_negotiate_protocol_version"
436                        );
437                    }
438                }
439                BackendMessage::BackendKeyData {
440                    process_id,
441                    secret_key,
442                } => {
443                    if !saw_auth_ok {
444                        return Err(PgError::Protocol(
445                            "Received BackendKeyData before AuthenticationOk".to_string(),
446                        ));
447                    }
448                    self.process_id = process_id;
449                    self.cancel_key_bytes = secret_key;
450                }
451                BackendMessage::ReadyForQuery(TransactionStatus::Idle) => {
452                    if !saw_auth_ok {
453                        return Err(PgError::Protocol(
454                            "Startup completed without AuthenticationOk".to_string(),
455                        ));
456                    }
457                    return Ok(());
458                }
459                BackendMessage::ReadyForQuery(status) => {
460                    if !saw_auth_ok {
461                        return Err(PgError::Protocol(
462                            "Startup completed without AuthenticationOk".to_string(),
463                        ));
464                    }
465                    return Err(PgError::Protocol(format!(
466                        "Startup completed with non-idle transaction status: {:?}",
467                        status
468                    )));
469                }
470                BackendMessage::ErrorResponse(err) => {
471                    return Err(PgError::Connection(err.message));
472                }
473                BackendMessage::NoticeResponse(_) => {}
474                _ => {
475                    return Err(PgError::Protocol(
476                        "Unexpected backend message during startup".to_string(),
477                    ));
478                }
479            }
480        }
481    }
482
483    /// Build SCRAM `tls-server-end-point` channel-binding bytes from the server leaf cert.
484    ///
485    /// PostgreSQL expects the hash of the peer certificate DER for
486    /// `SCRAM-SHA-256-PLUS` channel binding. We currently use SHA-256 here.
487    fn tls_server_end_point_channel_binding(&self) -> Option<Vec<u8>> {
488        let PgStream::Tls(tls) = &self.stream else {
489            return None;
490        };
491
492        let (_, conn) = tls.get_ref();
493        let certs = conn.peer_certificates()?;
494        let leaf_cert = certs.first()?;
495
496        let mut hasher = Sha256::new();
497        hasher.update(leaf_cert.as_ref());
498        Some(hasher.finalize().to_vec())
499    }
500
501    /// Gracefully close the connection by sending a Terminate message.
502    /// This tells the server we're done and allows proper cleanup.
503    pub async fn close(mut self) -> PgResult<()> {
504        use crate::protocol::PgEncoder;
505
506        // Send Terminate packet ('X')
507        let terminate = PgEncoder::encode_terminate();
508        self.write_all_with_timeout(&terminate, "stream write")
509            .await?;
510        self.flush_with_timeout("stream flush").await?;
511
512        Ok(())
513    }
514
515    /// Maximum prepared statements per connection before LRU eviction kicks in.
516    ///
517    /// This prevents memory spikes from dynamic batch filters generating
518    /// thousands of unique SQL shapes within a single request. Using LRU
519    /// eviction instead of nuclear `.clear()` preserves hot statements.
520    pub(crate) const MAX_PREPARED_PER_CONN: usize = 128;
521
522    /// Evict the least-recently-used prepared statement if at capacity.
523    ///
524    /// Called before every new statement registration to enforce
525    /// `MAX_PREPARED_PER_CONN`. Both `stmt_cache` (LRU ordering) and
526    /// `prepared_statements` (name→SQL map) are kept in sync.
527    pub(crate) fn evict_prepared_if_full(&mut self) {
528        if self.prepared_statements.len() >= Self::MAX_PREPARED_PER_CONN {
529            // Pop the LRU entry from the cache
530            if let Some((evicted_hash, evicted_name)) = self.stmt_cache.pop_lru() {
531                self.prepared_statements.remove(&evicted_name);
532                self.column_info_cache.remove(&evicted_hash);
533                self.pending_statement_closes.push(evicted_name);
534            } else {
535                // stmt_cache is empty but prepared_statements is full —
536                // shouldn't happen in normal flow, but handle defensively
537                // by clearing the oldest entry from the HashMap.
538                if let Some(key) = self.prepared_statements.keys().next().cloned() {
539                    self.prepared_statements.remove(&key);
540                    self.pending_statement_closes.push(key);
541                }
542            }
543        }
544    }
545
546    /// Clear all local prepared-statement state for this connection.
547    ///
548    /// Used by one-shot self-heal paths when server-side statement state
549    /// becomes invalid after DDL or failover.
550    pub(crate) fn clear_prepared_statement_state(&mut self) {
551        self.stmt_cache.clear();
552        self.prepared_statements.clear();
553        self.column_info_cache.clear();
554        self.pending_statement_closes.clear();
555    }
556}
qail_pg/driver/connection/startup.rs

qail_pg/driver/connection/
startup.rs
