Implement protocol stage 1

Cargo.toml

@@ -1,4 +1,5 @@
 [workspace]
+resolver = "2"
 members = [
     "cli",
     "server",

server/src/engine/mem/buf.rs

@@ -29,36 +29,137 @@ use core::{ptr, slice};
 #[derive(Debug)]
 pub struct BufferedScanner<'a> {
     d: &'a [u8],
-    i: usize,
+    __cursor: usize,
 }
 
 impl<'a> BufferedScanner<'a> {
     pub const fn new(d: &'a [u8]) -> Self {
-        Self { d, i: 0 }
+        unsafe { Self::new_with_cursor(d, 0) }
+    }
+    pub const unsafe fn new_with_cursor(d: &'a [u8], i: usize) -> Self {
+        Self { d, __cursor: i }
     }
     pub const fn remaining(&self) -> usize {
-        self.d.len() - self.i
+        self.d.len() - self.__cursor
     }
     pub const fn consumed(&self) -> usize {
-        self.i
+        self.__cursor
     }
     pub const fn cursor(&self) -> usize {
-        self.i
+        self.__cursor
-    }
-    pub(crate) fn has_left(&self, sizeof: usize) -> bool {
-        self.remaining() >= sizeof
     }
-    unsafe fn _cursor(&self) -> *const u8 {
+    pub fn current(&self) -> &[u8] {
-        self.d.as_ptr().add(self.i)
+        &self.d[self.__cursor..]
     }
     pub fn eof(&self) -> bool {
         self.remaining() == 0
     }
+    pub fn has_left(&self, sizeof: usize) -> bool {
+        self.remaining() >= sizeof
+    }
+    pub fn matches_cursor_rounded(&self, f: impl Fn(u8) -> bool) -> bool {
+        f(self.d[self.d.len().min(self.__cursor)])
+    }
+    pub fn matches_cursor_rounded_and_not_eof(&self, f: impl Fn(u8) -> bool) -> bool {
+        self.matches_cursor_rounded(f) & !self.eof()
+    }
+}
+
+impl<'a> BufferedScanner<'a> {
+    pub unsafe fn set_cursor(&mut self, i: usize) {
+        self.__cursor = i;
+    }
+    pub unsafe fn move_ahead_by(&mut self, by: usize) {
+        self._incr(by)
+    }
+    pub unsafe fn move_back(&mut self) {
+        self.move_back_by(1)
+    }
+    pub unsafe fn move_back_by(&mut self, by: usize) {
+        self.__cursor -= by;
+    }
     unsafe fn _incr(&mut self, by: usize) {
-        self.i += by;
+        self.__cursor += by;
     }
-    pub fn current(&self) -> &[u8] {
+    unsafe fn _cursor(&self) -> *const u8 {
-        &self.d[self.i..]
+        self.d.as_ptr().add(self.__cursor)
+    }
+}
+
+impl<'a> BufferedScanner<'a> {
+    pub fn try_next_byte(&mut self) -> Option<u8> {
+        if self.eof() {
+            None
+        } else {
+            Some(unsafe { self.next_byte() })
+        }
+    }
+    pub fn try_next_block<const N: usize>(&mut self) -> Option<[u8; N]> {
+        if self.has_left(N) {
+            Some(unsafe { self.next_chunk() })
+        } else {
+            None
+        }
+    }
+    pub fn try_next_variable_block(&'a mut self, len: usize) -> Option<&'a [u8]> {
+        if self.has_left(len) {
+            Some(unsafe { self.next_chunk_variable(len) })
+        } else {
+            None
+        }
+    }
+}
+
+pub enum BufferedReadResult<T> {
+    Value(T),
+    NeedMore,
+    Error,
+}
+
+impl<'a> BufferedScanner<'a> {
+    /// Attempt to parse a `\n` terminated (we move past the LF, so you can't see it)
+    ///
+    /// If we were unable to read in the integer, then the cursor will be restored to its starting position
+    // TODO(@ohsayan): optimize
+    pub fn try_next_ascii_u64_lf_separated(&mut self) -> BufferedReadResult<u64> {
+        let mut okay = true;
+        let start = self.cursor();
+        let mut ret = 0u64;
+        while self.matches_cursor_rounded_and_not_eof(|b| b != b'\n') & okay {
+            let b = self.d[self.cursor()];
+            okay &= b.is_ascii_digit();
+            ret = match ret.checked_mul(10) {
+                Some(r) => r,
+                None => {
+                    okay = false;
+                    break;
+                }
+            };
+            ret = match ret.checked_add((b & 0x0F) as u64) {
+                Some(r) => r,
+                None => {
+                    okay = false;
+                    break;
+                }
+            };
+            unsafe { self._incr(1) }
+        }
+        let payload_ok = okay;
+        let null_ok = self.matches_cursor_rounded_and_not_eof(|b| b == b'\n');
+        okay &= null_ok;
+        unsafe { self._incr(okay as _) }; // skip LF
+        if okay {
+            BufferedReadResult::Value(ret)
+        } else {
+            unsafe { self.set_cursor(start) }
+            if payload_ok {
+                // payload was ok, but we missed a null
+                BufferedReadResult::NeedMore
+            } else {
+                // payload was NOT ok
+                BufferedReadResult::Error
+            }
+        }
     }
 }
 

server/src/engine/mem/mod.rs

@@ -25,7 +25,7 @@
 */
 
 mod astr;
-mod buf;
+pub mod buf;
 mod ll;
 mod stackop;
 mod uarray;

server/src/engine/mod.rs

@@ -34,6 +34,7 @@ mod error;
 mod fractal;
 mod idx;
 mod mem;
+mod net;
 mod ql;
 mod storage;
 mod sync;

server/src/engine/net/mod.rs

@@ -0,0 +1,27 @@
+/*
+ * Created on Fri Sep 15 2023
+ *
+ * This file is a part of Skytable
+ * Skytable (formerly known as TerrabaseDB or Skybase) is a free and open-source
+ * NoSQL database written by Sayan Nandan ("the Author") with the
+ * vision to provide flexibility in data modelling without compromising
+ * on performance, queryability or scalability.
+ *
+ * Copyright (c) 2023, Sayan Nandan <ohsayan@outlook.com>
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU Affero General Public License as published by
+ * the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+ * GNU Affero General Public License for more details.
+ *
+ * You should have received a copy of the GNU Affero General Public License
+ * along with this program. If not, see <https://www.gnu.org/licenses/>.
+ *
+*/
+
+mod protocol;

server/src/engine/net/protocol/handshake.rs

@@ -0,0 +1,406 @@
+/*
+ * Created on Mon Sep 18 2023
+ *
+ * This file is a part of Skytable
+ * Skytable (formerly known as TerrabaseDB or Skybase) is a free and open-source
+ * NoSQL database written by Sayan Nandan ("the Author") with the
+ * vision to provide flexibility in data modelling without compromising
+ * on performance, queryability or scalability.
+ *
+ * Copyright (c) 2023, Sayan Nandan <ohsayan@outlook.com>
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU Affero General Public License as published by
+ * the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+ * GNU Affero General Public License for more details.
+ *
+ * You should have received a copy of the GNU Affero General Public License
+ * along with this program. If not, see <https://www.gnu.org/licenses/>.
+ *
+*/
+
+use {
+    crate::{
+        engine::mem::buf::{BufferedReadResult, BufferedScanner},
+        util::compiler,
+    },
+    std::slice,
+};
+
+#[derive(Debug, PartialEq, Eq, Clone, Copy, sky_macros::EnumMethods)]
+#[repr(u8)]
+/// Low-level protocol errors
+pub enum ProtocolError {
+    /// packet has incorrect structure
+    CorruptedHSPacket = 0,
+    /// incorrect handshake version
+    RejectHSVersion = 1,
+    /// invalid protocol version
+    RejectProtocol = 2,
+    /// invalid exchange mode
+    RejectExchangeMode = 3,
+    /// invalid query mode
+    RejectQueryMode = 4,
+    /// invalid auth details
+    ///
+    /// **NB**: this can be due to either an incorrect auth flag, or incorrect auth data or disallowed auth mode. we keep it
+    /// in one error for purposes of security
+    RejectAuth = 5,
+}
+
+/*
+    handshake meta
+*/
+
+#[derive(Debug, PartialEq, Eq, Clone, Copy, sky_macros::EnumMethods)]
+#[repr(u8)]
+/// the handshake version
+pub enum HandshakeVersion {
+    /// Skyhash/2.0 HS
+    Original = 0,
+}
+
+#[derive(Debug, PartialEq, Eq, Clone, Copy, sky_macros::EnumMethods)]
+#[repr(u8)]
+/// the skyhash protocol version
+pub enum ProtocolVersion {
+    /// Skyhash/2.0 protocol
+    Original = 0,
+}
+
+#[derive(Debug, PartialEq, Eq, Clone, Copy, sky_macros::EnumMethods)]
+#[repr(u8)]
+/// the data exchange mode
+pub enum DataExchangeMode {
+    /// query-time data exchange mode
+    QueryTime = 0,
+}
+
+#[derive(Debug, PartialEq, Eq, Clone, Copy, sky_macros::EnumMethods)]
+#[repr(u8)]
+/// the query mode
+pub enum QueryMode {
+    /// BQL-1 query mode
+    Bql1 = 0,
+}
+
+#[derive(Debug, PartialEq, Eq, Clone, Copy, sky_macros::EnumMethods)]
+#[repr(u8)]
+/// the authentication mode
+pub enum AuthMode {
+    Anonymous = 0,
+    Password = 1,
+}
+
+impl AuthMode {
+    unsafe fn from_raw(v: u8) -> Self {
+        core::mem::transmute(v)
+    }
+    /// returns the minimum number of metadata bytes need to parse the payload for this auth mode
+    const fn min_payload_bytes(&self) -> usize {
+        match self {
+            Self::Anonymous => 1,
+            Self::Password => 4,
+        }
+    }
+}
+
+/*
+    client handshake
+*/
+
+/// The handshake state
+#[derive(Debug, PartialEq, Clone)]
+pub enum HandshakeState {
+    /// we just began the handshake
+    Initial,
+    /// we just processed the static block
+    StaticBlock(CHandshakeStatic),
+    /// Expecting some more auth meta
+    ExpectingMetaForVariableBlock {
+        /// static block
+        static_hs: CHandshakeStatic,
+        /// uname len
+        uname_l: usize,
+    },
+    /// we're expecting to finish the handshake
+    ExpectingVariableBlock {
+        /// static block
+        static_hs: CHandshakeStatic,
+        /// uname len
+        uname_l: usize,
+        /// pwd len
+        pwd_l: usize,
+    },
+}
+
+impl Default for HandshakeState {
+    fn default() -> Self {
+        Self::Initial
+    }
+}
+
+/// The static segment of the handshake
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub struct CHandshakeStatic {
+    /// the handshake version
+    hs_version: HandshakeVersion,
+    /// protocol version
+    protocol: ProtocolVersion,
+    /// exchange mode
+    exchange_mode: DataExchangeMode,
+    /// query mode
+    query_mode: QueryMode,
+    /// authentication mode
+    auth_mode: AuthMode,
+}
+
+impl CHandshakeStatic {
+    pub const fn new(
+        hs_version: HandshakeVersion,
+        protocol: ProtocolVersion,
+        exchange_mode: DataExchangeMode,
+        query_mode: QueryMode,
+        auth_mode: AuthMode,
+    ) -> Self {
+        Self {
+            hs_version,
+            protocol,
+            exchange_mode,
+            query_mode,
+            auth_mode,
+        }
+    }
+}
+
+/// handshake authentication
+// TODO(@ohsayan): enum?
+#[derive(Debug, PartialEq)]
+pub struct CHandshakeAuth<'a> {
+    username: &'a [u8],
+    password: &'a [u8],
+}
+
+impl<'a> CHandshakeAuth<'a> {
+    pub fn new(username: &'a [u8], password: &'a [u8]) -> Self {
+        Self { username, password }
+    }
+}
+
+#[derive(Debug, PartialEq)]
+pub enum HandshakeResult<'a> {
+    /// Finished handshake
+    Completed(CHandshake<'a>),
+    /// Update handshake state
+    ///
+    /// **NOTE:** expect does not take into account the current amount of buffered data (hence the unbuffered part must be computed!)
+    ChangeState {
+        new_state: HandshakeState,
+        expect: usize,
+    },
+    /// An error occurred
+    Error(ProtocolError),
+}
+
+/// The client's handshake record
+#[derive(Debug, PartialEq)]
+pub struct CHandshake<'a> {
+    /// the static segment of the handshake
+    hs_static: CHandshakeStatic,
+    /// the auth section of the dynamic segment of the handshake
+    hs_auth: Option<CHandshakeAuth<'a>>,
+}
+
+impl<'a> CHandshake<'a> {
+    pub const INITIAL_READ: usize = 6;
+    const CLIENT_HELLO: u8 = b'H';
+    pub fn new(hs_static: CHandshakeStatic, hs_auth: Option<CHandshakeAuth<'a>>) -> Self {
+        Self { hs_static, hs_auth }
+    }
+    /// Resume handshake with the given state and buffer
+    pub fn resume_with(
+        scanner: &mut BufferedScanner<'a>,
+        state: HandshakeState,
+    ) -> HandshakeResult<'a> {
+        match state {
+            // nothing buffered yet
+            HandshakeState::Initial => Self::resume_initial(scanner),
+            // buffered static block
+            HandshakeState::StaticBlock(static_block) => {
+                Self::resume_at_auth_metadata1(scanner, static_block)
+            }
+            // buffered some auth meta
+            HandshakeState::ExpectingMetaForVariableBlock { static_hs, uname_l } => {
+                Self::resume_at_auth_metadata2(scanner, static_hs, uname_l)
+            }
+            // buffered full auth meta
+            HandshakeState::ExpectingVariableBlock {
+                static_hs,
+                uname_l,
+                pwd_l,
+            } => Self::resume_at_variable_block_payload(scanner, static_hs, uname_l, pwd_l),
+        }
+    }
+}
+
+impl<'a> CHandshake<'a> {
+    /// Resume from the initial state (nothing buffered yet)
+    fn resume_initial(scanner: &mut BufferedScanner<'a>) -> HandshakeResult<'a> {
+        // get our block
+        if cfg!(debug_assertions) {
+            if scanner.remaining() < Self::INITIAL_READ {
+                return HandshakeResult::ChangeState {
+                    new_state: HandshakeState::Initial,
+                    expect: Self::INITIAL_READ,
+                };
+            }
+        } else {
+            assert!(scanner.remaining() >= Self::INITIAL_READ);
+        }
+        let buf: [u8; CHandshake::INITIAL_READ] = unsafe { scanner.next_chunk() };
+        let invalid_first_byte = buf[0] != Self::CLIENT_HELLO;
+        let invalid_hs_version = buf[1] > HandshakeVersion::MAX;
+        let invalid_proto_version = buf[2] > ProtocolVersion::MAX;
+        let invalid_exchange_mode = buf[3] > DataExchangeMode::MAX;
+        let invalid_query_mode = buf[4] > QueryMode::MAX;
+        let invalid_auth_mode = buf[5] > AuthMode::MAX;
+        // check block
+        if compiler::unlikely(
+            invalid_first_byte
+                | invalid_hs_version
+                | invalid_proto_version
+                | invalid_exchange_mode
+                | invalid_query_mode
+                | invalid_auth_mode,
+        ) {
+            static ERROR: [ProtocolError; 6] = [
+                ProtocolError::CorruptedHSPacket,
+                ProtocolError::RejectHSVersion,
+                ProtocolError::RejectProtocol,
+                ProtocolError::RejectExchangeMode,
+                ProtocolError::RejectQueryMode,
+                ProtocolError::RejectAuth,
+            ];
+            return HandshakeResult::Error(
+                ERROR[((invalid_first_byte as u8 * 1)
+                    | (invalid_hs_version as u8 * 2)
+                    | (invalid_proto_version as u8 * 3)
+                    | (invalid_exchange_mode as u8 * 4)
+                    | (invalid_query_mode as u8 * 5)
+                    | (invalid_auth_mode as u8) * 6) as usize
+                    - 1usize],
+            );
+        }
+        // init header
+        let static_header = CHandshakeStatic::new(
+            HandshakeVersion::Original,
+            ProtocolVersion::Original,
+            DataExchangeMode::QueryTime,
+            QueryMode::Bql1,
+            unsafe {
+                // UNSAFE(@ohsayan): already checked
+                AuthMode::from_raw(buf[5])
+            },
+        );
+        // check if we have auth data
+        Self::resume_at_auth_metadata1(scanner, static_header)
+    }
+    fn resume_at_variable_block_payload(
+        scanner: &mut BufferedScanner<'a>,
+        static_hs: CHandshakeStatic,
+        uname_l: usize,
+        pwd_l: usize,
+    ) -> HandshakeResult<'a> {
+        if scanner.has_left(uname_l + pwd_l) {
+            // we're done here
+            return unsafe {
+                // UNSAFE(@ohsayan): we just checked buffered size
+                let uname = slice::from_raw_parts(scanner.current().as_ptr(), uname_l);
+                let pwd = slice::from_raw_parts(scanner.current().as_ptr().add(uname_l), pwd_l);
+                scanner.move_ahead_by(uname_l + pwd_l);
+                HandshakeResult::Completed(Self::new(
+                    static_hs,
+                    Some(CHandshakeAuth::new(uname, pwd)),
+                ))
+            };
+        }
+        HandshakeResult::ChangeState {
+            new_state: HandshakeState::ExpectingVariableBlock {
+                static_hs,
+                uname_l,
+                pwd_l,
+            },
+            expect: (uname_l + pwd_l),
+        }
+    }
+}
+
+impl<'a> CHandshake<'a> {
+    /// Resume parsing at the first part of the auth metadata
+    fn resume_at_auth_metadata1(
+        scanner: &mut BufferedScanner<'a>,
+        static_header: CHandshakeStatic,
+    ) -> HandshakeResult<'a> {
+        // now let's see if we have buffered enough data for auth
+        if scanner.remaining() < static_header.auth_mode.min_payload_bytes() {
+            // we need more data
+            return HandshakeResult::ChangeState {
+                new_state: HandshakeState::StaticBlock(static_header),
+                expect: static_header.auth_mode.min_payload_bytes(),
+            };
+        }
+        // we seem to have enough data for this auth mode
+        match static_header.auth_mode {
+            AuthMode::Anonymous => {
+                if unsafe { scanner.next_byte() } == 0 {
+                    // matched
+                    return HandshakeResult::Completed(Self::new(static_header, None));
+                }
+                // we can only accept a NUL byte
+                return HandshakeResult::Error(ProtocolError::RejectAuth);
+            }
+            AuthMode::Password => {}
+        }
+        // let us see if we can parse the username length
+        let uname_l = match scanner.try_next_ascii_u64_lf_separated() {
+            BufferedReadResult::NeedMore => {
+                return HandshakeResult::ChangeState {
+                    new_state: HandshakeState::StaticBlock(static_header),
+                    expect: AuthMode::Password.min_payload_bytes(), // 2 for uname_l and 2 for pwd_l
+                };
+            }
+            BufferedReadResult::Value(v) => v as usize,
+            BufferedReadResult::Error => {
+                return HandshakeResult::Error(ProtocolError::CorruptedHSPacket)
+            }
+        };
+        Self::resume_at_auth_metadata2(scanner, static_header, uname_l)
+    }
+    /// Resume at trying to get the final part of the auth metadata
+    fn resume_at_auth_metadata2(
+        scanner: &mut BufferedScanner<'a>,
+        static_hs: CHandshakeStatic,
+        uname_l: usize,
+    ) -> HandshakeResult<'a> {
+        // we just have to get the password len
+        let pwd_l = match scanner.try_next_ascii_u64_lf_separated() {
+            BufferedReadResult::Value(v) => v as usize,
+            BufferedReadResult::NeedMore => {
+                // newline missing (or maybe there's more?)
+                return HandshakeResult::ChangeState {
+                    new_state: HandshakeState::ExpectingMetaForVariableBlock { static_hs, uname_l },
+                    expect: uname_l + 2, // space for username + password len
+                };
+            }
+            BufferedReadResult::Error => {
+                return HandshakeResult::Error(ProtocolError::CorruptedHSPacket)
+            }
+        };
+        Self::resume_at_variable_block_payload(scanner, static_hs, uname_l, pwd_l)
+    }
+}

