Read message

Author: sfackler

src/lib.rs

@@ -50,7 +50,6 @@ extern crate log;
 extern crate phf;
 extern crate postgres_protocol;
 
-use bufstream::BufStream;
 use md5::Md5;
 use std::cell::{Cell, RefCell};
 use std::collections::{VecDeque, HashMap};
@@ -64,7 +63,7 @@ use std::time::Duration;
 use postgres_protocol::message::frontend;
 
 use error::{Error, ConnectError, SqlState, DbError};
-use io::{TlsStream, TlsHandshake};
+use io::TlsHandshake;
 use message::{Backend, RowDescriptionEntry, ReadMessage};
 use notification::{Notifications, Notification};
 use params::{ConnectParams, IntoConnectParams, UserInfo};
@@ -304,7 +303,7 @@ impl InnerConnection {
     fn read_message_with_notification(&mut self) -> std_io::Result<Backend> {
         debug_assert!(!self.desynchronized);
         loop {
-            match try_desync!(self, self.stream.read_message()) {
+            match try_desync!(self, ReadMessage::read_message(&mut self.stream)) {
                 Backend::NoticeResponse { fields } => {
                     if let Ok(err) = DbError::new_raw(fields) {
                         self.notice_handler.handle_notice(err);

src/priv_io.rs

@@ -12,13 +12,15 @@ use std::os::unix::io::{AsRawFd, RawFd};
 #[cfg(windows)]
 use std::os::windows::io::{AsRawSocket, RawSocket};
 use postgres_protocol::message::frontend;
+use postgres_protocol::message::backend::{self, ParseResult};
 
 use TlsMode;
 use params::{ConnectParams, ConnectTarget};
 use error::ConnectError;
 use io::TlsStream;
 
 const DEFAULT_PORT: u16 = 5432;
+const MESSAGE_HEADER_SIZE: usize = 5;
 
 pub struct MessageStream {
     stream: BufStream<Box<TlsStream>>,
@@ -43,6 +45,27 @@ impl MessageStream {
         self.stream.write_all(&self.buf)
     }
 
+    pub fn read_message<'a>(&'a mut self) -> io::Result<backend::Message<'a>> {
+        self.buf.resize(MESSAGE_HEADER_SIZE, 0);
+        try!(self.stream.read_exact(&mut self.buf));
+
+        let len = match try!(backend::Message::parse(&self.buf)) {
+            // FIXME this is dumb but an explicit return runs into borrowck issues :(
+            ParseResult::Complete { .. } => None,
+            ParseResult::Incomplete { required_size } => Some(required_size.unwrap()),
+        };
+
+        if let Some(len) = len {
+            self.buf.resize(len, 0);
+            try!(self.stream.read_exact(&mut self.buf[MESSAGE_HEADER_SIZE..]));
+        };
+
+        match try!(backend::Message::parse(&self.buf)) {
+            ParseResult::Complete { message, .. } => Ok(message),
+            ParseResult::Incomplete { .. } => unreachable!(),
+        }
+    }
+
     pub fn flush(&mut self) -> io::Result<()> {
         self.stream.flush()
     }

src/rows.rs

@@ -3,7 +3,6 @@
 use std::ascii::AsciiExt;
 use std::borrow::Cow;
 use std::collections::VecDeque;
-use std::io::Write;
 use std::fmt;
 use std::ops::Deref;
 use std::slice;