50 lines
1.4 KiB
Rust
50 lines
1.4 KiB
Rust
|
use bytes::{Buf, BytesMut};
|
||
|
|
use tokio_util::codec::{Decoder, Encoder};
|
||
|
|
|
||
|
|
use crate::error::WireError;
|
||
|
|
use crate::parser::{parse_message, ParsedCommand};
|
||
|
|
|
||
|
|
/// Tokio codec for framing wire protocol messages on a TCP stream.
|
||
|
|
///
|
||
|
|
/// The wire protocol is naturally length-prefixed:
|
||
|
|
/// the first 4 bytes of each message contain the total message length.
|
||
|
|
pub struct WireCodec;
|
||
|
|
|
||
|
|
impl Decoder for WireCodec {
|
||
|
|
type Item = ParsedCommand;
|
||
|
|
type Error = WireError;
|
||
|
|
|
||
|
|
fn decode(&mut self, src: &mut BytesMut) -> Result<Option<Self::Item>, Self::Error> {
|
||
|
|
if src.len() < 4 {
|
||
|
|
return Ok(None);
|
||
|
|
}
|
||
|
|
|
||
|
|
// Peek at message length
|
||
|
|
let msg_len = i32::from_le_bytes([src[0], src[1], src[2], src[3]]) as usize;
|
||
|
|
|
||
|
|
if src.len() < msg_len {
|
||
|
|
// Reserve space for the rest of the message
|
||
|
|
src.reserve(msg_len - src.len());
|
||
|
|
return Ok(None);
|
||
|
|
}
|
||
|
|
|
||
|
|
match parse_message(src)? {
|
||
|
|
Some((cmd, bytes_consumed)) => {
|
||
|
|
src.advance(bytes_consumed);
|
||
|
|
Ok(Some(cmd))
|
||
|
|
}
|
||
|
|
None => Ok(None),
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
/// Encoder for raw byte responses (already serialized by the command handlers).
|
||
|
|
impl Encoder<Vec<u8>> for WireCodec {
|
||
|
|
type Error = WireError;
|
||
|
|
|
||
|
|
fn encode(&mut self, item: Vec<u8>, dst: &mut BytesMut) -> Result<(), Self::Error> {
|
||
|
|
dst.extend_from_slice(&item);
|
||
|
|
Ok(())
|
||
|
|
}
|
||
|
|
}
|