409/archive
1
1
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 1 Wiki Activity

refactor client byte write / read

Browse Source
This commit is contained in:
Shawn
2025-07-01 16:59:39 +02:00
parent 87eb32eb5d
commit cff5c37a40
11 changed files with 80 additions and 138 deletions
Download Patch File Download Diff File Expand all files Collapse all files

33
src/buffer.rs
View File

@@ -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();

162
src/client.rs
View File

@@ -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_slice(b"get");
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());
bytes.put_short_string("get")?;
bytes.put_short_string(key)?;
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 {
return Err(AppError::InvalidCommandResponse);
}
let value = r.try_get_option(ArchiveBuf::try_get_bytes)?;
Some(r.copy_to_bytes(len))
}
_ => return Err(AppError::InvalidCommandResponse),
};
Ok(response)
Ok(value)
}
pub async fn set(
@@ -58,130 +44,74 @@ impl Client {
) -> Result<()> {
let mut bytes = BytesMut::new();
bytes.put_u16(3);
bytes.put_slice(b"set");
bytes.put_short_string("set")?;
bytes.put_short_string(key)?;
let key_length: u16 = key
.len()
.try_into()
.map_err(|_| AppError::KeyLength(key.len()))?;
bytes.put_bytes_with_length(data);
bytes.put_u16(key_length);
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);
}
}
bytes.put_option(expiration_secs, BytesMut::put_u64);
self.connection.write(bytes.into()).await?;
let mut r = self.get_response().await?;
match r.try_get_u8()? {
1 => return Ok(()),
_ => return Err(AppError::InvalidCommandResponse),
if !r.try_get_bool()? {
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_slice(b"delete");
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());
bytes.put_short_string("delete")?;
bytes.put_short_string(key)?;
self.connection.write(bytes.into()).await?;
let mut r = self.get_response().await?;
let response = match r.try_get_u8()? {
1 => {
let len = r.try_get_u32()?;
let bytes = r.copy_to_bytes(len as usize);
let value = r.try_get_option(ArchiveBuf::try_get_bytes)?;
Some(bytes)
}
0 => None,
_ => return Err(AppError::InvalidCommandResponse),
};
Ok(response)
Ok(value)
}
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
.len()
.try_into()
.map_err(|_| AppError::KeyLength(key.len()))?;
bytes.put_u16(key_length);
bytes.put_slice(key.as_bytes());
bytes.put_short_string("has")?;
bytes.put_short_string(key)?;
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
.len()
.try_into()
.map_err(|_| AppError::KeyLength(key.len()))?;
bytes.put_u16(key_length);
bytes.put_slice(key.as_bytes());
bytes.put_short_string("ttl")?;
bytes.put_short_string(key)?;
self.connection.write(bytes.into()).await?;
let mut r = self.get_response().await?;
let ttl = match r.try_get_u8()? {
1 => Some(r.try_get_u64()?),
0 => None,
_ => return Err(AppError::InvalidCommandResponse),
};
let ttl = r.try_get_option(Bytes::try_get_u64)?;
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
.len()
.try_into()
.map_err(|_| AppError::KeyLength(key.len()))?;
bytes.put_u16(key_length);
bytes.put_slice(key.as_bytes());
bytes.put_short_string("expire")?;
bytes.put_short_string(key)?;
bytes.put_u64(seconds);
@@ -189,39 +119,19 @@ impl Client {
let mut r = self.get_response().await?;
let success = match r.try_get_u8()? {
1 => true,
0 => false,
_ => return Err(AppError::InvalidCommandResponse),
};
Ok(success)
Ok(r.try_get_bool()?)
}
pub async fn persist(&mut self, key: &str) -> Result<bool> {
let mut bytes = BytesMut::new();
bytes.put_u16(7);
bytes.put_slice(b"persist");
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());
bytes.put_short_string("persist")?;
bytes.put_short_string(key)?;
self.connection.write(bytes.into()).await?;
let mut r = self.get_response().await?;
let success = match r.try_get_u8()? {
1 => true,
0 => false,
_ => return Err(AppError::InvalidCommandResponse),
};
Ok(success)
Ok(r.try_get_bool()?)
}
async fn get_response(&mut self) -> Result<Bytes> {

2
src/commands/delete.rs
View File

@@ -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 })
}

2
src/commands/expire.rs
View File

@@ -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 })

2
src/commands/get.rs
View File

@@ -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 })
}

2
src/commands/has.rs
View File

@@ -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 })
}

3
src/commands/mod.rs
View File

@@ -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()?;
println!("Command name: {name}, buf: {buf:?}");
let name = buf.try_get_short_string()?;
Self::parse_inner(name, &mut buf)
}

2
src/commands/persist.rs
View File

@@ -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 })
}

2
src/commands/set.rs
View File

@@ -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()?;

2
src/commands/ttl.rs
View File

@@ -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 })
}

6
src/errors.rs
View File

@@ -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")]
UnexpectedCommandData,
#[error("The binary data is not structured correctly")]
UnexpectedData,
#[error("Failed to convert integer")]
TryFromInt(#[from] std::num::TryFromIntError),
}