qail_pg/driver/connection/
types.rs

1//! PostgreSQL Connection
2//!
3//! Low-level TCP connection with wire protocol handling.
4//! This is Layer 3 (async I/O).
5//!
6//! Methods are split across modules for easier maintenance:
7//! - `io.rs` - Core I/O (send, recv)
8//! - `query.rs` - Query execution
9//! - `transaction.rs` - Transaction control
10//! - `cursor.rs` - Streaming cursors
11//! - `copy.rs` - COPY protocol
12//! - `pipeline.rs` - High-performance pipelining
13//! - `cancel.rs` - Query cancellation
14
15use super::super::notification::Notification;
16use super::super::stream::PgStream;
17use super::super::{AuthSettings, EnterpriseAuthMechanism};
18use crate::protocol::PROTOCOL_VERSION_3_2;
19use bytes::BytesMut;
20use std::collections::{HashMap, VecDeque};
21use std::num::NonZeroUsize;
22use std::sync::Arc;
23use std::sync::atomic::AtomicU64;
24use tokio::net::TcpStream;
25
26/// Statement cache capacity per connection.
27pub(super) const STMT_CACHE_CAPACITY: NonZeroUsize = NonZeroUsize::new(100).unwrap();
28
29/// Small, allocation-bounded prepared statement cache.
30///
31/// This mirrors the subset of `lru::LruCache` APIs used by the driver while
32/// avoiding external unsoundness advisories on `IterMut` (which we don't use).
33#[derive(Debug)]
34pub(crate) struct StatementCache {
35    capacity: NonZeroUsize,
36    entries: HashMap<u64, String>,
37    order: VecDeque<u64>, // Front = LRU, back = MRU
38}
39
40impl StatementCache {
41    pub(crate) fn new(capacity: NonZeroUsize) -> Self {
42        Self {
43            capacity,
44            entries: HashMap::with_capacity(capacity.get()),
45            order: VecDeque::with_capacity(capacity.get()),
46        }
47    }
48
49    pub(crate) fn len(&self) -> usize {
50        self.entries.len()
51    }
52
53    pub(crate) fn cap(&self) -> NonZeroUsize {
54        self.capacity
55    }
56
57    pub(crate) fn contains(&self, key: &u64) -> bool {
58        self.entries.contains_key(key)
59    }
60
61    pub(crate) fn get(&mut self, key: &u64) -> Option<String> {
62        let value = self.entries.get(key).cloned()?;
63        self.touch(*key);
64        Some(value)
65    }
66
67    pub(crate) fn put(&mut self, key: u64, value: String) {
68        if let std::collections::hash_map::Entry::Occupied(mut e) = self.entries.entry(key) {
69            e.insert(value);
70            self.touch(key);
71            return;
72        }
73
74        if self.entries.len() >= self.capacity.get() {
75            let _ = self.pop_lru();
76        }
77
78        self.entries.insert(key, value);
79        self.order.push_back(key);
80    }
81
82    pub(crate) fn pop_lru(&mut self) -> Option<(u64, String)> {
83        while let Some(key) = self.order.pop_front() {
84            if let Some(value) = self.entries.remove(&key) {
85                return Some((key, value));
86            }
87        }
88        None
89    }
90
91    pub(crate) fn remove(&mut self, key: &u64) -> Option<String> {
92        let removed = self.entries.remove(key);
93        if removed.is_some() {
94            self.order.retain(|k| k != key);
95        }
96        removed
97    }
98
99    pub(crate) fn clear(&mut self) {
100        self.entries.clear();
101        self.order.clear();
102    }
103
104    fn touch(&mut self, key: u64) {
105        self.order.retain(|k| *k != key);
106        self.order.push_back(key);
107    }
108}
109
110/// Initial buffer capacity (64KB for pipeline performance)
111pub(crate) const BUFFER_CAPACITY: usize = 65536;
112
113/// SSLRequest message bytes (request code: 80877103)
114pub(super) const SSL_REQUEST: [u8; 8] = [0, 0, 0, 8, 4, 210, 22, 47];
115
116/// GSSENCRequest message bytes (request code: 80877104)
117/// Byte breakdown: length=8 (00 00 00 08), code=80877104 (04 D2 16 30)
118pub(super) const GSSENC_REQUEST: [u8; 8] = [0, 0, 0, 8, 4, 210, 22, 48];
119
120/// Result of sending a GSSENCRequest to the server.
121#[derive(Debug)]
122pub(super) enum GssEncNegotiationResult {
123    /// Server responded 'G' — willing to perform GSSAPI encryption.
124    /// The TCP stream is returned for the caller to establish the
125    /// GSSAPI security context and wrap all subsequent traffic.
126    Accepted(TcpStream),
127    /// Server responded 'N' — unwilling to perform GSSAPI encryption.
128    Rejected,
129    /// Server sent an ErrorMessage — must not be displayed to user
130    /// (CVE-2024-10977: server not yet authenticated).
131    ServerError,
132}
133
134/// CancelRequest protocol code: 80877102
135pub(crate) const CANCEL_REQUEST_CODE: i32 = 80877102;
136
137/// Monotonic session id source for stateful GSS provider callbacks.
138pub(super) static GSS_SESSION_COUNTER: AtomicU64 = AtomicU64::new(1);
139
140/// Default timeout for TCP connect + PostgreSQL handshake.
141/// Prevents Slowloris DoS where a malicious server accepts TCP but never responds.
142pub(crate) const DEFAULT_CONNECT_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(10);
143pub(super) const CONNECT_TRANSPORT_PLAIN: &str = "plain";
144pub(super) const CONNECT_TRANSPORT_TLS: &str = "tls";
145pub(super) const CONNECT_TRANSPORT_MTLS: &str = "mtls";
146pub(super) const CONNECT_TRANSPORT_GSSENC: &str = "gssenc";
147pub(super) const CONNECT_BACKEND_TOKIO: &str = "tokio";
148#[cfg(all(target_os = "linux", feature = "io_uring"))]
149pub(super) const CONNECT_BACKEND_IO_URING: &str = "io_uring";
150
151/// TLS configuration for mutual TLS (client certificate authentication).
152#[derive(Debug, Clone)]
153pub struct TlsConfig {
154    /// Client certificate in PEM format
155    pub client_cert_pem: Vec<u8>,
156    /// Client private key in PEM format
157    pub client_key_pem: Vec<u8>,
158    /// Optional CA certificate for server verification (uses system certs if None)
159    pub ca_cert_pem: Option<Vec<u8>>,
160}
161
162impl TlsConfig {
163    /// Create a new TLS config from file paths.
164    pub fn from_files(
165        cert_path: impl AsRef<std::path::Path>,
166        key_path: impl AsRef<std::path::Path>,
167        ca_path: Option<impl AsRef<std::path::Path>>,
168    ) -> std::io::Result<Self> {
169        Ok(Self {
170            client_cert_pem: std::fs::read(cert_path)?,
171            client_key_pem: std::fs::read(key_path)?,
172            ca_cert_pem: ca_path.map(|p| std::fs::read(p)).transpose()?,
173        })
174    }
175}
176
177/// Bundled connection parameters for internal functions.
178///
179/// Groups the 8 common arguments to avoid exceeding clippy's
180/// `too_many_arguments` threshold.
181#[derive(Clone)]
182pub(super) struct ConnectParams<'a> {
183    pub(super) host: &'a str,
184    pub(super) port: u16,
185    pub(super) user: &'a str,
186    pub(super) database: &'a str,
187    pub(super) password: Option<&'a str>,
188    pub(super) auth_settings: AuthSettings,
189    pub(super) gss_token_provider: Option<super::super::GssTokenProvider>,
190    pub(super) gss_token_provider_ex: Option<super::super::GssTokenProviderEx>,
191    pub(super) protocol_minor: u16,
192    pub(super) startup_params: Vec<(String, String)>,
193}
194
195#[inline]
196pub(super) fn has_logical_replication_startup_mode(startup_params: &[(String, String)]) -> bool {
197    startup_params
198        .iter()
199        .any(|(k, v)| k.eq_ignore_ascii_case("replication") && v.eq_ignore_ascii_case("database"))
200}
201
202#[derive(Debug, Clone, Copy, PartialEq, Eq)]
203pub(super) enum StartupAuthFlow {
204    CleartextPassword,
205    Md5Password,
206    Scram { server_final_seen: bool },
207    EnterpriseGss { mechanism: EnterpriseAuthMechanism },
208}
209
210impl StartupAuthFlow {
211    pub(super) fn label(self) -> &'static str {
212        match self {
213            Self::CleartextPassword => "cleartext-password",
214            Self::Md5Password => "md5-password",
215            Self::Scram { .. } => "scram",
216            Self::EnterpriseGss { mechanism } => match mechanism {
217                EnterpriseAuthMechanism::KerberosV5 => "kerberos-v5",
218                EnterpriseAuthMechanism::GssApi => "gssapi",
219                EnterpriseAuthMechanism::Sspi => "sspi",
220            },
221        }
222    }
223}
224
225/// A raw PostgreSQL connection.
226pub struct PgConnection {
227    pub(crate) stream: PgStream,
228    pub(crate) buffer: BytesMut,
229    pub(crate) write_buf: BytesMut,
230    pub(crate) sql_buf: BytesMut,
231    pub(crate) params_buf: Vec<Option<Vec<u8>>>,
232    pub(crate) prepared_statements: HashMap<String, String>,
233    pub(crate) stmt_cache: StatementCache,
234    /// Cache of column metadata (RowDescription) per statement hash.
235    /// PostgreSQL only sends RowDescription after Parse, not on subsequent Bind+Execute.
236    /// This cache ensures by-name column access works even for cached prepared statements.
237    pub(crate) column_info_cache: HashMap<u64, Arc<super::super::ColumnInfo>>,
238    pub(crate) process_id: i32,
239    /// Legacy 4-byte cancel secret key (protocol 3.0-compatible wrappers).
240    ///
241    /// For protocol 3.2 extended key lengths, this remains `0` and callers
242    /// must use `cancel_key_bytes`.
243    pub(crate) secret_key: i32,
244    /// Full cancel key bytes (`4..=256`) from BackendKeyData.
245    pub(crate) cancel_key_bytes: Vec<u8>,
246    /// Startup protocol minor requested by this connection (for example `2` for 3.2).
247    pub(crate) requested_protocol_minor: u16,
248    /// Startup protocol minor negotiated with the server.
249    pub(crate) negotiated_protocol_minor: u16,
250    /// Buffer for asynchronous LISTEN/NOTIFY notifications.
251    /// Populated by `recv()` when it encounters NotificationResponse messages.
252    pub(crate) notifications: VecDeque<Notification>,
253    /// True while a logical replication CopyBoth stream is active.
254    pub(crate) replication_stream_active: bool,
255    /// True when StartupMessage was sent with `replication=database`.
256    pub(crate) replication_mode_enabled: bool,
257    /// Last seen wal_end from a replication XLogData frame.
258    pub(crate) last_replication_wal_end: Option<u64>,
259    /// Sticky fail-closed flag for uncertain protocol/I-O state.
260    /// Once set, the connection must not return to pool reuse.
261    pub(crate) io_desynced: bool,
262    /// Statement names scheduled for server-side `Close` on next write.
263    /// This keeps backend prepared state aligned with local LRU eviction.
264    pub(crate) pending_statement_closes: Vec<String>,
265    /// Reentrancy guard for pending-close drain path.
266    pub(crate) draining_statement_closes: bool,
267}
268
269impl PgConnection {
270    #[inline]
271    pub(crate) fn default_protocol_minor() -> u16 {
272        (PROTOCOL_VERSION_3_2 & 0xFFFF) as u16
273    }
274
275    /// Startup protocol minor requested by this connection.
276    #[inline]
277    pub fn requested_protocol_minor(&self) -> u16 {
278        self.requested_protocol_minor
279    }
280
281    /// Startup protocol minor negotiated with the server.
282    #[inline]
283    pub fn negotiated_protocol_minor(&self) -> u16 {
284        self.negotiated_protocol_minor
285    }
286}
qail_pg/driver/connection/types.rs

qail_pg/driver/connection/
types.rs
