From b6af918ac1317aa4db4530bdd8d1b51920483cc8 Mon Sep 17 00:00:00 2001
From: Daniel Akhterov <daniel@launchbadge.com>
Date: Fri, 14 Jun 2019 15:40:50 -0700
Subject: [PATCH] Add HandshakeResponsePacket struct

---
 mason-mariadb/src/protocol/client.rs | 117 ++++++++++++++++++++++++++-
 1 file changed, 114 insertions(+), 3 deletions(-)

diff --git a/mason-mariadb/src/protocol/client.rs b/mason-mariadb/src/protocol/client.rs
index ebbc8893..230d2dc5 100644
--- a/mason-mariadb/src/protocol/client.rs
+++ b/mason-mariadb/src/protocol/client.rs
@@ -1,8 +1,16 @@
 // Reference: https://mariadb.com/kb/en/library/connection
+// Packets: https://mariadb.com/kb/en/library/0-packet
+
+// TODO: Handle lengths which are greater than 3 bytes
+// Either break the backet into several smaller ones, or 
+// return error
+// TODO: Handle different Capabilities for server and client
+// TODO: Handle when capability is set, but field is None
 
 use super::server::Capabilities;
 use byteorder::ByteOrder;
 use byteorder::LittleEndian;
+use bytes::Bytes;
 
 pub trait Serialize {
     fn serialize(&self, buf: &mut Vec<u8>);
@@ -17,17 +25,32 @@ pub struct SSLRequestPacket {
     pub extended_capabilities: Option<Capabilities>,
 }
 
+#[derive(Default, Debug)]
+pub struct HandshakeResponsePacket {
+    pub sequence_number: u8,
+    pub capabilities: Capabilities,
+    pub max_packet_size: u32,
+    pub collation: u8,
+    pub extended_capabilities: Option<Capabilities>,
+    pub username: Bytes,
+    pub auth_data: Option<Bytes>,
+    pub auth_response_len: Option<u8>,
+    pub auth_response: Option<Bytes>,
+    pub database: Option<Bytes>,
+    pub auth_plugin_name: Option<Bytes>,
+    pub conn_attr_len: Option<usize>,
+    pub conn_attr: Option<Vec<(Bytes, Bytes)>>,
+}
+
 impl Serialize for SSLRequestPacket {
     fn serialize(&self, buf: &mut Vec<u8>) {
-        // https://mariadb.com/kb/en/library/0-packet
-
         // Temporary storage for length: 3 bytes
         buf.push(0);
         buf.push(0);
         buf.push(0);
 
         // Sequence Numer
-        buf.push(0);
+        buf.push(self.sequence_number);
 
         LittleEndian::write_u32(buf, self.capabilities.bits() as u32);
 
@@ -52,3 +75,91 @@ impl Serialize for SSLRequestPacket {
         buf[2] = buf.len().to_le_bytes()[2];
     }
 }
+
+impl Serialize for HandshakeResponsePacket {
+    fn serialize(&self, buf: &mut Vec<u8>) {
+        // Temporary storage for length: 3 bytes
+        buf.push(0);
+        buf.push(0);
+        buf.push(0);
+
+        // Sequence Numer
+        buf.push(self.sequence_number);
+
+        LittleEndian::write_u32(buf, self.capabilities.bits() as u32);
+
+        LittleEndian::write_u32(buf, self.max_packet_size);
+
+        buf.push(self.collation);
+
+        buf.extend_from_slice(&[0u8;19]);
+
+        if !(self.capabilities & Capabilities::CLIENT_MYSQL).is_empty() {
+            if let Some(capabilities) = self.extended_capabilities {
+                LittleEndian::write_u32(buf, capabilities.bits() as u32);
+            }
+        } else {
+            buf.extend_from_slice(&[0u8;4]);
+        }
+
+        // Username: string<NUL>
+        buf.extend_from_slice(&self.username);
+        buf.push(0);
+
+        if !(self.capabilities & Capabilities::PLUGIN_AUTH_LENENC_CLIENT_DATA).is_empty() {
+            if let Some(auth_data) = &self.auth_data {
+                // string<lenenc>
+                buf.push(auth_data.len().to_le_bytes()[0]);
+                buf.push(auth_data.len().to_le_bytes()[1]);
+                buf.push(auth_data.len().to_le_bytes()[2]);
+                buf.extend_from_slice(&auth_data);
+            }
+        } else if !(self.capabilities & Capabilities::SECURE_CONNECTION).is_empty() {
+            if let Some(auth_response) = &self.auth_response {
+                buf.push(self.auth_response_len.unwrap());
+                buf.extend_from_slice(&auth_response);
+            }
+        } else {
+            buf.push(0);
+        }
+
+        if !(self.capabilities & Capabilities::CONNECT_WITH_DB).is_empty() {
+            if let Some(database) = &self.database {
+                // string<NUL>
+                buf.extend_from_slice(&database);
+                buf.push(0);
+            }
+        }
+
+        if !(self.capabilities & Capabilities::PLUGIN_AUTH).is_empty() {
+            if let Some(auth_plugin_name) = &self.auth_plugin_name {
+                // string<NUL>
+                buf.extend_from_slice(&auth_plugin_name);
+                buf.push(0);
+            }
+        }
+
+        if !(self.capabilities & Capabilities::CONNECT_ATTRS).is_empty() {
+            if let (Some(conn_attr_len), Some(conn_attr)) = (&self.conn_attr_len, &self.conn_attr) {
+                // int<lenenc>
+                buf.push(conn_attr_len.to_le_bytes().len().to_le_bytes()[0]);
+                buf.extend_from_slice(&conn_attr_len.to_le_bytes());
+
+                // Loop
+                for (key, value) in conn_attr {
+                    // string<lenenc>
+                    buf.push(key.len().to_le_bytes()[0]);
+                    buf.push(key.len().to_le_bytes()[1]);
+                    buf.push(key.len().to_le_bytes()[2]);
+                    buf.extend_from_slice(&key);
+
+                    // string<lenenc>
+                    buf.push(value.len().to_le_bytes()[0]);
+                    buf.push(value.len().to_le_bytes()[1]);
+                    buf.push(value.len().to_le_bytes()[2]);
+                    buf.extend_from_slice(&value);
+                }
+            }
+        }
+    }
+}