refactor string / byte reads from buffer

src/buffer.rs

@@ -0,0 +1,25 @@
+use bytes::{Buf, Bytes};
+
+use crate::{Result, errors::AppError};
+
+pub fn try_get_string<B: Buf>(buf: &mut B) -> Result<String> {
+    let len = buf.try_get_u16()? as usize;
+
+    if buf.remaining() <= len {
+        return Err(AppError::IncompleteBuffer);
+    }
+
+    Ok(String::from_utf8(buf.copy_to_bytes(len).to_vec())?)
+}
+
+pub fn try_get_bytes<B: Buf>(buf: &mut B) -> Result<Bytes> {
+    let len = buf.try_get_u32()? as usize;
+
+    if buf.remaining() < len {
+        return Err(AppError::IncompleteBuffer);
+    }
+
+    let data = buf.copy_to_bytes(len);
+
+    Ok(data)
+}

src/commands/delete.rs

@@ -1,8 +1,8 @@
 use std::io::Cursor;
 
-use bytes::{Buf as _, BufMut as _, BytesMut};
+use bytes::{BufMut as _, BytesMut};
 
-use crate::{Result, connection::Connection, database::Database, errors::AppError};
+use crate::{Result, buffer::try_get_string, connection::Connection, database::Database};
 
 #[derive(Debug, Clone)]
 pub struct Delete {
@@ -29,13 +29,7 @@ impl Delete {
     }
 
     pub fn parse(bytes: &mut Cursor<&[u8]>) -> Result<Self> {
-        let key_length = bytes.try_get_u16()? as usize;
-
-        if bytes.remaining() < key_length {
-            return Err(AppError::IncompleteCommandBuffer);
-        }
-
-        let key = String::from_utf8(bytes.copy_to_bytes(key_length).to_vec())?;
+        let key = try_get_string(bytes)?;
 
         Ok(Self { key })
     }

src/commands/expire.rs

@@ -2,7 +2,7 @@ use std::io::Cursor;
 
 use bytes::{Buf as _, Bytes};
 
-use crate::{Result, connection::Connection, database::Database, errors::AppError};
+use crate::{Result, buffer::try_get_string, connection::Connection, database::Database};
 
 #[derive(Debug, Clone)]
 pub struct Expire {
@@ -22,14 +22,7 @@ impl Expire {
     }
 
     pub fn parse(bytes: &mut Cursor<&[u8]>) -> Result<Self> {
-        let key_length = bytes.try_get_u16()? as usize;
-
-        if bytes.remaining() < key_length {
-            return Err(AppError::IncompleteCommandBuffer);
-        }
-
-        let key = String::from_utf8(bytes.copy_to_bytes(key_length).to_vec())?;
-
+        let key = try_get_string(bytes)?;
         let seconds = bytes.try_get_u64()?;
 
         Ok(Self { key, seconds })

src/commands/get.rs

@@ -1,8 +1,8 @@
 use std::io::Cursor;
 
-use bytes::{Buf as _, BufMut as _, BytesMut};
+use bytes::{BufMut as _, BytesMut};
 
-use crate::{Result, connection::Connection, database::Database, errors::AppError};
+use crate::{Result, buffer::try_get_string, connection::Connection, database::Database};
 
 #[derive(Debug, Clone)]
 pub struct Get {
@@ -31,13 +31,7 @@ impl Get {
     }
 
     pub fn parse(bytes: &mut Cursor<&[u8]>) -> Result<Self> {
-        let key_length = bytes.try_get_u16()? as usize;
-
-        if bytes.remaining() < key_length {
-            return Err(AppError::IncompleteCommandBuffer);
-        }
-
-        let key = String::from_utf8(bytes.copy_to_bytes(key_length).to_vec())?;
+        let key = try_get_string(bytes)?;
 
         Ok(Self { key })
     }

src/commands/has.rs

@@ -1,8 +1,8 @@
 use std::io::Cursor;
 
-use bytes::{Buf as _, Bytes};
+use bytes::Bytes;
 
-use crate::{Result, connection::Connection, database::Database, errors::AppError};
+use crate::{Result, buffer::try_get_string, connection::Connection, database::Database};
 
 #[derive(Debug, Clone)]
 pub struct Has {
@@ -21,13 +21,7 @@ impl Has {
     }
 
     pub fn parse(bytes: &mut Cursor<&[u8]>) -> Result<Self> {
-        let key_length = bytes.try_get_u16()? as usize;
-
-        if bytes.remaining() < key_length {
-            return Err(AppError::IncompleteCommandBuffer);
-        }
-
-        let key = String::from_utf8(bytes.copy_to_bytes(key_length).to_vec())?;
+        let key = try_get_string(bytes)?;
 
         Ok(Self { key })
     }

src/commands/mod.rs

@@ -8,7 +8,7 @@ mod ttl;
 
 use std::io::Cursor;
 
-use bytes::{Buf, BytesMut};
+use bytes::BytesMut;
 use delete::Delete;
 use expire::Expire;
 use get::Get;
@@ -17,7 +17,9 @@ use persist::Persist;
 use set::Set;
 use ttl::Ttl;
 
-use crate::{Result, connection::Connection, database::Database, errors::AppError};
+use crate::{
+    Result, buffer::try_get_string, connection::Connection, database::Database, errors::AppError,
+};
 
 #[derive(Debug)]
 pub enum Command {
@@ -46,13 +48,7 @@ impl Command {
     pub fn parse(bytes: &BytesMut) -> Result<(Self, u64)> {
         let mut buffer = Cursor::new(&bytes[..]);
 
-        let name_length = buffer.try_get_u16()? as usize;
-
-        if buffer.remaining() < name_length {
-            return Err(AppError::IncompleteCommandBuffer);
-        }
-
-        let name = String::from_utf8(buffer.copy_to_bytes(name_length).to_vec())?;
+        let name = try_get_string(&mut buffer)?;
 
         Self::parse_inner(name, &mut buffer)
     }

src/commands/persist.rs

@@ -1,8 +1,8 @@
 use std::io::Cursor;
 
-use bytes::{Buf as _, Bytes};
+use bytes::Bytes;
 
-use crate::{Result, connection::Connection, database::Database, errors::AppError};
+use crate::{Result, buffer::try_get_string, connection::Connection, database::Database};
 
 #[derive(Debug, Clone)]
 pub struct Persist {
@@ -21,13 +21,7 @@ impl Persist {
     }
 
     pub fn parse(bytes: &mut Cursor<&[u8]>) -> Result<Self> {
-        let key_length = bytes.try_get_u16()? as usize;
-
-        if bytes.remaining() < key_length {
-            return Err(AppError::IncompleteCommandBuffer);
-        }
-
-        let key = String::from_utf8(bytes.copy_to_bytes(key_length).to_vec())?;
+        let key = try_get_string(bytes)?;
 
         Ok(Self { key })
     }

src/commands/set.rs

@@ -1,6 +1,12 @@
 use std::io::Cursor;
 
-use crate::{Result, connection::Connection, database::Database, errors::AppError};
+use crate::{
+    Result,
+    buffer::{try_get_bytes, try_get_string},
+    connection::Connection,
+    database::Database,
+    errors::AppError,
+};
 use bytes::{Buf as _, Bytes};
 
 #[derive(Debug, Clone)]
@@ -20,21 +26,9 @@ impl Set {
     }
 
     pub fn parse(bytes: &mut Cursor<&[u8]>) -> Result<Self> {
-        let key_length = bytes.try_get_u16()? as usize;
+        let key = try_get_string(bytes)?;
 
-        if bytes.remaining() < key_length {
-            return Err(AppError::IncompleteCommandBuffer);
-        }
-
-        let key = String::from_utf8(bytes.copy_to_bytes(key_length).to_vec())?;
-
-        let value_length = bytes.try_get_u32()? as usize;
-
-        if bytes.remaining() < value_length {
-            return Err(AppError::IncompleteCommandBuffer);
-        }
-
-        let data = bytes.copy_to_bytes(value_length);
+        let data = try_get_bytes(bytes)?;
 
         let expiration: Option<u64> = match bytes.try_get_u8()? {
             1 => Some(bytes.try_get_u64()?),

src/commands/ttl.rs

@@ -1,8 +1,8 @@
 use std::io::Cursor;
 
-use bytes::{Buf as _, BufMut, Bytes, BytesMut};
+use bytes::{BufMut, Bytes, BytesMut};
 
-use crate::{Result, connection::Connection, database::Database, errors::AppError};
+use crate::{Result, buffer::try_get_string, connection::Connection, database::Database};
 
 #[derive(Debug, Clone)]
 pub struct Ttl {
@@ -29,13 +29,7 @@ impl Ttl {
     }
 
     pub fn parse(bytes: &mut Cursor<&[u8]>) -> Result<Self> {
-        let key_length = bytes.try_get_u16()? as usize;
-
-        if bytes.remaining() < key_length {
-            return Err(AppError::IncompleteCommandBuffer);
-        }
-
-        let key = String::from_utf8(bytes.copy_to_bytes(key_length).to_vec())?;
+        let key = try_get_string(bytes)?;
 
         Ok(Self { key })
     }

src/errors.rs

@@ -6,8 +6,8 @@ pub enum AppError {
     Io(#[from] std::io::Error),
     #[error("A TryGetError occurred")]
     TryGet(#[from] bytes::TryGetError),
-    #[error("The buffer is missing data for a complete command")]
-    IncompleteCommandBuffer,
+    #[error("The buffer is missing data")]
+    IncompleteBuffer,
     #[error("A Utf8Error occurred")]
     FromUtf8(#[from] std::string::FromUtf8Error),
     #[error("The command {0} was not recognized")]

src/lib.rs

@@ -3,6 +3,7 @@ use errors::AppError;
 #[cfg(test)]
 pub mod tests;
 
+pub mod buffer;
 pub mod client;
 pub mod commands;
 pub mod config;