From eb3cd5f4eea4ab3ca38e46f438e9aa69185ae398 Mon Sep 17 00:00:00 2001
From: Daniel Akhterov <akhterovd@gmail.com>
Date: Sun, 23 Jun 2019 20:20:06 -0700
Subject: [PATCH] Serialization length

---
 mason-mariadb/src/protocol/client.rs    |  9 ++--
 mason-mariadb/src/protocol/serialize.rs | 59 +++++++++++++++++++++----
 2 files changed, 54 insertions(+), 14 deletions(-)

diff --git a/mason-mariadb/src/protocol/client.rs b/mason-mariadb/src/protocol/client.rs
index 0ab1a57b..40e17241 100644
--- a/mason-mariadb/src/protocol/client.rs
+++ b/mason-mariadb/src/protocol/client.rs
@@ -8,9 +8,9 @@
 // TODO: Handle when capability is set, but field is None
 
 use super::server::Capabilities;
-use byteorder::ByteOrder;
-use byteorder::LittleEndian;
+use byteorder::{ByteOrder, LittleEndian, WriteBytesExt};
 use bytes::Bytes;
+use crate::protocol::serialize::*;
 
 pub trait Serialize {
     fn serialize(&self, buf: &mut Vec<u8>);
@@ -54,9 +54,7 @@ pub struct AuthenticationSwitchRequestPacket {
 impl Serialize for SSLRequestPacket {
     fn serialize(&self, buf: &mut Vec<u8>) {
         // Temporary storage for length: 3 bytes
-        buf.push(0);
-        buf.push(0);
-        buf.push(0);
+        buf.write_u24::<LittleEndian>(0);
 
         // Sequence Numer
         buf.push(self.sequence_number);
@@ -82,6 +80,7 @@ impl Serialize for SSLRequestPacket {
         buf[0] = buf.len().to_le_bytes()[0];
         buf[1] = buf.len().to_le_bytes()[1];
         buf[2] = buf.len().to_le_bytes()[2];
+        serialize_length(buf);
     }
 }
 
diff --git a/mason-mariadb/src/protocol/serialize.rs b/mason-mariadb/src/protocol/serialize.rs
index d95d95b2..15e04d08 100644
--- a/mason-mariadb/src/protocol/serialize.rs
+++ b/mason-mariadb/src/protocol/serialize.rs
@@ -5,6 +5,24 @@ use failure::err_msg;
 
 const U24_MAX: usize = 0xFF_FF_FF;
 
+#[inline]
+pub fn serialize_length(buf: &mut Vec<u8>) {
+    let mut length =  [0;  3];
+    if buf.len() > U24_MAX {
+        panic!("Buffer too long");
+    } else if buf.len() <= 4 {
+        panic!("Buffer too short. Only contains packet length and sequence number")
+    }
+
+    LittleEndian::write_u24(&mut length, buf.len() as u32 - 4);
+
+    // Set length at the start of the buffer
+    // sadly there is no `prepend` for rust Vec
+    buf[0] = length[0];
+    buf[1] = length[1];
+    buf[2] = length[2];
+}
+
 #[inline]
 pub fn serialize_int_8(buf: &mut Vec<u8>, value: u64) {
     buf.write_u64::<LittleEndian>(value).unwrap();
@@ -112,16 +130,39 @@ pub fn serialize_byte_eof(buf: &mut Vec<u8>, bytes: &Bytes) {
 #[cfg(test)]
 mod tests {
     use super::*;
+    use byteorder::{ByteOrder, LittleEndian, WriteBytesExt};
 
-    // [ ] serialize_int_lenenc
-    // [ ] serialize_int_lenenc
-    // [ ] serialize_int_lenenc
-    // [ ] serialize_int_lenenc
-    // [ ] serialize_int_lenenc
-    // [ ] serialize_int_lenenc
-    // [ ] serialize_int_lenenc
-    // [ ] serialize_int_lenenc
-    // [ ] serialize_int_lenenc
+    // [X] serialize_int_lenenc_u64
+    // [X] serialize_int_lenenc_u32
+    // [X] serialize_int_lenenc_u24
+    // [X] serialize_int_lenenc_u16
+    // [X] serialize_int_lenenc_u8
+    // [X] serialize_int_u64
+    // [X] serialize_int_u32
+    // [X] serialize_int_u24
+    // [X] serialize_int_u16
+    // [X] serialize_int_u8
+    // [X] serialize_string_lenenc
+    // [X] serialize_string_fix
+    // [X] serialize_string_null
+    // [X] serialize_string_eof
+    // [X] serialize_byte_lenenc
+    // [X] serialize_byte_fix
+    // [X] serialize_byte_eof
+
+    #[test]
+    fn it_encodes_length() {
+        let mut buf: Vec<u8> = Vec::new();
+        // Reserve space of length
+        buf.write_u24::<LittleEndian>(0);
+        // Sequence number; typically 0
+        buf.write_u8(0x00);
+        // Contents of buffer
+        buf.write_u8(0xFF);
+        serialize_length(&mut buf);
+
+        assert_eq!(buf, b"\x01\0\0\0\xFF".to_vec());
+    }
 
     #[test]
     fn it_encodes_int_lenenc_none() {