Merge pull request #137 from wllenyj/message_header

Improve MessageHeader parsing.

Author: Tim-Zhang

src/asynchronous/stream.rs

@@ -3,8 +3,6 @@
 // SPDX-License-Identifier: Apache-2.0
 //
 
-use byteorder::{BigEndian, ByteOrder};
-
 use crate::common::{MESSAGE_HEADER_LENGTH, MESSAGE_LENGTH_MAX, MESSAGE_TYPE_RESPONSE};
 use crate::error::{get_rpc_status, sock_error_msg, Error, Result};
 use crate::proto::{Code, Response, Status};
@@ -38,15 +36,7 @@ where
         ));
     }
 
-    let mut mh = MessageHeader::default();
-    let mut covbuf: &[u8] = &buf[..4];
-    mh.length = byteorder::ReadBytesExt::read_u32::<BigEndian>(&mut covbuf)
-        .map_err(err_to_rpc_err!(Code::INVALID_ARGUMENT, e, ""))?;
-    let mut covbuf: &[u8] = &buf[4..8];
-    mh.stream_id = byteorder::ReadBytesExt::read_u32::<BigEndian>(&mut covbuf)
-        .map_err(err_to_rpc_err!(Code::INVALID_ARGUMENT, e, ""))?;
-    mh.type_ = buf[8];
-    mh.flags = buf[9];
+    let mh = MessageHeader::from(&buf);
 
     Ok(mh)
 }
