Skip to main content

qail_pg/driver/
io.rs

1//! Core I/O operations for PostgreSQL connection.
2//!
3//! This module provides low-level send/receive methods.
4
5use super::{PgBytesRow, PgConnection, PgError, PgResult, is_ignorable_session_message};
6use crate::protocol::{BackendMessage, FrontendMessage, PgEncoder};
7use bytes::{Bytes, BytesMut};
8use tokio::io::{AsyncReadExt, AsyncWriteExt};
9
10pub(crate) const MAX_MESSAGE_SIZE: usize = 64 * 1024 * 1024; // 64 MB — prevents OOM from malicious server messages
11
12/// Maximum undrained LISTEN/NOTIFY notifications buffered per connection.
13/// `MAX_MESSAGE_SIZE` caps a single frame; this caps the aggregate queue so a
14/// server streaming unsolicited NotificationResponse frames while the client
15/// is blocked in a receive loop cannot grow memory without bound.
16pub(crate) const MAX_BUFFERED_NOTIFICATIONS: usize = 8192;
17
18/// Maximum channel + payload bytes for a single buffered notification.
19/// A real PostgreSQL server caps NOTIFY payloads at ~8000 bytes and channel
20/// names at 63; a larger frame is a hostile peer inflating the queue, so the
21/// aggregate buffered memory stays bounded by
22/// `MAX_BUFFERED_NOTIFICATIONS * MAX_NOTIFICATION_BYTES`.
23pub(crate) const MAX_NOTIFICATION_BYTES: usize = 16 * 1024;
24
25/// Default read timeout for individual socket reads.
26/// Prevents Slowloris DoS where a server sends partial data then goes silent.
27const DEFAULT_READ_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(30);
28/// Default write timeout for individual socket writes/flushes.
29/// Prevents indefinitely blocked writes from pinning pool slots.
30const DEFAULT_WRITE_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(30);
31const READ_SPARE_LOW_WATERMARK: usize = 64 * 1024;
32
33#[inline]
34fn reserve_read_spare_capacity(buffer: &mut BytesMut) {
35    let spare = buffer.capacity().saturating_sub(buffer.len());
36    if spare < READ_SPARE_LOW_WATERMARK {
37        let target_spare = READ_SPARE_LOW_WATERMARK.max(buffer.capacity());
38        buffer.reserve(target_spare.saturating_sub(spare));
39    }
40}
41
42#[inline]
43fn parse_data_row_payload_owned(payload: &[u8]) -> PgResult<Vec<Option<Vec<u8>>>> {
44    if payload.len() < 2 {
45        return Err(PgError::Protocol("DataRow payload too short".into()));
46    }
47
48    let raw_count = i16::from_be_bytes([payload[0], payload[1]]);
49    if raw_count < 0 {
50        return Err(PgError::Protocol(format!(
51            "DataRow invalid column count: {}",
52            raw_count
53        )));
54    }
55    let column_count = raw_count as usize;
56    if column_count > (payload.len() - 2) / 4 + 1 {
57        return Err(PgError::Protocol(format!(
58            "DataRow claims {} columns but payload is only {} bytes",
59            column_count,
60            payload.len()
61        )));
62    }
63
64    let mut columns = Vec::with_capacity(column_count);
65    let mut pos = 2;
66    for _ in 0..column_count {
67        if pos + 4 > payload.len() {
68            return Err(PgError::Protocol(
69                "DataRow truncated: missing column length".into(),
70            ));
71        }
72
73        let len = i32::from_be_bytes([
74            payload[pos],
75            payload[pos + 1],
76            payload[pos + 2],
77            payload[pos + 3],
78        ]);
79        pos += 4;
80
81        if len == -1 {
82            columns.push(None);
83            continue;
84        }
85        if len < -1 {
86            return Err(PgError::Protocol(format!(
87                "DataRow invalid column length: {}",
88                len
89            )));
90        }
91
92        let len = len as usize;
93        if len > payload.len().saturating_sub(pos) {
94            return Err(PgError::Protocol(
95                "DataRow truncated: column data exceeds payload".into(),
96            ));
97        }
98        columns.push(Some(payload[pos..pos + len].to_vec()));
99        pos += len;
100    }
101
102    if pos != payload.len() {
103        return Err(PgError::Protocol("DataRow has trailing bytes".into()));
104    }
105
106    Ok(columns)
107}
108
109#[inline]
110fn parse_data_row_payload_reuse(
111    payload: &[u8],
112    columns: &mut Vec<Option<Vec<u8>>>,
113) -> PgResult<()> {
114    if payload.len() < 2 {
115        return Err(PgError::Protocol("DataRow payload too short".into()));
116    }
117
118    let raw_count = i16::from_be_bytes([payload[0], payload[1]]);
119    if raw_count < 0 {
120        return Err(PgError::Protocol(format!(
121            "DataRow invalid column count: {}",
122            raw_count
123        )));
124    }
125    let column_count = raw_count as usize;
126    if column_count > (payload.len() - 2) / 4 + 1 {
127        return Err(PgError::Protocol(format!(
128            "DataRow claims {} columns but payload is only {} bytes",
129            column_count,
130            payload.len()
131        )));
132    }
133
134    let previous_len = columns.len();
135    if previous_len < column_count {
136        columns.reserve(column_count - previous_len);
137    }
138
139    let mut pos = 2usize;
140    for idx in 0..column_count {
141        if pos + 4 > payload.len() {
142            return Err(PgError::Protocol(
143                "DataRow truncated: missing column length".into(),
144            ));
145        }
146
147        let len = i32::from_be_bytes([
148            payload[pos],
149            payload[pos + 1],
150            payload[pos + 2],
151            payload[pos + 3],
152        ]);
153        pos += 4;
154
155        if len == -1 {
156            if idx < previous_len {
157                columns[idx] = None;
158            } else {
159                columns.push(None);
160            }
161            continue;
162        }
163        if len < -1 {
164            return Err(PgError::Protocol(format!(
165                "DataRow invalid column length: {}",
166                len
167            )));
168        }
169
170        let len = len as usize;
171        if len > payload.len().saturating_sub(pos) {
172            return Err(PgError::Protocol(
173                "DataRow truncated: column data exceeds payload".into(),
174            ));
175        }
176        let value = &payload[pos..pos + len];
177        pos += len;
178
179        if idx < previous_len {
180            match &mut columns[idx] {
181                Some(buf) => {
182                    buf.clear();
183                    buf.extend_from_slice(value);
184                }
185                None => columns[idx] = Some(value.to_vec()),
186            }
187        } else {
188            columns.push(Some(value.to_vec()));
189        }
190    }
191
192    if columns.len() > column_count {
193        columns.truncate(column_count);
194    }
195
196    if pos != payload.len() {
197        return Err(PgError::Protocol("DataRow has trailing bytes".into()));
198    }
199
200    Ok(())
201}
202
203#[inline]
204fn parse_data_row_payload_zerocopy(payload: Bytes, row: &mut PgBytesRow) -> PgResult<()> {
205    if payload.len() < 2 {
206        return Err(PgError::Protocol("DataRow payload too short".into()));
207    }
208
209    let raw_count = i16::from_be_bytes([payload[0], payload[1]]);
210    if raw_count < 0 {
211        return Err(PgError::Protocol(format!(
212            "DataRow invalid column count: {}",
213            raw_count
214        )));
215    }
216    let column_count = raw_count as usize;
217    if column_count > (payload.len() - 2) / 4 + 1 {
218        return Err(PgError::Protocol(format!(
219            "DataRow claims {} columns but payload is only {} bytes",
220            column_count,
221            payload.len()
222        )));
223    }
224
225    row.payload = payload;
226    row.spans.clear();
227    if row.spans.capacity() < column_count {
228        row.spans.reserve(column_count - row.spans.capacity());
229    }
230
231    let mut pos = 2usize;
232    for _ in 0..column_count {
233        if pos + 4 > row.payload.len() {
234            return Err(PgError::Protocol(
235                "DataRow truncated: missing column length".into(),
236            ));
237        }
238
239        let len = i32::from_be_bytes([
240            row.payload[pos],
241            row.payload[pos + 1],
242            row.payload[pos + 2],
243            row.payload[pos + 3],
244        ]);
245        pos += 4;
246
247        if len == -1 {
248            row.spans.push(None);
249            continue;
250        }
251        if len < -1 {
252            return Err(PgError::Protocol(format!(
253                "DataRow invalid column length: {}",
254                len
255            )));
256        }
257
258        let len = len as usize;
259        if len > row.payload.len().saturating_sub(pos) {
260            return Err(PgError::Protocol(
261                "DataRow truncated: column data exceeds payload".into(),
262            ));
263        }
264        row.spans.push(Some((pos, len)));
265        pos += len;
266    }
267
268    if pos != row.payload.len() {
269        return Err(PgError::Protocol("DataRow has trailing bytes".into()));
270    }
271
272    Ok(())
273}
274
275#[inline]
276fn parse_first_column_payload_zerocopy(payload: Bytes) -> PgResult<Option<Bytes>> {
277    if payload.len() < 2 {
278        return Err(PgError::Protocol("DataRow payload too short".into()));
279    }
280
281    let raw_count = i16::from_be_bytes([payload[0], payload[1]]);
282    if raw_count < 0 {
283        return Err(PgError::Protocol(format!(
284            "DataRow invalid column count: {}",
285            raw_count
286        )));
287    }
288    let column_count = raw_count as usize;
289    if column_count > (payload.len() - 2) / 4 + 1 {
290        return Err(PgError::Protocol(format!(
291            "DataRow claims {} columns but payload is only {} bytes",
292            column_count,
293            payload.len()
294        )));
295    }
296
297    let mut pos = 2usize;
298    let mut first_column = None;
299
300    for idx in 0..column_count {
301        if pos + 4 > payload.len() {
302            return Err(PgError::Protocol(
303                "DataRow truncated: missing column length".into(),
304            ));
305        }
306
307        let len = i32::from_be_bytes([
308            payload[pos],
309            payload[pos + 1],
310            payload[pos + 2],
311            payload[pos + 3],
312        ]);
313        pos += 4;
314
315        if len == -1 {
316            if idx == 0 {
317                first_column = None;
318            }
319            continue;
320        }
321        if len < -1 {
322            return Err(PgError::Protocol(format!(
323                "DataRow invalid column length: {}",
324                len
325            )));
326        }
327
328        let len = len as usize;
329        if len > payload.len().saturating_sub(pos) {
330            return Err(PgError::Protocol(
331                "DataRow truncated: column data exceeds payload".into(),
332            ));
333        }
334
335        if idx == 0 {
336            first_column = Some(payload.slice(pos..pos + len));
337        }
338        pos += len;
339    }
340
341    if pos != payload.len() {
342        return Err(PgError::Protocol("DataRow has trailing bytes".into()));
343    }
344
345    Ok(first_column)
346}
347
348#[inline]
349fn parse_first_four_columns_payload_zerocopy(
350    payload: Bytes,
351    columns: &mut [Option<Bytes>; 4],
352) -> PgResult<()> {
353    if payload.len() < 2 {
354        return Err(PgError::Protocol("DataRow payload too short".into()));
355    }
356
357    let raw_count = i16::from_be_bytes([payload[0], payload[1]]);
358    if raw_count < 0 {
359        return Err(PgError::Protocol(format!(
360            "DataRow invalid column count: {}",
361            raw_count
362        )));
363    }
364    let column_count = raw_count as usize;
365    if column_count > (payload.len() - 2) / 4 + 1 {
366        return Err(PgError::Protocol(format!(
367            "DataRow claims {} columns but payload is only {} bytes",
368            column_count,
369            payload.len()
370        )));
371    }
372    if column_count != 4 {
373        return Err(PgError::Protocol(format!(
374            "DataRow fast-path expects exactly 4 columns, got {}",
375            column_count
376        )));
377    }
378
379    let mut pos = 2usize;
380    for slot in columns.iter_mut() {
381        if pos + 4 > payload.len() {
382            return Err(PgError::Protocol(
383                "DataRow truncated: missing column length".into(),
384            ));
385        }
386
387        let len = i32::from_be_bytes([
388            payload[pos],
389            payload[pos + 1],
390            payload[pos + 2],
391            payload[pos + 3],
392        ]);
393        pos += 4;
394
395        if len == -1 {
396            *slot = None;
397            continue;
398        }
399        if len < -1 {
400            return Err(PgError::Protocol(format!(
401                "DataRow invalid column length: {}",
402                len
403            )));
404        }
405
406        let len = len as usize;
407        if len > payload.len().saturating_sub(pos) {
408            return Err(PgError::Protocol(
409                "DataRow truncated: column data exceeds payload".into(),
410            ));
411        }
412        *slot = Some(payload.slice(pos..pos + len));
413        pos += len;
414    }
415
416    if pos != payload.len() {
417        return Err(PgError::Protocol("DataRow has trailing bytes".into()));
418    }
419
420    Ok(())
421}
422
423impl PgConnection {
424    #[inline]
425    fn stream_requires_flush(&self) -> bool {
426        use super::stream::PgStream;
427
428        match &self.stream {
429            PgStream::Tcp(_) => false,
430            PgStream::Tls(_) => true,
431            #[cfg(all(target_os = "linux", feature = "native-io-uring"))]
432            PgStream::Uring(_) => false,
433            #[cfg(unix)]
434            PgStream::Unix(_) => false,
435            #[cfg(all(feature = "enterprise-gssapi", target_os = "linux"))]
436            PgStream::GssEnc(_) => true,
437        }
438    }
439
440    #[inline]
441    pub(crate) fn mark_io_desynced(&mut self) {
442        self.io_desynced = true;
443    }
444
445    #[inline]
446    pub(crate) fn is_io_desynced(&self) -> bool {
447        self.io_desynced
448    }
449
450    #[inline]
451    fn protocol_desync<T>(&mut self, msg: String) -> PgResult<T> {
452        self.mark_io_desynced();
453        Err(PgError::Protocol(msg))
454    }
455
456    #[inline]
457    fn protocol_desync_error<T>(&mut self, err: PgError) -> PgResult<T> {
458        match err {
459            PgError::Protocol(msg) => self.protocol_desync(msg),
460            err => {
461                self.mark_io_desynced();
462                Err(err)
463            }
464        }
465    }
466
467    #[inline]
468    fn connection_desync<T>(&mut self, msg: String) -> PgResult<T> {
469        self.mark_io_desynced();
470        Err(PgError::Connection(msg))
471    }
472
473    /// Send queued statement `Close` messages and drain until `ReadyForQuery`.
474    ///
475    /// We ignore `26000 prepared statement ... does not exist` because this
476    /// can happen after failover or server-side invalidation, and in that case
477    /// local state is already being reconciled by retry paths.
478    async fn flush_pending_statement_closes(&mut self) -> PgResult<()> {
479        if self.draining_statement_closes || self.pending_statement_closes.is_empty() {
480            return Ok(());
481        }
482
483        self.draining_statement_closes = true;
484        let close_names = std::mem::take(&mut self.pending_statement_closes);
485
486        let estimated_payload_len: usize = close_names
487            .iter()
488            .map(|name| 16usize.saturating_add(name.len()))
489            .sum();
490        let mut buf = BytesMut::with_capacity(estimated_payload_len.saturating_add(5));
491        for stmt_name in &close_names {
492            let close_msg = PgEncoder::try_encode_close(false, stmt_name)
493                .map_err(|e| PgError::Encode(e.to_string()))?;
494            buf.extend_from_slice(&close_msg);
495        }
496        PgEncoder::encode_sync_to(&mut buf);
497
498        if let Err(err) = self
499            .write_all_with_timeout_inner(&buf, "pending statement close write")
500            .await
501        {
502            self.draining_statement_closes = false;
503            return Err(err);
504        }
505        if let Err(err) = self
506            .flush_with_timeout("pending statement close flush")
507            .await
508        {
509            self.draining_statement_closes = false;
510            return Err(err);
511        }
512
513        let mut error: Option<PgError> = None;
514        loop {
515            let msg = match self.recv().await {
516                Ok(msg) => msg,
517                Err(err) => {
518                    self.draining_statement_closes = false;
519                    return Err(err);
520                }
521            };
522            match msg {
523                BackendMessage::CloseComplete => {}
524                BackendMessage::ReadyForQuery(_) => {
525                    self.draining_statement_closes = false;
526                    if let Some(err) = error {
527                        return Err(err);
528                    }
529                    return Ok(());
530                }
531                BackendMessage::ErrorResponse(err_fields) => {
532                    if error.is_none() {
533                        let code_26000 = err_fields.code.eq_ignore_ascii_case("26000");
534                        let msg_lower = err_fields.message.to_ascii_lowercase();
535                        let missing_prepared = msg_lower.contains("prepared statement")
536                            && msg_lower.contains("does not exist");
537                        if !(code_26000 && missing_prepared) {
538                            error = Some(PgError::QueryServer(err_fields.into()));
539                        }
540                    }
541                }
542                msg if is_ignorable_session_message(&msg) => {}
543                other => {
544                    self.draining_statement_closes = false;
545                    return self.protocol_desync(format!(
546                        "Unexpected backend message during pending statement close drain: {:?}",
547                        other
548                    ));
549                }
550            }
551        }
552    }
553
554    /// Write all bytes with a timeout guard.
555    ///
556    /// Prevents stuck kernel send buffers or dead sockets from hanging forever.
557    pub(crate) async fn write_all_with_timeout(
558        &mut self,
559        bytes: &[u8],
560        operation: &str,
561    ) -> PgResult<()> {
562        if !self.draining_statement_closes && !self.pending_statement_closes.is_empty() {
563            self.flush_pending_statement_closes().await?;
564        }
565        self.write_all_with_timeout_inner(bytes, operation).await
566    }
567
568    async fn write_all_with_timeout_inner(
569        &mut self,
570        bytes: &[u8],
571        operation: &str,
572    ) -> PgResult<()> {
573        if bytes.is_empty() {
574            return Err(PgError::Encode(
575                "refusing to send empty frontend payload".to_string(),
576            ));
577        }
578        use super::stream::PgStream;
579        let mut mark_desync = false;
580        let result = match &mut self.stream {
581            PgStream::Tcp(stream) => {
582                match tokio::time::timeout(DEFAULT_WRITE_TIMEOUT, stream.write_all(bytes)).await {
583                    Ok(Ok(())) => Ok(()),
584                    Ok(Err(e)) => {
585                        mark_desync = true;
586                        Err(PgError::Connection(format!("Write error: {}", e)))
587                    }
588                    Err(_) => {
589                        mark_desync = true;
590                        Err(PgError::Timeout(format!(
591                            "{} timeout after {:?}",
592                            operation, DEFAULT_WRITE_TIMEOUT
593                        )))
594                    }
595                }
596            }
597            PgStream::Tls(stream) => {
598                match tokio::time::timeout(DEFAULT_WRITE_TIMEOUT, stream.write_all(bytes)).await {
599                    Ok(Ok(())) => Ok(()),
600                    Ok(Err(e)) => {
601                        mark_desync = true;
602                        Err(PgError::Connection(format!("Write error: {}", e)))
603                    }
604                    Err(_) => {
605                        mark_desync = true;
606                        Err(PgError::Timeout(format!(
607                            "{} timeout after {:?}",
608                            operation, DEFAULT_WRITE_TIMEOUT
609                        )))
610                    }
611                }
612            }
613            #[cfg(all(target_os = "linux", feature = "native-io-uring"))]
614            PgStream::Uring(stream) => {
615                match tokio::time::timeout(DEFAULT_WRITE_TIMEOUT, stream.write_all(bytes)).await {
616                    Ok(Ok(())) => Ok(()),
617                    Ok(Err(e)) => {
618                        mark_desync = true;
619                        Err(PgError::Connection(format!("Write error: {}", e)))
620                    }
621                    Err(_) => {
622                        mark_desync = true;
623                        let _ = stream.abort_inflight();
624                        Err(PgError::Timeout(format!(
625                            "{} timeout after {:?}",
626                            operation, DEFAULT_WRITE_TIMEOUT
627                        )))
628                    }
629                }
630            }
631            #[cfg(unix)]
632            PgStream::Unix(stream) => {
633                match tokio::time::timeout(DEFAULT_WRITE_TIMEOUT, stream.write_all(bytes)).await {
634                    Ok(Ok(())) => Ok(()),
635                    Ok(Err(e)) => {
636                        mark_desync = true;
637                        Err(PgError::Connection(format!("Write error: {}", e)))
638                    }
639                    Err(_) => {
640                        mark_desync = true;
641                        Err(PgError::Timeout(format!(
642                            "{} timeout after {:?}",
643                            operation, DEFAULT_WRITE_TIMEOUT
644                        )))
645                    }
646                }
647            }
648            #[cfg(all(feature = "enterprise-gssapi", target_os = "linux"))]
649            PgStream::GssEnc(stream) => {
650                match tokio::time::timeout(DEFAULT_WRITE_TIMEOUT, stream.write_all(bytes)).await {
651                    Ok(Ok(())) => Ok(()),
652                    Ok(Err(e)) => {
653                        mark_desync = true;
654                        Err(PgError::Connection(format!("Write error: {}", e)))
655                    }
656                    Err(_) => {
657                        mark_desync = true;
658                        Err(PgError::Timeout(format!(
659                            "{} timeout after {:?}",
660                            operation, DEFAULT_WRITE_TIMEOUT
661                        )))
662                    }
663                }
664            }
665        };
666        if mark_desync {
667            self.mark_io_desynced();
668        }
669        result
670    }
671
672    /// Flush with a timeout guard.
673    pub(crate) async fn flush_with_timeout(&mut self, operation: &str) -> PgResult<()> {
674        if !self.stream_requires_flush() {
675            return Ok(());
676        }
677
678        use super::stream::PgStream;
679        let mut mark_desync = false;
680        let result = match &mut self.stream {
681            PgStream::Tcp(stream) => {
682                match tokio::time::timeout(DEFAULT_WRITE_TIMEOUT, stream.flush()).await {
683                    Ok(Ok(())) => Ok(()),
684                    Ok(Err(e)) => {
685                        mark_desync = true;
686                        Err(PgError::Connection(format!("Flush error: {}", e)))
687                    }
688                    Err(_) => {
689                        mark_desync = true;
690                        Err(PgError::Timeout(format!(
691                            "{} timeout after {:?}",
692                            operation, DEFAULT_WRITE_TIMEOUT
693                        )))
694                    }
695                }
696            }
697            PgStream::Tls(stream) => {
698                match tokio::time::timeout(DEFAULT_WRITE_TIMEOUT, stream.flush()).await {
699                    Ok(Ok(())) => Ok(()),
700                    Ok(Err(e)) => {
701                        mark_desync = true;
702                        Err(PgError::Connection(format!("Flush error: {}", e)))
703                    }
704                    Err(_) => {
705                        mark_desync = true;
706                        Err(PgError::Timeout(format!(
707                            "{} timeout after {:?}",
708                            operation, DEFAULT_WRITE_TIMEOUT
709                        )))
710                    }
711                }
712            }
713            #[cfg(all(target_os = "linux", feature = "native-io-uring"))]
714            PgStream::Uring(stream) => {
715                match tokio::time::timeout(DEFAULT_WRITE_TIMEOUT, stream.flush()).await {
716                    Ok(Ok(())) => Ok(()),
717                    Ok(Err(e)) => {
718                        mark_desync = true;
719                        Err(PgError::Connection(format!("Flush error: {}", e)))
720                    }
721                    Err(_) => {
722                        mark_desync = true;
723                        let _ = stream.abort_inflight();
724                        Err(PgError::Timeout(format!(
725                            "{} timeout after {:?}",
726                            operation, DEFAULT_WRITE_TIMEOUT
727                        )))
728                    }
729                }
730            }
731            #[cfg(unix)]
732            PgStream::Unix(stream) => {
733                match tokio::time::timeout(DEFAULT_WRITE_TIMEOUT, stream.flush()).await {
734                    Ok(Ok(())) => Ok(()),
735                    Ok(Err(e)) => {
736                        mark_desync = true;
737                        Err(PgError::Connection(format!("Flush error: {}", e)))
738                    }
739                    Err(_) => {
740                        mark_desync = true;
741                        Err(PgError::Timeout(format!(
742                            "{} timeout after {:?}",
743                            operation, DEFAULT_WRITE_TIMEOUT
744                        )))
745                    }
746                }
747            }
748            #[cfg(all(feature = "enterprise-gssapi", target_os = "linux"))]
749            PgStream::GssEnc(stream) => {
750                match tokio::time::timeout(DEFAULT_WRITE_TIMEOUT, stream.flush()).await {
751                    Ok(Ok(())) => Ok(()),
752                    Ok(Err(e)) => {
753                        mark_desync = true;
754                        Err(PgError::Connection(format!("Flush error: {}", e)))
755                    }
756                    Err(_) => {
757                        mark_desync = true;
758                        Err(PgError::Timeout(format!(
759                            "{} timeout after {:?}",
760                            operation, DEFAULT_WRITE_TIMEOUT
761                        )))
762                    }
763                }
764            }
765        };
766        if mark_desync {
767            self.mark_io_desynced();
768        }
769        result
770    }
771
772    /// Send a frontend message.
773    pub async fn send(&mut self, msg: FrontendMessage) -> PgResult<()> {
774        let bytes = msg
775            .encode_checked()
776            .map_err(|e| PgError::Encode(e.to_string()))?;
777        self.send_bytes(&bytes).await?;
778        Ok(())
779    }
780
781    /// Buffer an async notification, failing closed once the undrained queue
782    /// reaches [`MAX_BUFFERED_NOTIFICATIONS`] or a single notification
783    /// exceeds [`MAX_NOTIFICATION_BYTES`].
784    pub(crate) fn buffer_notification(
785        &mut self,
786        notification: super::notification::Notification,
787    ) -> PgResult<()> {
788        let size = notification.channel.len() + notification.payload.len();
789        if size > MAX_NOTIFICATION_BYTES {
790            return self.protocol_desync(format!(
791                "notification exceeds {} bytes ({} received)",
792                MAX_NOTIFICATION_BYTES, size
793            ));
794        }
795        if self.notifications.len() >= MAX_BUFFERED_NOTIFICATIONS {
796            tracing::warn!(
797                channel = %notification.channel,
798                buffered = self.notifications.len(),
799                "notification_buffer_overflow: undrained queue at cap, desyncing connection"
800            );
801            return self.protocol_desync(format!(
802                "notification buffer overflow: {} undrained notifications (max {})",
803                self.notifications.len(),
804                MAX_BUFFERED_NOTIFICATIONS
805            ));
806        }
807        self.notifications.push_back(notification);
808        Ok(())
809    }
810
811    /// Loops until a complete message is available.
812    /// Automatically buffers NotificationResponse messages for LISTEN/NOTIFY.
813    pub async fn recv(&mut self) -> PgResult<BackendMessage> {
814        loop {
815            // Try to decode from buffer first
816            if self.buffer.len() >= 5 {
817                let msg_len = u32::from_be_bytes([
818                    self.buffer[1],
819                    self.buffer[2],
820                    self.buffer[3],
821                    self.buffer[4],
822                ]) as usize;
823
824                if msg_len < 4 {
825                    return self.protocol_desync(format!(
826                        "Invalid message length: {} (minimum 4)",
827                        msg_len
828                    ));
829                }
830
831                if msg_len > MAX_MESSAGE_SIZE {
832                    return self.protocol_desync(format!(
833                        "Message too large: {} bytes (max {})",
834                        msg_len, MAX_MESSAGE_SIZE
835                    ));
836                }
837
838                if self.buffer.len() > msg_len {
839                    // We have a complete message - zero-copy split
840                    let msg_bytes = self.buffer.split_to(msg_len + 1);
841                    let (msg, _) = match BackendMessage::decode(&msg_bytes) {
842                        Ok(decoded) => decoded,
843                        Err(e) => return self.protocol_desync(e),
844                    };
845
846                    // Intercept async notifications — buffer them instead of returning
847                    if let BackendMessage::NotificationResponse {
848                        process_id,
849                        channel,
850                        payload,
851                    } = msg
852                    {
853                        self.buffer_notification(super::notification::Notification {
854                            process_id,
855                            channel,
856                            payload,
857                        })?;
858                        continue; // Keep reading for the actual response
859                    }
860
861                    return Ok(msg);
862                }
863            }
864
865            let n = self.read_with_timeout().await?;
866            if n == 0 {
867                return self.connection_desync("Connection closed".to_string());
868            }
869        }
870    }
871
872    /// Receive a backend message with idle-friendly timeout behavior.
873    ///
874    /// For long-lived idle streams (e.g. logical replication), an empty
875    /// buffer uses no-timeout reads so inactivity does not fail the stream.
876    /// If a backend frame is already partially buffered, switch back to the
877    /// normal read timeout to fail-closed on partial-frame stalls.
878    pub(crate) async fn recv_without_timeout(&mut self) -> PgResult<BackendMessage> {
879        loop {
880            if self.buffer.len() >= 5 {
881                let msg_len = u32::from_be_bytes([
882                    self.buffer[1],
883                    self.buffer[2],
884                    self.buffer[3],
885                    self.buffer[4],
886                ]) as usize;
887
888                if msg_len < 4 {
889                    return self.protocol_desync(format!(
890                        "Invalid message length: {} (minimum 4)",
891                        msg_len
892                    ));
893                }
894
895                if msg_len > MAX_MESSAGE_SIZE {
896                    return self.protocol_desync(format!(
897                        "Message too large: {} bytes (max {})",
898                        msg_len, MAX_MESSAGE_SIZE
899                    ));
900                }
901
902                if self.buffer.len() > msg_len {
903                    let msg_bytes = self.buffer.split_to(msg_len + 1);
904                    let (msg, _) = match BackendMessage::decode(&msg_bytes) {
905                        Ok(decoded) => decoded,
906                        Err(e) => return self.protocol_desync(e),
907                    };
908
909                    if let BackendMessage::NotificationResponse {
910                        process_id,
911                        channel,
912                        payload,
913                    } = msg
914                    {
915                        self.buffer_notification(super::notification::Notification {
916                            process_id,
917                            channel,
918                            payload,
919                        })?;
920                        continue;
921                    }
922
923                    return Ok(msg);
924                }
925            }
926
927            let n = if self.buffer.is_empty() {
928                self.read_without_timeout().await?
929            } else {
930                self.read_with_timeout().await?
931            };
932            if n == 0 {
933                return self.connection_desync("Connection closed".to_string());
934            }
935        }
936    }
937
938    /// Read from the socket with a timeout guard.
939    /// Returns the number of bytes read, or an error if the timeout fires.
940    /// This prevents Slowloris DoS attacks where a malicious server sends
941    /// partial data then goes silent, causing the driver to hang forever.
942    #[inline]
943    pub(crate) async fn read_with_timeout(&mut self) -> PgResult<usize> {
944        reserve_read_spare_capacity(&mut self.buffer);
945
946        use super::stream::PgStream;
947        let (stream, buffer) = (&mut self.stream, &mut self.buffer);
948        let mut mark_desync = false;
949        let result = match stream {
950            PgStream::Tcp(stream) => {
951                match tokio::time::timeout(DEFAULT_READ_TIMEOUT, stream.read_buf(buffer)).await {
952                    Ok(Ok(n)) => Ok(n),
953                    Ok(Err(e)) => {
954                        mark_desync = true;
955                        Err(PgError::Connection(format!("Read error: {}", e)))
956                    }
957                    Err(_) => {
958                        mark_desync = true;
959                        Err(PgError::Connection(format!(
960                            "Read timeout after {:?} — possible Slowloris attack or dead connection",
961                            DEFAULT_READ_TIMEOUT
962                        )))
963                    }
964                }
965            }
966            PgStream::Tls(stream) => {
967                match tokio::time::timeout(DEFAULT_READ_TIMEOUT, stream.read_buf(buffer)).await {
968                    Ok(Ok(n)) => Ok(n),
969                    Ok(Err(e)) => {
970                        mark_desync = true;
971                        Err(PgError::Connection(format!("Read error: {}", e)))
972                    }
973                    Err(_) => {
974                        mark_desync = true;
975                        Err(PgError::Connection(format!(
976                            "Read timeout after {:?} — possible Slowloris attack or dead connection",
977                            DEFAULT_READ_TIMEOUT
978                        )))
979                    }
980                }
981            }
982            #[cfg(all(target_os = "linux", feature = "native-io-uring"))]
983            PgStream::Uring(stream) => {
984                match tokio::time::timeout(DEFAULT_READ_TIMEOUT, stream.read_into(buffer, 131072))
985                    .await
986                {
987                    Ok(Ok(n)) => Ok(n),
988                    Ok(Err(e)) => {
989                        mark_desync = true;
990                        Err(PgError::Connection(format!("Read error: {}", e)))
991                    }
992                    Err(_) => {
993                        mark_desync = true;
994                        let _ = stream.abort_inflight();
995                        Err(PgError::Connection(format!(
996                            "Read timeout after {:?} — possible Slowloris attack or dead connection",
997                            DEFAULT_READ_TIMEOUT
998                        )))
999                    }
1000                }
1001            }
1002            #[cfg(unix)]
1003            PgStream::Unix(stream) => {
1004                match tokio::time::timeout(DEFAULT_READ_TIMEOUT, stream.read_buf(buffer)).await {
1005                    Ok(Ok(n)) => Ok(n),
1006                    Ok(Err(e)) => {
1007                        mark_desync = true;
1008                        Err(PgError::Connection(format!("Read error: {}", e)))
1009                    }
1010                    Err(_) => {
1011                        mark_desync = true;
1012                        Err(PgError::Connection(format!(
1013                            "Read timeout after {:?} — possible Slowloris attack or dead connection",
1014                            DEFAULT_READ_TIMEOUT
1015                        )))
1016                    }
1017                }
1018            }
1019            #[cfg(all(feature = "enterprise-gssapi", target_os = "linux"))]
1020            PgStream::GssEnc(stream) => {
1021                match tokio::time::timeout(DEFAULT_READ_TIMEOUT, stream.read_buf(buffer)).await {
1022                    Ok(Ok(n)) => Ok(n),
1023                    Ok(Err(e)) => {
1024                        mark_desync = true;
1025                        Err(PgError::Connection(format!("Read error: {}", e)))
1026                    }
1027                    Err(_) => {
1028                        mark_desync = true;
1029                        Err(PgError::Connection(format!(
1030                            "Read timeout after {:?} — possible Slowloris attack or dead connection",
1031                            DEFAULT_READ_TIMEOUT
1032                        )))
1033                    }
1034                }
1035            }
1036        };
1037        if mark_desync {
1038            self.mark_io_desynced();
1039        }
1040        result
1041    }
1042
1043    /// Read from socket without timeout guard.
1044    ///
1045    /// Used for long-idle LISTEN/NOTIFY connections.
1046    pub(crate) async fn read_without_timeout(&mut self) -> PgResult<usize> {
1047        reserve_read_spare_capacity(&mut self.buffer);
1048
1049        use super::stream::PgStream;
1050        let (stream, buffer) = (&mut self.stream, &mut self.buffer);
1051        let read_result = match stream {
1052            PgStream::Tcp(stream) => stream.read_buf(buffer).await,
1053            PgStream::Tls(stream) => stream.read_buf(buffer).await,
1054            #[cfg(all(target_os = "linux", feature = "native-io-uring"))]
1055            PgStream::Uring(stream) => stream.read_into(buffer, 131072).await,
1056            #[cfg(unix)]
1057            PgStream::Unix(stream) => stream.read_buf(buffer).await,
1058            #[cfg(all(feature = "enterprise-gssapi", target_os = "linux"))]
1059            PgStream::GssEnc(stream) => stream.read_buf(buffer).await,
1060        };
1061
1062        match read_result {
1063            Ok(n) => Ok(n),
1064            Err(e) => {
1065                self.mark_io_desynced();
1066                Err(PgError::Connection(format!("Read error: {}", e)))
1067            }
1068        }
1069    }
1070
1071    /// Send raw bytes to the stream.
1072    /// Includes flush for TLS safety — TLS buffers internally and
1073    /// needs flush to push encrypted data to the underlying TCP socket.
1074    pub async fn send_bytes(&mut self, bytes: &[u8]) -> PgResult<()> {
1075        self.write_all_with_timeout(bytes, "send raw bytes").await?;
1076        self.flush_with_timeout("flush raw bytes").await?;
1077        Ok(())
1078    }
1079
1080    #[inline]
1081    fn decode_fast_message_type(&mut self, msg_bytes: BytesMut) -> PgResult<Option<u8>> {
1082        let msg_type = msg_bytes[0];
1083        let (msg, _) = match BackendMessage::decode(&msg_bytes) {
1084            Ok(decoded) => decoded,
1085            Err(e) => return self.protocol_desync(e),
1086        };
1087        match msg {
1088            BackendMessage::ErrorResponse(err) => Err(PgError::QueryServer(err.into())),
1089            BackendMessage::NotificationResponse {
1090                process_id,
1091                channel,
1092                payload,
1093            } => {
1094                self.buffer_notification(super::notification::Notification {
1095                    process_id,
1096                    channel,
1097                    payload,
1098                })?;
1099                Ok(None)
1100            }
1101            _ => Ok(Some(msg_type)),
1102        }
1103    }
1104
1105    // ==================== BUFFERED WRITE API (High Performance) ====================
1106
1107    /// Buffer bytes for later flush (NO SYSCALL).
1108    /// Use flush_write_buf() to send all buffered data.
1109    #[inline]
1110    pub fn buffer_bytes(&mut self, bytes: &[u8]) {
1111        self.write_buf.extend_from_slice(bytes);
1112    }
1113
1114    /// Flush the write buffer to the stream (single write_all + flush).
1115    /// The flush is critical for TLS connections.
1116    pub async fn flush_write_buf(&mut self) -> PgResult<()> {
1117        if !self.write_buf.is_empty() {
1118            let payload = self.write_buf.split().freeze();
1119            self.write_all_with_timeout(&payload, "flush write buffer")
1120                .await?;
1121            self.flush_with_timeout("flush write buffer").await?;
1122        }
1123        Ok(())
1124    }
1125
1126    /// FAST receive - returns only message type byte, skips parsing.
1127    /// This is ~10x faster than recv() for pipelining benchmarks.
1128    /// Returns: message_type
1129    #[inline]
1130    pub(crate) async fn recv_msg_type_fast(&mut self) -> PgResult<u8> {
1131        loop {
1132            if self.buffer.len() >= 5 {
1133                let msg_len = u32::from_be_bytes([
1134                    self.buffer[1],
1135                    self.buffer[2],
1136                    self.buffer[3],
1137                    self.buffer[4],
1138                ]) as usize;
1139
1140                if msg_len < 4 {
1141                    return self.protocol_desync(format!(
1142                        "Invalid message length: {} (minimum 4)",
1143                        msg_len
1144                    ));
1145                }
1146
1147                if msg_len > MAX_MESSAGE_SIZE {
1148                    return self.protocol_desync(format!(
1149                        "Message too large: {} bytes (max {})",
1150                        msg_len, MAX_MESSAGE_SIZE
1151                    ));
1152                }
1153
1154                if self.buffer.len() > msg_len {
1155                    let msg_bytes = self.buffer.split_to(msg_len + 1);
1156                    if let Some(msg_type) = self.decode_fast_message_type(msg_bytes)? {
1157                        return Ok(msg_type);
1158                    }
1159                    continue;
1160                }
1161            }
1162
1163            let n = self.read_with_timeout().await?;
1164            if n == 0 {
1165                return self.connection_desync("Connection closed".to_string());
1166            }
1167        }
1168    }
1169
1170    /// FAST receive for result consumption - inline DataRow parsing.
1171    /// Returns: (msg_type, Option<row_data>)
1172    /// For 'D' (DataRow): returns parsed columns
1173    /// For other types: returns None
1174    /// This avoids BackendMessage enum allocation for non-DataRow messages.
1175    #[inline]
1176    pub(crate) async fn recv_with_data_fast(
1177        &mut self,
1178    ) -> PgResult<(u8, Option<Vec<Option<Vec<u8>>>>)> {
1179        loop {
1180            if self.buffer.len() >= 5 {
1181                let msg_len = u32::from_be_bytes([
1182                    self.buffer[1],
1183                    self.buffer[2],
1184                    self.buffer[3],
1185                    self.buffer[4],
1186                ]) as usize;
1187
1188                if msg_len < 4 {
1189                    return self.protocol_desync(format!(
1190                        "Invalid message length: {} (minimum 4)",
1191                        msg_len
1192                    ));
1193                }
1194
1195                if msg_len > MAX_MESSAGE_SIZE {
1196                    return self.protocol_desync(format!(
1197                        "Message too large: {} bytes (max {})",
1198                        msg_len, MAX_MESSAGE_SIZE
1199                    ));
1200                }
1201
1202                if self.buffer.len() > msg_len {
1203                    let msg_type = self.buffer[0];
1204
1205                    // Fast path: DataRow - parse inline
1206                    if msg_type == b'D' {
1207                        let parse_result = {
1208                            let payload = &self.buffer[5..msg_len + 1];
1209                            parse_data_row_payload_owned(payload)
1210                        };
1211
1212                        let _ = self.buffer.split_to(msg_len + 1);
1213                        match parse_result {
1214                            Ok(columns) => return Ok((msg_type, Some(columns))),
1215                            Err(err) => return self.protocol_desync_error(err),
1216                        }
1217                    }
1218
1219                    let msg_bytes = self.buffer.split_to(msg_len + 1);
1220                    if let Some(msg_type) = self.decode_fast_message_type(msg_bytes)? {
1221                        return Ok((msg_type, None));
1222                    }
1223                    continue;
1224                }
1225            }
1226
1227            let n = self.read_with_timeout().await?;
1228            if n == 0 {
1229                return self.connection_desync("Connection closed".to_string());
1230            }
1231        }
1232    }
1233
1234    /// FAST receive for result consumption into a reusable row buffer.
1235    ///
1236    /// This preserves owned row semantics while reusing allocations across
1237    /// `DataRow` messages.
1238    #[inline]
1239    pub(crate) async fn recv_fill_data_row_fast(
1240        &mut self,
1241        row_buf: &mut Vec<Option<Vec<u8>>>,
1242    ) -> PgResult<u8> {
1243        loop {
1244            if self.buffer.len() >= 5 {
1245                let msg_len = u32::from_be_bytes([
1246                    self.buffer[1],
1247                    self.buffer[2],
1248                    self.buffer[3],
1249                    self.buffer[4],
1250                ]) as usize;
1251
1252                if msg_len < 4 {
1253                    return self.protocol_desync(format!(
1254                        "Invalid message length: {} (minimum 4)",
1255                        msg_len
1256                    ));
1257                }
1258
1259                if msg_len > MAX_MESSAGE_SIZE {
1260                    return self.protocol_desync(format!(
1261                        "Message too large: {} bytes (max {})",
1262                        msg_len, MAX_MESSAGE_SIZE
1263                    ));
1264                }
1265
1266                if self.buffer.len() > msg_len {
1267                    let msg_type = self.buffer[0];
1268
1269                    if msg_type == b'D' {
1270                        let parse_result = {
1271                            let payload = &self.buffer[5..msg_len + 1];
1272                            parse_data_row_payload_reuse(payload, row_buf)
1273                        };
1274
1275                        let _ = self.buffer.split_to(msg_len + 1);
1276                        if let Err(err) = parse_result {
1277                            return self.protocol_desync_error(err);
1278                        }
1279                        return Ok(msg_type);
1280                    }
1281
1282                    let msg_bytes = self.buffer.split_to(msg_len + 1);
1283                    if let Some(msg_type) = self.decode_fast_message_type(msg_bytes)? {
1284                        return Ok(msg_type);
1285                    }
1286                    continue;
1287                }
1288            }
1289
1290            let n = self.read_with_timeout().await?;
1291            if n == 0 {
1292                return self.connection_desync("Connection closed".to_string());
1293            }
1294        }
1295    }
1296
1297    /// FAST receive for result consumption into a reusable zero-copy row.
1298    #[inline]
1299    pub(crate) async fn recv_fill_zerocopy_row_fast(
1300        &mut self,
1301        row: &mut PgBytesRow,
1302    ) -> PgResult<u8> {
1303        loop {
1304            if self.buffer.len() >= 5 {
1305                let msg_len = u32::from_be_bytes([
1306                    self.buffer[1],
1307                    self.buffer[2],
1308                    self.buffer[3],
1309                    self.buffer[4],
1310                ]) as usize;
1311
1312                if msg_len < 4 {
1313                    return self.protocol_desync(format!(
1314                        "Invalid message length: {} (minimum 4)",
1315                        msg_len
1316                    ));
1317                }
1318
1319                if msg_len > MAX_MESSAGE_SIZE {
1320                    return self.protocol_desync(format!(
1321                        "Message too large: {} bytes (max {})",
1322                        msg_len, MAX_MESSAGE_SIZE
1323                    ));
1324                }
1325
1326                if self.buffer.len() > msg_len {
1327                    let msg_type = self.buffer[0];
1328
1329                    if msg_type == b'D' {
1330                        let msg_bytes = self.buffer.split_to(msg_len + 1).freeze();
1331                        let payload = msg_bytes.slice(5..);
1332                        if let Err(err) = parse_data_row_payload_zerocopy(payload, row) {
1333                            return self.protocol_desync_error(err);
1334                        }
1335                        return Ok(msg_type);
1336                    }
1337
1338                    let msg_bytes = self.buffer.split_to(msg_len + 1);
1339                    if let Some(msg_type) = self.decode_fast_message_type(msg_bytes)? {
1340                        return Ok(msg_type);
1341                    }
1342                    continue;
1343                }
1344            }
1345
1346            let n = self.read_with_timeout().await?;
1347            if n == 0 {
1348                return self.connection_desync("Connection closed".to_string());
1349            }
1350        }
1351    }
1352
1353    /// FAST receive for scalar result consumption into a reusable first-column buffer.
1354    #[inline]
1355    pub(crate) async fn recv_fill_first_column_zerocopy_fast(
1356        &mut self,
1357        first_column: &mut Option<Bytes>,
1358    ) -> PgResult<u8> {
1359        loop {
1360            if self.buffer.len() >= 5 {
1361                let msg_len = u32::from_be_bytes([
1362                    self.buffer[1],
1363                    self.buffer[2],
1364                    self.buffer[3],
1365                    self.buffer[4],
1366                ]) as usize;
1367
1368                if msg_len < 4 {
1369                    return self.protocol_desync(format!(
1370                        "Invalid message length: {} (minimum 4)",
1371                        msg_len
1372                    ));
1373                }
1374
1375                if msg_len > MAX_MESSAGE_SIZE {
1376                    return self.protocol_desync(format!(
1377                        "Message too large: {} bytes (max {})",
1378                        msg_len, MAX_MESSAGE_SIZE
1379                    ));
1380                }
1381
1382                if self.buffer.len() > msg_len {
1383                    let msg_type = self.buffer[0];
1384
1385                    if msg_type == b'D' {
1386                        let msg_bytes = self.buffer.split_to(msg_len + 1).freeze();
1387                        let payload = msg_bytes.slice(5..);
1388                        match parse_first_column_payload_zerocopy(payload) {
1389                            Ok(column) => *first_column = column,
1390                            Err(err) => return self.protocol_desync_error(err),
1391                        }
1392                        return Ok(msg_type);
1393                    }
1394
1395                    let msg_bytes = self.buffer.split_to(msg_len + 1);
1396                    if let Some(msg_type) = self.decode_fast_message_type(msg_bytes)? {
1397                        return Ok(msg_type);
1398                    }
1399                    continue;
1400                }
1401            }
1402
1403            let n = self.read_with_timeout().await?;
1404            if n == 0 {
1405                return self.connection_desync("Connection closed".to_string());
1406            }
1407        }
1408    }
1409
1410    /// FAST receive for fixed 4-column scalar result sets.
1411    #[inline]
1412    pub(crate) async fn recv_fill_first_four_columns_zerocopy_fast(
1413        &mut self,
1414        columns: &mut [Option<Bytes>; 4],
1415    ) -> PgResult<u8> {
1416        loop {
1417            if self.buffer.len() >= 5 {
1418                let msg_len = u32::from_be_bytes([
1419                    self.buffer[1],
1420                    self.buffer[2],
1421                    self.buffer[3],
1422                    self.buffer[4],
1423                ]) as usize;
1424
1425                if msg_len < 4 {
1426                    return self.protocol_desync(format!(
1427                        "Invalid message length: {} (minimum 4)",
1428                        msg_len
1429                    ));
1430                }
1431
1432                if msg_len > MAX_MESSAGE_SIZE {
1433                    return self.protocol_desync(format!(
1434                        "Message too large: {} bytes (max {})",
1435                        msg_len, MAX_MESSAGE_SIZE
1436                    ));
1437                }
1438
1439                if self.buffer.len() > msg_len {
1440                    let msg_type = self.buffer[0];
1441
1442                    if msg_type == b'D' {
1443                        let msg_bytes = self.buffer.split_to(msg_len + 1).freeze();
1444                        let payload = msg_bytes.slice(5..);
1445                        if let Err(err) =
1446                            parse_first_four_columns_payload_zerocopy(payload, columns)
1447                        {
1448                            return self.protocol_desync_error(err);
1449                        }
1450                        return Ok(msg_type);
1451                    }
1452
1453                    let msg_bytes = self.buffer.split_to(msg_len + 1);
1454                    if let Some(msg_type) = self.decode_fast_message_type(msg_bytes)? {
1455                        return Ok(msg_type);
1456                    }
1457                    continue;
1458                }
1459            }
1460
1461            let n = self.read_with_timeout().await?;
1462            if n == 0 {
1463                return self.connection_desync("Connection closed".to_string());
1464            }
1465        }
1466    }
1467
1468    /// ZERO-COPY receive for DataRow.
1469    /// Uses bytes::Bytes for reference-counted slicing instead of Vec copy.
1470    /// Returns: (msg_type, Option<row_data>)
1471    /// For 'D' (DataRow): returns Bytes slices (no copy!)
1472    /// For other types: returns None
1473    #[inline]
1474    pub(crate) async fn recv_data_zerocopy(
1475        &mut self,
1476    ) -> PgResult<(u8, Option<Vec<Option<bytes::Bytes>>>)> {
1477        use bytes::Buf;
1478
1479        loop {
1480            if self.buffer.len() >= 5 {
1481                let msg_len = u32::from_be_bytes([
1482                    self.buffer[1],
1483                    self.buffer[2],
1484                    self.buffer[3],
1485                    self.buffer[4],
1486                ]) as usize;
1487
1488                if msg_len < 4 {
1489                    return self.protocol_desync(format!(
1490                        "Invalid message length: {} (minimum 4)",
1491                        msg_len
1492                    ));
1493                }
1494
1495                if msg_len > MAX_MESSAGE_SIZE {
1496                    return self.protocol_desync(format!(
1497                        "Message too large: {} bytes (max {})",
1498                        msg_len, MAX_MESSAGE_SIZE
1499                    ));
1500                }
1501
1502                if self.buffer.len() > msg_len {
1503                    let msg_type = self.buffer[0];
1504
1505                    // Fast path: DataRow - ZERO-COPY using Bytes
1506                    if msg_type == b'D' {
1507                        // Split off the entire message
1508                        let mut msg_bytes = self.buffer.split_to(msg_len + 1);
1509
1510                        // Skip type byte (1) + length (4) = 5 bytes
1511                        msg_bytes.advance(5);
1512
1513                        if msg_bytes.len() >= 2 {
1514                            let raw_count = msg_bytes.get_i16();
1515                            if raw_count < 0 {
1516                                return self.protocol_desync(format!(
1517                                    "DataRow invalid column count: {}",
1518                                    raw_count
1519                                ));
1520                            }
1521                            let column_count = raw_count as usize;
1522                            if column_count > msg_bytes.remaining() / 4 + 1 {
1523                                return self.protocol_desync(format!(
1524                                    "DataRow claims {} columns but payload is only {} bytes",
1525                                    column_count,
1526                                    msg_bytes.remaining() + 2
1527                                ));
1528                            }
1529                            let mut columns = Vec::with_capacity(column_count);
1530
1531                            for _ in 0..column_count {
1532                                if msg_bytes.remaining() < 4 {
1533                                    return self.protocol_desync(
1534                                        "DataRow truncated: missing column length".into(),
1535                                    );
1536                                }
1537
1538                                let len = msg_bytes.get_i32();
1539
1540                                if len == -1 {
1541                                    columns.push(None);
1542                                } else {
1543                                    if len < -1 {
1544                                        return self.protocol_desync(format!(
1545                                            "DataRow invalid column length: {}",
1546                                            len
1547                                        ));
1548                                    }
1549                                    let len = len as usize;
1550                                    if msg_bytes.remaining() < len {
1551                                        return self.protocol_desync(
1552                                            "DataRow truncated: column data exceeds payload".into(),
1553                                        );
1554                                    }
1555                                    let col_data = msg_bytes.split_to(len).freeze();
1556                                    columns.push(Some(col_data));
1557                                }
1558                            }
1559
1560                            if msg_bytes.remaining() != 0 {
1561                                return self.protocol_desync("DataRow has trailing bytes".into());
1562                            }
1563
1564                            return Ok((msg_type, Some(columns)));
1565                        }
1566                        return self.protocol_desync("DataRow payload too short".into());
1567                    }
1568
1569                    let msg_bytes = self.buffer.split_to(msg_len + 1);
1570                    if let Some(msg_type) = self.decode_fast_message_type(msg_bytes)? {
1571                        return Ok((msg_type, None));
1572                    }
1573                    continue;
1574                }
1575            }
1576
1577            let n = self.read_with_timeout().await?;
1578            if n == 0 {
1579                return self.connection_desync("Connection closed".to_string());
1580            }
1581        }
1582    }
1583
1584    /// ULTRA-FAST receive for 2-column DataRow (id, name pattern).
1585    /// Uses fixed-size array instead of Vec allocation.
1586    /// Returns: (msg_type, Option<(col0, col1)>)
1587    #[inline(always)]
1588    pub(crate) async fn recv_data_ultra(
1589        &mut self,
1590    ) -> PgResult<(u8, Option<(bytes::Bytes, bytes::Bytes)>)> {
1591        use bytes::Buf;
1592
1593        loop {
1594            if self.buffer.len() >= 5 {
1595                let msg_len = u32::from_be_bytes([
1596                    self.buffer[1],
1597                    self.buffer[2],
1598                    self.buffer[3],
1599                    self.buffer[4],
1600                ]) as usize;
1601
1602                if msg_len < 4 {
1603                    return self.protocol_desync(format!(
1604                        "Invalid message length: {} (minimum 4)",
1605                        msg_len
1606                    ));
1607                }
1608
1609                if msg_len > MAX_MESSAGE_SIZE {
1610                    return self.protocol_desync(format!(
1611                        "Message too large: {} bytes (max {})",
1612                        msg_len, MAX_MESSAGE_SIZE
1613                    ));
1614                }
1615
1616                if self.buffer.len() > msg_len {
1617                    let msg_type = self.buffer[0];
1618
1619                    if msg_type == b'D' {
1620                        let mut msg_bytes = self.buffer.split_to(msg_len + 1);
1621                        msg_bytes.advance(5); // Skip type + length
1622
1623                        // Bounds checks to prevent panic on truncated DataRow
1624                        if msg_bytes.remaining() < 2 {
1625                            return self.protocol_desync(
1626                                "DataRow ultra: too short for column count".into(),
1627                            );
1628                        }
1629
1630                        // Read column count (expect 2)
1631                        let col_count = msg_bytes.get_i16();
1632                        if col_count != 2 {
1633                            return self.protocol_desync(format!(
1634                                "DataRow ultra expects exactly 2 columns, got {}",
1635                                col_count
1636                            ));
1637                        }
1638
1639                        if msg_bytes.remaining() < 4 {
1640                            return self.protocol_desync(
1641                                "DataRow ultra: truncated before col0 length".into(),
1642                            );
1643                        }
1644                        let len0 = msg_bytes.get_i32();
1645                        let col0 = if len0 > 0 {
1646                            let len0 = len0 as usize;
1647                            if msg_bytes.remaining() < len0 {
1648                                return self.protocol_desync(
1649                                    "DataRow ultra: col0 data exceeds payload".into(),
1650                                );
1651                            }
1652                            msg_bytes.split_to(len0).freeze()
1653                        } else if len0 == 0 {
1654                            bytes::Bytes::new()
1655                        } else if len0 == -1 {
1656                            return self.protocol_desync(
1657                                "DataRow ultra does not support NULL columns".into(),
1658                            );
1659                        } else {
1660                            return self.protocol_desync(format!(
1661                                "DataRow ultra: invalid col0 length {}",
1662                                len0
1663                            ));
1664                        };
1665
1666                        if msg_bytes.remaining() < 4 {
1667                            return self.protocol_desync(
1668                                "DataRow ultra: truncated before col1 length".into(),
1669                            );
1670                        }
1671                        let len1 = msg_bytes.get_i32();
1672                        let col1 = if len1 > 0 {
1673                            let len1 = len1 as usize;
1674                            if msg_bytes.remaining() < len1 {
1675                                return self.protocol_desync(
1676                                    "DataRow ultra: col1 data exceeds payload".into(),
1677                                );
1678                            }
1679                            msg_bytes.split_to(len1).freeze()
1680                        } else if len1 == 0 {
1681                            bytes::Bytes::new()
1682                        } else if len1 == -1 {
1683                            return self.protocol_desync(
1684                                "DataRow ultra does not support NULL columns".into(),
1685                            );
1686                        } else {
1687                            return self.protocol_desync(format!(
1688                                "DataRow ultra: invalid col1 length {}",
1689                                len1
1690                            ));
1691                        };
1692
1693                        if msg_bytes.remaining() != 0 {
1694                            return self.protocol_desync(
1695                                "DataRow ultra: trailing bytes after expected columns".into(),
1696                            );
1697                        }
1698
1699                        return Ok((msg_type, Some((col0, col1))));
1700                    }
1701
1702                    let msg_bytes = self.buffer.split_to(msg_len + 1);
1703                    if let Some(msg_type) = self.decode_fast_message_type(msg_bytes)? {
1704                        return Ok((msg_type, None));
1705                    }
1706                    continue;
1707                }
1708            }
1709
1710            let n = self.read_with_timeout().await?;
1711            if n == 0 {
1712                return self.connection_desync("Connection closed".to_string());
1713            }
1714        }
1715    }
1716}
1717
1718#[cfg(test)]
1719mod tests {
1720    use super::*;
1721
1722    #[cfg(unix)]
1723    fn test_conn() -> PgConnection {
1724        use crate::driver::connection::StatementCache;
1725        use crate::driver::stream::PgStream;
1726        use std::collections::{HashMap, VecDeque};
1727        use std::num::NonZeroUsize;
1728        use tokio::net::UnixStream;
1729
1730        let (unix_stream, _peer) = UnixStream::pair().expect("unix stream pair");
1731        PgConnection {
1732            stream: PgStream::Unix(unix_stream),
1733            buffer: BytesMut::with_capacity(1024),
1734            write_buf: BytesMut::with_capacity(1024),
1735            sql_buf: BytesMut::with_capacity(256),
1736            params_buf: Vec::new(),
1737            prepared_statements: HashMap::new(),
1738            stmt_cache: StatementCache::new(NonZeroUsize::new(2).expect("non-zero")),
1739            column_info_cache: HashMap::new(),
1740            process_id: 0,
1741            cancel_key_bytes: Vec::new(),
1742            requested_protocol_minor: PgConnection::default_protocol_minor(),
1743            negotiated_protocol_minor: PgConnection::default_protocol_minor(),
1744            notifications: VecDeque::new(),
1745            replication_stream_active: false,
1746            replication_mode_enabled: false,
1747            last_replication_wal_end: None,
1748            io_desynced: false,
1749            pending_statement_closes: Vec::new(),
1750            draining_statement_closes: false,
1751        }
1752    }
1753
1754    fn build_data_row_payload(columns: &[Option<&[u8]>]) -> Bytes {
1755        let mut payload = Vec::new();
1756        payload.extend_from_slice(&(columns.len() as i16).to_be_bytes());
1757        for column in columns {
1758            match column {
1759                Some(bytes) => {
1760                    payload.extend_from_slice(&(bytes.len() as i32).to_be_bytes());
1761                    payload.extend_from_slice(bytes);
1762                }
1763                None => payload.extend_from_slice(&(-1i32).to_be_bytes()),
1764            }
1765        }
1766        Bytes::from(payload)
1767    }
1768
1769    fn push_data_row_frame(conn: &mut PgConnection, payload: &[u8]) {
1770        let msg_len = payload.len() + 4;
1771        conn.buffer.extend_from_slice(b"D");
1772        conn.buffer
1773            .extend_from_slice(&(msg_len as u32).to_be_bytes());
1774        conn.buffer.extend_from_slice(payload);
1775    }
1776
1777    fn push_one_column_datarow_without_column_length(conn: &mut PgConnection) {
1778        push_data_row_frame(conn, &[0, 1]);
1779    }
1780
1781    fn assert_protocol_error_contains(err: PgError, expected: &str) {
1782        match err {
1783            PgError::Protocol(msg) => assert!(
1784                msg.contains(expected),
1785                "expected protocol error containing {expected:?}, got {msg:?}"
1786            ),
1787            err => panic!("expected protocol error containing {expected:?}, got {err:?}"),
1788        }
1789    }
1790
1791    #[test]
1792    fn parse_first_four_columns_payload_zerocopy_reads_values() {
1793        let payload = build_data_row_payload(&[Some(b"10"), None, Some(b"30"), Some(b"")]);
1794        let mut columns = [None, None, None, None];
1795
1796        parse_first_four_columns_payload_zerocopy(payload, &mut columns).unwrap();
1797
1798        assert_eq!(columns[0].as_deref(), Some(&b"10"[..]));
1799        assert_eq!(columns[1].as_deref(), None);
1800        assert_eq!(columns[2].as_deref(), Some(&b"30"[..]));
1801        assert_eq!(columns[3].as_deref(), Some(&b""[..]));
1802    }
1803
1804    #[test]
1805    fn parse_first_four_columns_payload_zerocopy_rejects_wrong_arity() {
1806        let payload = build_data_row_payload(&[Some(b"1"), Some(b"2"), Some(b"3")]);
1807        let mut columns = [None, None, None, None];
1808
1809        let err = parse_first_four_columns_payload_zerocopy(payload, &mut columns).unwrap_err();
1810
1811        assert!(
1812            err.to_string()
1813                .contains("fast-path expects exactly 4 columns")
1814        );
1815    }
1816
1817    #[cfg(unix)]
1818    #[tokio::test]
1819    async fn recv_data_zerocopy_rejects_datarow_length_4() {
1820        let mut conn = test_conn();
1821        conn.buffer.extend_from_slice(&[b'D', 0, 0, 0, 4]);
1822
1823        let err = conn.recv_data_zerocopy().await.unwrap_err();
1824
1825        assert!(err.to_string().contains("DataRow payload too short"));
1826        assert!(conn.is_io_desynced());
1827    }
1828
1829    #[cfg(unix)]
1830    #[tokio::test]
1831    async fn recv_data_zerocopy_rejects_datarow_length_5() {
1832        let mut conn = test_conn();
1833        conn.buffer.extend_from_slice(&[b'D', 0, 0, 0, 5, 0]);
1834
1835        let err = conn.recv_data_zerocopy().await.unwrap_err();
1836
1837        assert!(err.to_string().contains("DataRow payload too short"));
1838        assert!(conn.is_io_desynced());
1839    }
1840
1841    #[cfg(unix)]
1842    #[tokio::test]
1843    async fn recv_with_data_fast_desyncs_on_malformed_datarow() {
1844        let mut conn = test_conn();
1845        push_one_column_datarow_without_column_length(&mut conn);
1846
1847        let err = conn.recv_with_data_fast().await.unwrap_err();
1848
1849        assert_protocol_error_contains(err, "DataRow truncated");
1850        assert!(conn.is_io_desynced());
1851    }
1852
1853    #[cfg(unix)]
1854    #[tokio::test]
1855    async fn recv_fill_data_row_fast_desyncs_on_malformed_datarow() {
1856        let mut conn = test_conn();
1857        let mut row = Vec::new();
1858        push_one_column_datarow_without_column_length(&mut conn);
1859
1860        let err = conn.recv_fill_data_row_fast(&mut row).await.unwrap_err();
1861
1862        assert_protocol_error_contains(err, "DataRow truncated");
1863        assert!(conn.is_io_desynced());
1864    }
1865
1866    #[cfg(unix)]
1867    #[tokio::test]
1868    async fn recv_fill_zerocopy_row_fast_desyncs_on_malformed_datarow() {
1869        let mut conn = test_conn();
1870        let mut row = PgBytesRow::default();
1871        push_one_column_datarow_without_column_length(&mut conn);
1872
1873        let err = conn
1874            .recv_fill_zerocopy_row_fast(&mut row)
1875            .await
1876            .unwrap_err();
1877
1878        assert_protocol_error_contains(err, "DataRow truncated");
1879        assert!(conn.is_io_desynced());
1880    }
1881
1882    #[cfg(unix)]
1883    #[tokio::test]
1884    async fn recv_fill_first_column_zerocopy_fast_desyncs_on_malformed_datarow() {
1885        let mut conn = test_conn();
1886        let mut first_column = None;
1887        push_one_column_datarow_without_column_length(&mut conn);
1888
1889        let err = conn
1890            .recv_fill_first_column_zerocopy_fast(&mut first_column)
1891            .await
1892            .unwrap_err();
1893
1894        assert_protocol_error_contains(err, "DataRow truncated");
1895        assert!(conn.is_io_desynced());
1896    }
1897
1898    #[cfg(unix)]
1899    #[tokio::test]
1900    async fn recv_fill_first_four_columns_zerocopy_fast_desyncs_on_malformed_datarow() {
1901        let mut conn = test_conn();
1902        let mut columns = [None, None, None, None];
1903        push_one_column_datarow_without_column_length(&mut conn);
1904
1905        let err = conn
1906            .recv_fill_first_four_columns_zerocopy_fast(&mut columns)
1907            .await
1908            .unwrap_err();
1909
1910        assert_protocol_error_contains(err, "DataRow fast-path expects exactly 4 columns");
1911        assert!(conn.is_io_desynced());
1912    }
1913
1914    #[cfg(unix)]
1915    #[tokio::test]
1916    async fn recv_data_ultra_desyncs_on_malformed_datarow() {
1917        let mut conn = test_conn();
1918        push_one_column_datarow_without_column_length(&mut conn);
1919
1920        let err = conn.recv_data_ultra().await.unwrap_err();
1921
1922        assert_protocol_error_contains(err, "DataRow ultra expects exactly 2 columns");
1923        assert!(conn.is_io_desynced());
1924    }
1925
1926    #[cfg(unix)]
1927    #[tokio::test]
1928    async fn recv_msg_type_fast_rejects_malformed_ready_for_query() {
1929        let mut conn = test_conn();
1930        conn.buffer.extend_from_slice(&[b'Z', 0, 0, 0, 5, b'X']);
1931
1932        let err = conn.recv_msg_type_fast().await.unwrap_err();
1933
1934        assert!(err.to_string().contains("Unknown transaction status"));
1935        assert!(conn.is_io_desynced());
1936    }
1937
1938    #[cfg(unix)]
1939    #[tokio::test]
1940    async fn recv_msg_type_fast_rejects_malformed_command_complete() {
1941        let mut conn = test_conn();
1942        conn.buffer.extend_from_slice(&[
1943            b'C', 0, 0, 0, 12, b'S', b'E', b'L', b'E', b'C', b'T', b' ', b'1',
1944        ]);
1945
1946        let err = conn.recv_msg_type_fast().await.unwrap_err();
1947
1948        assert!(
1949            err.to_string()
1950                .contains("CommandComplete missing null terminator")
1951        );
1952        assert!(conn.is_io_desynced());
1953    }
1954
1955    #[cfg(unix)]
1956    fn push_notification_frame(conn: &mut PgConnection, channel: &str, payload: &str) {
1957        let mut body = Vec::new();
1958        body.extend_from_slice(&7i32.to_be_bytes());
1959        body.extend_from_slice(channel.as_bytes());
1960        body.push(0);
1961        body.extend_from_slice(payload.as_bytes());
1962        body.push(0);
1963        conn.buffer.extend_from_slice(b"A");
1964        conn.buffer
1965            .extend_from_slice(&((body.len() + 4) as u32).to_be_bytes());
1966        conn.buffer.extend_from_slice(&body);
1967    }
1968
1969    #[cfg(unix)]
1970    #[tokio::test]
1971    async fn recv_buffers_notifications_below_cap() {
1972        let mut conn = test_conn();
1973        for _ in 0..3 {
1974            push_notification_frame(&mut conn, "jobs", "42");
1975        }
1976        conn.buffer.extend_from_slice(&[b'Z', 0, 0, 0, 5, b'I']);
1977
1978        let msg = conn.recv().await.unwrap();
1979
1980        assert!(matches!(msg, BackendMessage::ReadyForQuery(_)));
1981        assert_eq!(conn.notifications.len(), 3);
1982        assert!(!conn.is_io_desynced());
1983    }
1984
1985    #[cfg(unix)]
1986    #[tokio::test]
1987    async fn recv_desyncs_on_notification_flood() {
1988        let mut conn = test_conn();
1989        for _ in 0..=MAX_BUFFERED_NOTIFICATIONS {
1990            push_notification_frame(&mut conn, "c", "");
1991        }
1992
1993        let err = conn.recv().await.unwrap_err();
1994
1995        assert!(err.to_string().contains("notification buffer overflow"));
1996        assert!(conn.is_io_desynced());
1997        assert_eq!(conn.notifications.len(), MAX_BUFFERED_NOTIFICATIONS);
1998    }
1999
2000    #[cfg(unix)]
2001    #[tokio::test]
2002    async fn recv_msg_type_fast_desyncs_on_notification_flood() {
2003        let mut conn = test_conn();
2004        for _ in 0..=MAX_BUFFERED_NOTIFICATIONS {
2005            push_notification_frame(&mut conn, "c", "");
2006        }
2007
2008        let err = conn.recv_msg_type_fast().await.unwrap_err();
2009
2010        assert!(err.to_string().contains("notification buffer overflow"));
2011        assert!(conn.is_io_desynced());
2012        assert_eq!(conn.notifications.len(), MAX_BUFFERED_NOTIFICATIONS);
2013    }
2014
2015    #[cfg(unix)]
2016    #[tokio::test]
2017    async fn recv_without_timeout_desyncs_on_notification_flood() {
2018        let mut conn = test_conn();
2019        for _ in 0..=MAX_BUFFERED_NOTIFICATIONS {
2020            push_notification_frame(&mut conn, "c", "");
2021        }
2022
2023        let err = conn.recv_without_timeout().await.unwrap_err();
2024
2025        assert!(err.to_string().contains("notification buffer overflow"));
2026        assert!(conn.is_io_desynced());
2027        assert_eq!(conn.notifications.len(), MAX_BUFFERED_NOTIFICATIONS);
2028    }
2029
2030    #[cfg(unix)]
2031    #[tokio::test]
2032    async fn recv_desyncs_on_oversized_notification() {
2033        let mut conn = test_conn();
2034        let payload = "x".repeat(MAX_NOTIFICATION_BYTES + 1);
2035        push_notification_frame(&mut conn, "c", &payload);
2036
2037        let err = conn.recv().await.unwrap_err();
2038
2039        assert!(err.to_string().contains("notification exceeds"));
2040        assert!(conn.is_io_desynced());
2041        assert!(conn.notifications.is_empty());
2042    }
2043}