From 902114b7565f07e3ee9f91d966befd1ebb3cd475 Mon Sep 17 00:00:00 2001 From: rich-murphey Date: Thu, 11 Feb 2021 22:12:35 -0600 Subject: [PATCH] postgres: peek at header then take whole message --- sqlx-postgres/src/connection/stream.rs | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/sqlx-postgres/src/connection/stream.rs b/sqlx-postgres/src/connection/stream.rs index 7559cb4f..59d8f878 100644 --- a/sqlx-postgres/src/connection/stream.rs +++ b/sqlx-postgres/src/connection/stream.rs @@ -53,13 +53,14 @@ macro_rules! read_packet { loop { read_packet!($(@$blocking)? @stream $self, 0, 5); - let mut header: Bytes = $self.stream.take(5); + // peek at the messaage type and payload size + let r#type = MessageType::try_from(*$self.stream.get(0,1))?; + let size = (u32::from_le_bytes($self.stream.get(1,4)) - 4) as usize; - let r#type = MessageType::try_from(header.get_u8())?; - let size = (header.get_u32() - 4) as usize; - - read_packet!($(@$blocking)? @stream $self, 4, size); + read_packet!($(@$blocking)? @stream $self, 5, size); + // take the whole packet + let header = $self.stream.take(5); let contents = $self.stream.take(size); let message = Message { r#type, contents };