@@ -82,16 +72,7 @@ where
 }
 
 fn header_to_buf(mh: MessageHeader) -> Vec<u8> {
-    let mut buf = vec![0u8; MESSAGE_HEADER_LENGTH];
-
-    let covbuf: &mut [u8] = &mut buf[..4];
-    BigEndian::write_u32(covbuf, mh.length);
-    let covbuf: &mut [u8] = &mut buf[4..8];
-    BigEndian::write_u32(covbuf, mh.stream_id);
-    buf[8] = mh.type_;
-    buf[9] = mh.flags;
-
-    buf
+    mh.into()
 }
 
 pub(crate) fn to_req_buf(stream_id: u32, mut body: Vec<u8>) -> Vec<u8> {

src/common.rs

@@ -6,6 +6,7 @@
 //! Common functions and macros.
 
 use crate::error::{Error, Result};
+use byteorder::{BigEndian, ByteOrder};
 #[cfg(any(feature = "async", not(target_os = "linux")))]
 use nix::fcntl::FdFlag;
 use nix::fcntl::{fcntl, FcntlArg, OFlag};
@@ -28,6 +29,44 @@ pub struct MessageHeader {
     pub flags: u8,
 }
 
+impl<T> From<T> for MessageHeader
+where
+    T: AsRef<[u8]>,
+{
+    fn from(buf: T) -> Self {
+        let buf = buf.as_ref();
+        debug_assert!(buf.len() >= MESSAGE_HEADER_LENGTH);
+        Self {
+            length: BigEndian::read_u32(&buf[..4]),
+            stream_id: BigEndian::read_u32(&buf[4..8]),
+            type_: buf[8],
+            flags: buf[9],
+        }
+    }
+}
+
+impl From<MessageHeader> for Vec<u8> {
+    fn from(mh: MessageHeader) -> Self {
+        let mut buf = vec![0u8; MESSAGE_HEADER_LENGTH];
+        mh.into_buf(&mut buf);
+        buf
+    }
+}
+
+impl MessageHeader {
+    pub(crate) fn into_buf(self, mut buf: impl AsMut<[u8]>) {
+        let buf = buf.as_mut();
+        debug_assert!(buf.len() >= MESSAGE_HEADER_LENGTH);
+
+        let covbuf: &mut [u8] = &mut buf[..4];
+        BigEndian::write_u32(covbuf, self.length);
+        let covbuf: &mut [u8] = &mut buf[4..8];
+        BigEndian::write_u32(covbuf, self.stream_id);
+        buf[8] = self.type_;
+        buf[9] = self.flags;
+    }
+}
+
 pub(crate) fn do_listen(listener: RawFd) -> Result<()> {
     if let Err(e) = fcntl(listener, FcntlArg::F_SETFL(OFlag::O_NONBLOCK)) {
         return Err(Error::Others(format!(
@@ -207,8 +246,7 @@ pub const MESSAGE_TYPE_RESPONSE: u8 = 0x2;
 
 #[cfg(test)]
 mod tests {
-    use super::parse_sockaddr;
-    use super::Domain;
+    use super::*;
 
     #[cfg(target_os = "linux")]
     #[test]
@@ -268,4 +306,23 @@ mod tests {
             }
         }
     }
+
+    #[test]
+    fn message_header() {
+        let buf = vec![
+            0x10, 0x0, 0x0, 0x0, // length
+            0x0, 0x0, 0x0, 0x03, // stream_id
+            0x2,  // type_
+            0xef, // flags
+        ];
+        let mh = MessageHeader::from(&buf);
+        assert_eq!(mh.length, 0x1000_0000);
+        assert_eq!(mh.stream_id, 0x3);
+        assert_eq!(mh.type_, MESSAGE_TYPE_RESPONSE);
+        assert_eq!(mh.flags, 0xef);
+
+        let mut buf2 = vec![0; MESSAGE_HEADER_LENGTH];
+        mh.into_buf(&mut buf2);
+        assert_eq!(&buf, &buf2);
+    }
 }

src/error.rs

@@ -59,12 +59,6 @@ pub fn sock_error_msg(size: usize, msg: String) -> Error {
     get_rpc_status(Code::INVALID_ARGUMENT, msg)
 }
 
-macro_rules! err_to_rpc_err {
-    ($c: expr, $e: ident, $s: expr) => {
-        |$e| get_rpc_status($c, $s.to_string() + &$e.to_string())
-    };
-}
-
 macro_rules! err_to_others_err {
     ($e: ident, $s: expr) => {
         |$e| Error::Others($s.to_string() + &$e.to_string())

src/sync/channel.rs

@@ -12,7 +12,6 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use byteorder::{BigEndian, ByteOrder, ReadBytesExt};
 use nix::sys::socket::*;
 use std::os::unix::io::RawFd;
 
@@ -96,19 +95,7 @@ fn read_message_header(fd: RawFd) -> Result<MessageHeader> {
         ));
     }
 
-    let mut mh = MessageHeader::default();
-    let mut covbuf: &[u8] = &buf[..4];
-    mh.length =
-        covbuf
-            .read_u32::<BigEndian>()
-            .map_err(err_to_rpc_err!(Code::INVALID_ARGUMENT, e, ""))?;
-    let mut covbuf: &[u8] = &buf[4..8];
-    mh.stream_id =
-        covbuf
-            .read_u32::<BigEndian>()
-            .map_err(err_to_rpc_err!(Code::INVALID_ARGUMENT, e, ""))?;
-    mh.type_ = buf[8];
-    mh.flags = buf[9];
+    let mh = MessageHeader::from(&buf);
 
     Ok(mh)
 }
@@ -141,14 +128,7 @@ pub fn read_message(fd: RawFd) -> Result<(MessageHeader, Vec<u8>)> {
 }
 
 fn write_message_header(fd: RawFd, mh: MessageHeader) -> Result<()> {
-    let mut buf = [0u8; MESSAGE_HEADER_LENGTH];
-
-    let covbuf: &mut [u8] = &mut buf[..4];
-    BigEndian::write_u32(covbuf, mh.length);
-    let covbuf: &mut [u8] = &mut buf[4..8];
-    BigEndian::write_u32(covbuf, mh.stream_id);
-    buf[8] = mh.type_;
-    buf[9] = mh.flags;
+    let buf: Vec<u8> = mh.into();
 
     let size = write_count(fd, &buf, MESSAGE_HEADER_LENGTH)?;
     if size != MESSAGE_HEADER_LENGTH {