server/src/engine/net/protocol/mod.rs

@@ -0,0 +1,29 @@
+/*
+ * Created on Fri Sep 15 2023
+ *
+ * This file is a part of Skytable
+ * Skytable (formerly known as TerrabaseDB or Skybase) is a free and open-source
+ * NoSQL database written by Sayan Nandan ("the Author") with the
+ * vision to provide flexibility in data modelling without compromising
+ * on performance, queryability or scalability.
+ *
+ * Copyright (c) 2023, Sayan Nandan <ohsayan@outlook.com>
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU Affero General Public License as published by
+ * the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+ * GNU Affero General Public License for more details.
+ *
+ * You should have received a copy of the GNU Affero General Public License
+ * along with this program. If not, see <https://www.gnu.org/licenses/>.
+ *
+*/
+
+mod handshake;
+#[cfg(test)]
+mod tests;

server/src/engine/net/protocol/tests.rs

@@ -0,0 +1,197 @@
+/*
+ * Created on Mon Sep 18 2023
+ *
+ * This file is a part of Skytable
+ * Skytable (formerly known as TerrabaseDB or Skybase) is a free and open-source
+ * NoSQL database written by Sayan Nandan ("the Author") with the
+ * vision to provide flexibility in data modelling without compromising
+ * on performance, queryability or scalability.
+ *
+ * Copyright (c) 2023, Sayan Nandan <ohsayan@outlook.com>
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU Affero General Public License as published by
+ * the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+ * GNU Affero General Public License for more details.
+ *
+ * You should have received a copy of the GNU Affero General Public License
+ * along with this program. If not, see <https://www.gnu.org/licenses/>.
+ *
+*/
+
+use crate::engine::mem::BufferedScanner;
+use crate::engine::net::protocol::handshake::{
+    AuthMode, CHandshake, CHandshakeAuth, CHandshakeStatic, DataExchangeMode, HandshakeResult,
+    HandshakeState, HandshakeVersion, ProtocolVersion, QueryMode,
+};
+
+const FULL_HANDSHAKE_NO_AUTH: [u8; 7] = [b'H', 0, 0, 0, 0, 0, 0];
+const FULL_HANDSHAKE_WITH_AUTH: [u8; 23] = *b"H\0\0\0\0\x015\n8\nsayanpass1234";
+
+const STATIC_HANDSHAKE_NO_AUTH: CHandshakeStatic = CHandshakeStatic::new(
+    HandshakeVersion::Original,
+    ProtocolVersion::Original,
+    DataExchangeMode::QueryTime,
+    QueryMode::Bql1,
+    AuthMode::Anonymous,
+);
+
+const STATIC_HANDSHAKE_WITH_AUTH: CHandshakeStatic = CHandshakeStatic::new(
+    HandshakeVersion::Original,
+    ProtocolVersion::Original,
+    DataExchangeMode::QueryTime,
+    QueryMode::Bql1,
+    AuthMode::Password,
+);
+
+/*
+    handshake with no state changes
+*/
+
+#[test]
+fn parse_staged_no_auth() {
+    for i in 0..FULL_HANDSHAKE_NO_AUTH.len() {
+        let buf = &FULL_HANDSHAKE_NO_AUTH[..i + 1];
+        let mut scanner = BufferedScanner::new(buf);
+        let result = CHandshake::resume_with(&mut scanner, HandshakeState::Initial);
+        match buf.len() {
+            1..=5 => {
+                assert_eq!(
+                    result,
+                    HandshakeResult::ChangeState {
+                        new_state: HandshakeState::Initial,
+                        expect: CHandshake::INITIAL_READ,
+                    }
+                );
+            }
+            6 => {
+                assert_eq!(
+                    result,
+                    HandshakeResult::ChangeState {
+                        new_state: HandshakeState::StaticBlock(STATIC_HANDSHAKE_NO_AUTH),
+                        expect: 1,
+                    }
+                );
+            }
+            7 => {
+                assert_eq!(
+                    result,
+                    HandshakeResult::Completed(CHandshake::new(STATIC_HANDSHAKE_NO_AUTH, None))
+                );
+            }
+            _ => unreachable!(),
+        }
+    }
+}
+
+#[test]
+fn parse_staged_with_auth() {
+    for i in 0..FULL_HANDSHAKE_WITH_AUTH.len() {
+        let buf = &FULL_HANDSHAKE_WITH_AUTH[..i + 1];
+        let mut s = BufferedScanner::new(buf);
+        let ref mut scanner = s;
+        let result = CHandshake::resume_with(scanner, HandshakeState::Initial);
+        match buf.len() {
+            1..=5 => {
+                assert_eq!(
+                    result,
+                    HandshakeResult::ChangeState {
+                        new_state: HandshakeState::Initial,
+                        expect: CHandshake::INITIAL_READ
+                    }
+                );
+            }
+            6..=9 => {
+                // might seem funny that we don't parse the second integer at all, but it's because
+                // of the relatively small size of the integers
+                assert_eq!(
+                    result,
+                    HandshakeResult::ChangeState {
+                        new_state: HandshakeState::StaticBlock(STATIC_HANDSHAKE_WITH_AUTH),
+                        expect: 4
+                    }
+                );
+            }
+            10..=22 => {
+                assert_eq!(
+                    result,
+                    HandshakeResult::ChangeState {
+                        new_state: HandshakeState::ExpectingVariableBlock {
+                            static_hs: STATIC_HANDSHAKE_WITH_AUTH,
+                            uname_l: 5,
+                            pwd_l: 8
+                        },
+                        expect: 13,
+                    }
+                );
+            }
+            23 => {
+                assert_eq!(
+                    result,
+                    HandshakeResult::Completed(CHandshake::new(
+                        STATIC_HANDSHAKE_WITH_AUTH,
+                        Some(CHandshakeAuth::new(b"sayan", b"pass1234"))
+                    ))
+                );
+            }
+            _ => unreachable!(),
+        }
+    }
+}
+
+/*
+    handshake with state changes
+*/
+
+fn run_state_changes_return_rounds(src: &[u8], expected_final_handshake: CHandshake) -> usize {
+    let mut rounds = 0;
+    let hs;
+    let mut state = HandshakeState::default();
+    let mut cursor = 0;
+    let mut expect_many = CHandshake::INITIAL_READ;
+    loop {
+        rounds += 1;
+        let buf = &src[..cursor + expect_many];
+        let mut scanner = unsafe { BufferedScanner::new_with_cursor(buf, cursor) };
+        match CHandshake::resume_with(&mut scanner, state) {
+            HandshakeResult::ChangeState { new_state, expect } => {
+                state = new_state;
+                expect_many = expect;
+                cursor = scanner.cursor();
+            }
+            HandshakeResult::Completed(c) => {
+                hs = c;
+                assert_eq!(hs, expected_final_handshake);
+                break;
+            }
+            HandshakeResult::Error(e) => panic!("unexpected handshake error: {:?}", e),
+        }
+    }
+    rounds
+}
+
+#[test]
+fn parse_no_auth_with_state_updates() {
+    let rounds = run_state_changes_return_rounds(
+        &FULL_HANDSHAKE_NO_AUTH,
+        CHandshake::new(STATIC_HANDSHAKE_NO_AUTH, None),
+    );
+    assert_eq!(rounds, 2); // r1 = initial, r2 = auth NUL
+}
+
+#[test]
+fn parse_auth_with_state_updates() {
+    let rounds = run_state_changes_return_rounds(
+        &FULL_HANDSHAKE_WITH_AUTH,
+        CHandshake::new(
+            STATIC_HANDSHAKE_WITH_AUTH,
+            Some(CHandshakeAuth::new(b"sayan", b"pass1234")),
+        ),
+    );
+    assert_eq!(rounds, 3); // r1 = initial read, r2 = lengths, r3 = items
+}

server/src/engine/ql/dml/upd.rs

@@ -238,7 +238,7 @@ mod impls {
     }
     #[cfg(test)]
     mod test {
-        use super::{super::AssignmentExpression, ASTNode, QueryResult, QueryData, State};
+        use super::{super::AssignmentExpression, ASTNode, QueryData, QueryResult, State};
         impl<'a> ASTNode<'a> for AssignmentExpression<'a> {
             // important: upstream must verify this
             const VERIFY: bool = true;