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

qail_pg/driver/connection/
helpers.rs
