refactor client byte write / read

src/buffer.rs

@@ -4,6 +4,7 @@ use crate::{Result, errors::AppError};
 
 pub trait ArchiveBuf<B: Buf> {
     fn try_get_string(&mut self) -> Result<String>;
+    fn try_get_short_string(&mut self) -> Result<String>;
     fn try_get_bytes(&mut self) -> Result<Bytes>;
     fn try_get_bool(&mut self) -> Result<bool>;
     fn try_get_option<T, F, E>(&mut self, f: F) -> Result<Option<T>>
@@ -14,6 +15,8 @@ pub trait ArchiveBuf<B: Buf> {
 }
 
 pub trait ArchiveBufMut<B: Buf> {
+    fn put_string(&mut self, s: &str) -> Result<()>;
+    fn put_short_string(&mut self, s: &str) -> Result<()>;
     fn put_bytes_with_length<T: AsRef<[u8]>>(&mut self, bytes: T);
     fn put_option<T, F>(&mut self, value: Option<T>, f: F)
     where
@@ -23,6 +26,16 @@ pub trait ArchiveBufMut<B: Buf> {
 
 impl<B: Buf> ArchiveBuf<B> for B {
     fn try_get_string(&mut self) -> Result<String> {
+        let len = self.try_get_u32()? as usize;
+
+        if self.remaining() < len {
+            return Err(AppError::IncompleteBuffer);
+        }
+
+        Ok(String::from_utf8(self.copy_to_bytes(len).to_vec())?)
+    }
+
+    fn try_get_short_string(&mut self) -> Result<String> {
         let len = self.try_get_u16()? as usize;
 
         if self.remaining() < len {
@@ -45,7 +58,11 @@ impl<B: Buf> ArchiveBuf<B> for B {
     }
 
     fn try_get_bool(&mut self) -> Result<bool> {
-        Ok(self.try_get_u8()? == 1)
+        Ok(match self.try_get_u8()? {
+            1 => true,
+            0 => false,
+            _ => return Err(AppError::UnexpectedData),
+        })
     }
 
     fn try_get_option<T, F, E>(&mut self, f: F) -> Result<Option<T>>
@@ -63,6 +80,20 @@ impl<B: Buf> ArchiveBuf<B> for B {
 }
 
 impl<B: Buf + BufMut> ArchiveBufMut<B> for B {
+    fn put_string(&mut self, s: &str) -> Result<()> {
+        self.put_u32(s.len().try_into()?);
+        self.put_slice(s.as_bytes());
+
+        Ok(())
+    }
+
+    fn put_short_string(&mut self, s: &str) -> Result<()> {
+        self.put_u16(s.len().try_into()?);
+        self.put_slice(s.as_bytes());
+
+        Ok(())
+    }
+
     fn put_bytes_with_length<T: AsRef<[u8]>>(&mut self, bytes: T) {
         let bytes = bytes.as_ref();
 

src/client.rs

@@ -1,7 +1,12 @@
 use bytes::{Buf, BufMut as _, Bytes, BytesMut};
 use tokio::net::{TcpStream, ToSocketAddrs};
 
-use crate::{Result, connection::Connection, errors::AppError};
+use crate::{
+    Result,
+    buffer::{ArchiveBuf, ArchiveBufMut as _},
+    connection::Connection,
+    errors::AppError,
+};
 
 pub struct Client {
     connection: Connection,
@@ -19,35 +24,16 @@ impl Client {
     pub async fn get(&mut self, key: &str) -> Result<Option<Bytes>> {
         let mut bytes = BytesMut::new();
 
-        bytes.put_u16(3);
+        bytes.put_short_string("get")?;
-        bytes.put_slice(b"get");
+        bytes.put_short_string(key)?;
-
-        let key_length: u16 = key
-            .len()
-            .try_into()
-            .map_err(|_| AppError::KeyLength(key.len()))?;
-
-        bytes.put_u16(key_length);
-        bytes.put_slice(key.as_bytes());
 
         self.connection.write(bytes.into()).await?;
 
         let mut r = self.get_response().await?;
-        let response = match r.try_get_u8()? {
-            0 => None,
-            1 => {
-                let len = r.try_get_u32()? as usize;
 
-                if r.remaining() < len {
+        let value = r.try_get_option(ArchiveBuf::try_get_bytes)?;
-                    return Err(AppError::InvalidCommandResponse);
-                }
 
-                Some(r.copy_to_bytes(len))
+        Ok(value)
-            }
-            _ => return Err(AppError::InvalidCommandResponse),
-        };
-
-        Ok(response)
     }
 
     pub async fn set(
@@ -58,130 +44,74 @@ impl Client {
     ) -> Result<()> {
         let mut bytes = BytesMut::new();
 
-        bytes.put_u16(3);
+        bytes.put_short_string("set")?;
-        bytes.put_slice(b"set");
+        bytes.put_short_string(key)?;
 
-        let key_length: u16 = key
+        bytes.put_bytes_with_length(data);
-            .len()
-            .try_into()
-            .map_err(|_| AppError::KeyLength(key.len()))?;
 
-        bytes.put_u16(key_length);
+        bytes.put_option(expiration_secs, BytesMut::put_u64);
-        bytes.put_slice(key.as_bytes());
-
-        bytes.put_u32(data.len() as u32);
-        bytes.put_slice(data);
-
-        match expiration_secs {
-            Some(seconds) => {
-                bytes.put_u8(1);
-                bytes.put_u64(seconds);
-            }
-            None => {
-                bytes.put_u8(0);
-            }
-        }
 
         self.connection.write(bytes.into()).await?;
 
         let mut r = self.get_response().await?;
 
-        match r.try_get_u8()? {
+        if !r.try_get_bool()? {
-            1 => return Ok(()),
+            return Err(AppError::InvalidCommandResponse);
-            _ => return Err(AppError::InvalidCommandResponse),
         }
+
+        Ok(())
     }
 
     pub async fn delete(&mut self, key: &str) -> Result<Option<Bytes>> {
         let mut bytes = BytesMut::new();
 
-        bytes.put_u16(6);
+        bytes.put_short_string("delete")?;
-        bytes.put_slice(b"delete");
+        bytes.put_short_string(key)?;
-
-        let key_length: u16 = key
-            .len()
-            .try_into()
-            .map_err(|_| AppError::KeyLength(key.len()))?;
-
-        bytes.put_u16(key_length);
-        bytes.put_slice(key.as_bytes());
 
         self.connection.write(bytes.into()).await?;
 
         let mut r = self.get_response().await?;
 
-        let response = match r.try_get_u8()? {
+        let value = r.try_get_option(ArchiveBuf::try_get_bytes)?;
-            1 => {
-                let len = r.try_get_u32()?;
-                let bytes = r.copy_to_bytes(len as usize);
 
-                Some(bytes)
+        Ok(value)
-            }
-            0 => None,
-            _ => return Err(AppError::InvalidCommandResponse),
-        };
-
-        Ok(response)
     }
 
     pub async fn has(&mut self, key: &str) -> Result<bool> {
         let mut bytes = BytesMut::new();
-        bytes.put_u16(3);
-        bytes.put_slice(b"has");
 
-        let key_length: u16 = key
+        bytes.put_short_string("has")?;
-            .len()
+        bytes.put_short_string(key)?;
-            .try_into()
-            .map_err(|_| AppError::KeyLength(key.len()))?;
-
-        bytes.put_u16(key_length);
-        bytes.put_slice(key.as_bytes());
 
         self.connection.write(bytes.into()).await?;
 
         let mut r = self.get_response().await?;
 
-        Ok(r.try_get_u8()? == 1)
+        let has = r.try_get_bool()?;
+
+        Ok(has)
     }
 
     pub async fn ttl(&mut self, key: &str) -> Result<Option<u64>> {
         let mut bytes = BytesMut::new();
-        bytes.put_u16(3);
-        bytes.put_slice(b"ttl");
 
-        let key_length: u16 = key
+        bytes.put_short_string("ttl")?;
-            .len()
+        bytes.put_short_string(key)?;
-            .try_into()
-            .map_err(|_| AppError::KeyLength(key.len()))?;
-
-        bytes.put_u16(key_length);
-        bytes.put_slice(key.as_bytes());
 
         self.connection.write(bytes.into()).await?;
 
         let mut r = self.get_response().await?;
 
-        let ttl = match r.try_get_u8()? {
+        let ttl = r.try_get_option(Bytes::try_get_u64)?;
-            1 => Some(r.try_get_u64()?),
-            0 => None,
-            _ => return Err(AppError::InvalidCommandResponse),
-        };
 
         Ok(ttl)
     }
 
     pub async fn expire(&mut self, key: &str, seconds: u64) -> Result<bool> {
         let mut bytes = BytesMut::new();
-        bytes.put_u16(6);
-        bytes.put_slice(b"expire");
 
-        let key_length: u16 = key
+        bytes.put_short_string("expire")?;
-            .len()
+        bytes.put_short_string(key)?;
-            .try_into()
-            .map_err(|_| AppError::KeyLength(key.len()))?;
-
-        bytes.put_u16(key_length);
-        bytes.put_slice(key.as_bytes());
 
         bytes.put_u64(seconds);
 
@@ -189,39 +119,19 @@ impl Client {
 
         let mut r = self.get_response().await?;
 
-        let success = match r.try_get_u8()? {
+        Ok(r.try_get_bool()?)
-            1 => true,
-            0 => false,
-            _ => return Err(AppError::InvalidCommandResponse),
-        };
-
-        Ok(success)
     }
 
     pub async fn persist(&mut self, key: &str) -> Result<bool> {
         let mut bytes = BytesMut::new();
-        bytes.put_u16(7);
+        bytes.put_short_string("persist")?;
-        bytes.put_slice(b"persist");
+        bytes.put_short_string(key)?;
-
-        let key_length: u16 = key
-            .len()
-            .try_into()
-            .map_err(|_| AppError::KeyLength(key.len()))?;
-
-        bytes.put_u16(key_length);
-        bytes.put_slice(key.as_bytes());
 
         self.connection.write(bytes.into()).await?;
 
         let mut r = self.get_response().await?;
 
-        let success = match r.try_get_u8()? {
+        Ok(r.try_get_bool()?)
-            1 => true,
-            0 => false,
-            _ => return Err(AppError::InvalidCommandResponse),
-        };
-
-        Ok(success)
     }
 
     async fn get_response(&mut self) -> Result<Bytes> {

src/commands/delete.rs

@@ -27,7 +27,7 @@ impl Delete {
     }
 
     pub fn parse(buf: &mut Cursor<&[u8]>) -> Result<Self> {
-        let key = buf.try_get_string()?;
+        let key = buf.try_get_short_string()?;
 
         Ok(Self { key })
     }

src/commands/expire.rs

@@ -22,7 +22,7 @@ impl Expire {
     }
 
     pub fn parse(buf: &mut Cursor<&[u8]>) -> Result<Self> {
-        let key = buf.try_get_string()?;
+        let key = buf.try_get_short_string()?;
         let seconds = buf.try_get_u64()?;
 
         Ok(Self { key, seconds })

src/commands/get.rs

@@ -27,7 +27,7 @@ impl Get {
     }
 
     pub fn parse(buf: &mut Cursor<&[u8]>) -> Result<Self> {
-        let key = buf.try_get_string()?;
+        let key = buf.try_get_short_string()?;
 
         Ok(Self { key })
     }

src/commands/has.rs

@@ -21,7 +21,7 @@ impl Has {
     }
 
     pub fn parse(buf: &mut Cursor<&[u8]>) -> Result<Self> {
-        let key = buf.try_get_string()?;
+        let key = buf.try_get_short_string()?;
 
         Ok(Self { key })
     }

src/commands/mod.rs

@@ -48,8 +48,7 @@ impl Command {
     pub fn parse(bytes: &BytesMut) -> Result<(Self, u64)> {
         let mut buf = Cursor::new(&bytes[..]);
 
-        let name = buf.try_get_string()?;
+        let name = buf.try_get_short_string()?;
-        println!("Command name: {name}, buf: {buf:?}");
 
         Self::parse_inner(name, &mut buf)
     }

src/commands/persist.rs

@@ -21,7 +21,7 @@ impl Persist {
     }
 
     pub fn parse(buf: &mut Cursor<&[u8]>) -> Result<Self> {
-        let key = buf.try_get_string()?;
+        let key = buf.try_get_short_string()?;
 
         Ok(Self { key })
     }

src/commands/set.rs

@@ -20,7 +20,7 @@ impl Set {
     }
 
     pub fn parse(buf: &mut Cursor<&[u8]>) -> Result<Self> {
-        let key = buf.try_get_string()?;
+        let key = buf.try_get_short_string()?;
 
         let data = buf.try_get_bytes()?;
 

src/commands/ttl.rs

@@ -27,7 +27,7 @@ impl Ttl {
     }
 
     pub fn parse(buf: &mut Cursor<&[u8]>) -> Result<Self> {
-        let key = buf.try_get_string()?;
+        let key = buf.try_get_short_string()?;
 
         Ok(Self { key })
     }

src/errors.rs

@@ -18,6 +18,8 @@ pub enum AppError {
     NoResponse,
     #[error("Expected a different response for the executed command")]
     InvalidCommandResponse,
-    #[error("The binary command data is not structured correctly")]
+    #[error("The binary data is not structured correctly")]
-    UnexpectedCommandData,
+    UnexpectedData,
+    #[error("Failed to convert integer")]
+    TryFromInt(#[from] std::num::TryFromIntError),
 }