qail_pg/driver/connection/
helpers.rs

1//! Free helper functions — GSS token gen, metrics, MD5 password, SCRAM selection, Drop.
2
3#[cfg(all(target_os = "linux", feature = "native-io-uring"))]
4use super::types::CONNECT_BACKEND_IO_URING;
5use super::types::{CONNECT_BACKEND_TOKIO, PgConnection};
6use crate::driver::stream::PgStream;
7use crate::driver::{
8    EnterpriseAuthMechanism, GssTokenProvider, GssTokenRequest, PgError, PgResult,
9    ScramChannelBindingMode,
10};
11
12pub(super) fn generate_gss_token(
13    session_id: u64,
14    mechanism: EnterpriseAuthMechanism,
15    server_token: Option<&[u8]>,
16    stateful_provider: Option<&GssTokenProvider>,
17) -> Result<Vec<u8>, String> {
18    if let Some(provider) = stateful_provider {
19        return provider(GssTokenRequest {
20            session_id,
21            mechanism,
22            server_token,
23        });
24    }
25
26    Err("No GSS token provider configured".to_string())
27}
28
29pub(super) fn plain_connect_attempt_backend(io_uring: bool) -> &'static str {
30    #[cfg(all(target_os = "linux", feature = "native-io-uring"))]
31    {
32        if should_try_uring_plain(io_uring) {
33            return CONNECT_BACKEND_IO_URING;
34        }
35    }
36    #[cfg(not(all(target_os = "linux", feature = "native-io-uring")))]
37    {
38        let _ = io_uring;
39    }
40    CONNECT_BACKEND_TOKIO
41}
42
43pub(crate) fn connect_backend_for_stream(stream: &PgStream) -> &'static str {
44    match stream {
45        PgStream::Tcp(_) => CONNECT_BACKEND_TOKIO,
46        #[cfg(all(target_os = "linux", feature = "native-io-uring"))]
47        PgStream::Uring(_) => CONNECT_BACKEND_IO_URING,
48        PgStream::Tls(_) => CONNECT_BACKEND_TOKIO,
49        #[cfg(unix)]
50        PgStream::Unix(_) => CONNECT_BACKEND_TOKIO,
51        #[cfg(all(feature = "enterprise-gssapi", target_os = "linux"))]
52        PgStream::GssEnc(_) => CONNECT_BACKEND_TOKIO,
53    }
54}
55
56pub(super) fn connect_error_kind(error: &PgError) -> &'static str {
57    match error {
58        PgError::Connection(_) => "connection",
59        PgError::Protocol(_) => "protocol",
60        PgError::Auth(_) => "auth",
61        PgError::Query(_) | PgError::QueryServer(_) => "query",
62        PgError::NoRows => "no_rows",
63        PgError::Io(_) => "io",
64        PgError::Encode(_) => "encode",
65        PgError::Timeout(_) => "timeout",
66        PgError::PoolExhausted { .. } => "pool_exhausted",
67        PgError::PoolClosed => "pool_closed",
68    }
69}
70
71pub(super) fn record_connect_attempt(transport: &'static str, backend: &'static str) {
72    metrics::counter!(
73        "qail_pg_connect_attempt_total",
74        "transport" => transport,
75        "backend" => backend
76    )
77    .increment(1);
78}
79
80pub(super) fn record_connect_result(
81    transport: &'static str,
82    backend: &'static str,
83    result: &PgResult<PgConnection>,
84    elapsed: std::time::Duration,
85) {
86    let outcome = if result.is_ok() { "success" } else { "error" };
87    metrics::histogram!(
88        "qail_pg_connect_duration_seconds",
89        "transport" => transport,
90        "backend" => backend,
91        "outcome" => outcome
92    )
93    .record(elapsed.as_secs_f64());
94
95    if let Err(error) = result {
96        metrics::counter!(
97            "qail_pg_connect_failure_total",
98            "transport" => transport,
99            "backend" => backend,
100            "error_kind" => connect_error_kind(error)
101        )
102        .increment(1);
103    } else {
104        metrics::counter!(
105            "qail_pg_connect_success_total",
106            "transport" => transport,
107            "backend" => backend
108        )
109        .increment(1);
110    }
111}
112
113pub(super) fn select_scram_mechanism(
114    mechanisms: &[String],
115    tls_server_end_point_binding: Option<Vec<u8>>,
116    channel_binding_mode: ScramChannelBindingMode,
117) -> Result<(String, Option<Vec<u8>>), String> {
118    let has_scram = mechanisms.iter().any(|m| m == "SCRAM-SHA-256");
119    let has_scram_plus = mechanisms.iter().any(|m| m == "SCRAM-SHA-256-PLUS");
120
121    match channel_binding_mode {
122        ScramChannelBindingMode::Disable => {
123            if has_scram {
124                return Ok(("SCRAM-SHA-256".to_string(), None));
125            }
126            Err(format!(
127                "channel_binding=disable, but server does not advertise SCRAM-SHA-256. Available: {:?}",
128                mechanisms
129            ))
130        }
131        ScramChannelBindingMode::Prefer => {
132            if has_scram_plus {
133                if let Some(binding) = tls_server_end_point_binding {
134                    return Ok(("SCRAM-SHA-256-PLUS".to_string(), Some(binding)));
135                }
136
137                if has_scram {
138                    return Ok(("SCRAM-SHA-256".to_string(), None));
139                }
140
141                return Err(
142                    "Server requires SCRAM-SHA-256-PLUS but TLS channel binding is unavailable"
143                        .to_string(),
144                );
145            }
146
147            if has_scram {
148                return Ok(("SCRAM-SHA-256".to_string(), None));
149            }
150
151            Err(format!(
152                "Server doesn't support SCRAM-SHA-256. Available: {:?}",
153                mechanisms
154            ))
155        }
156        ScramChannelBindingMode::Require => {
157            if !has_scram_plus {
158                return Err(
159                    "channel_binding=require, but server does not advertise SCRAM-SHA-256-PLUS"
160                        .to_string(),
161                );
162            }
163            let binding = tls_server_end_point_binding.ok_or_else(|| {
164                "channel_binding=require, but TLS channel binding data is unavailable".to_string()
165            })?;
166            Ok(("SCRAM-SHA-256-PLUS".to_string(), Some(binding)))
167        }
168    }
169}
170
171/// PostgreSQL MD5 password response: `md5` + md5(hex(md5(password + user)) + 4-byte salt).
172pub(super) fn md5_password_message(user: &str, password: &str, salt: [u8; 4]) -> String {
173    use md5::{Digest, Md5};
174
175    let mut inner = Md5::new();
176    inner.update(password.as_bytes());
177    inner.update(user.as_bytes());
178    let inner_hex = format!("{:x}", inner.finalize());
179
180    let mut outer = Md5::new();
181    outer.update(inner_hex.as_bytes());
182    outer.update(salt);
183    format!("md5{:x}", outer.finalize())
184}
185
186/// Drop implementation sends Terminate packet if possible.
187/// This ensures proper cleanup even without explicit close() call.
188impl Drop for PgConnection {
189    fn drop(&mut self) {
190        // Try to send Terminate packet synchronously using try_write
191        // This is best-effort - if it fails, TCP RST will handle cleanup
192        let terminate: [u8; 5] = [b'X', 0, 0, 0, 4];
193
194        match &mut self.stream {
195            PgStream::Tcp(tcp) => {
196                // try_write is non-blocking
197                let _ = tcp.try_write(&terminate);
198            }
199            #[cfg(all(target_os = "linux", feature = "native-io-uring"))]
200            PgStream::Uring(stream) => {
201                // io_uring transport owns I/O in a dedicated worker thread;
202                // terminate packet in Drop is not viable, but force socket
203                // shutdown so in-flight waits unblock promptly.
204                let _ = stream.abort_inflight();
205            }
206            PgStream::Tls(_) => {
207                // TLS requires async write which we can't do in Drop.
208                // The TCP connection close will still notify the server.
209                // For graceful TLS shutdown, use connection.close() explicitly.
210            }
211            #[cfg(unix)]
212            PgStream::Unix(unix) => {
213                let _ = unix.try_write(&terminate);
214            }
215            #[cfg(all(feature = "enterprise-gssapi", target_os = "linux"))]
216            PgStream::GssEnc(_) => {
217                // GSSENC requires async wrap+write; skip in Drop.
218            }
219        }
220    }
221}
222
223fn command_tag_carries_affected_rows(command: &str) -> bool {
224    matches!(
225        command,
226        "COPY" | "DELETE" | "FETCH" | "INSERT" | "MERGE" | "MOVE" | "SELECT" | "UPDATE"
227    )
228}
229
230pub(crate) fn parse_affected_rows(tag: &str) -> PgResult<u64> {
231    let parts: Vec<&str> = tag.split_whitespace().collect();
232    let Some(command) = parts.first().copied() else {
233        return Ok(0);
234    };
235    if !command_tag_carries_affected_rows(command) {
236        return Ok(0);
237    }
238
239    let count = match command {
240        "INSERT" => {
241            if parts.len() != 3 {
242                return Err(PgError::Protocol(format!(
243                    "CommandComplete tag '{}' has malformed INSERT shape",
244                    tag
245                )));
246            }
247            parts[2]
248        }
249        _ => {
250            if parts.len() != 2 {
251                return Err(PgError::Protocol(format!(
252                    "CommandComplete tag '{}' has malformed affected-row shape",
253                    tag
254                )));
255            }
256            parts[1]
257        }
258    };
259
260    count.parse::<u64>().map_err(|_| {
261        PgError::Protocol(format!(
262            "CommandComplete tag '{}' has invalid affected row count",
263            tag
264        ))
265    })
266}
267
268#[cfg(all(target_os = "linux", feature = "native-io-uring"))]
269pub(super) fn should_try_uring_plain(io_uring: bool) -> bool {
270    super::super::io_backend::should_use_uring_plain_transport(io_uring)
271}
qail_pg/driver/connection/helpers.rs

qail_pg/driver/connection/
helpers.rs
