tokio_netstring/
lib.rs

1use tokio_util::codec::{Decoder, Encoder};
2use bytes::{BytesMut, BufMut};
3use std::io::{Error, ErrorKind};
4
5
6#[allow(unused_imports)]
7use futures::sink::SinkExt;
8
9
10fn check(left: usize, right: usize, data: &[u8]) -> Result<(), Error> {
11    if let Ok(num_string) = String::from_utf8(Vec::from(&data[..left])) {
12        if let Ok(exp_size) = num_string.parse::<usize>() {
13            if right > left && (right - 1) - left == exp_size {
14                return Ok(());
15            }
16        }
17    }
18    Err(Error::new(ErrorKind::InvalidData, "Invalid data"))
19}
20
21
22fn extract_frameborders(src: &BytesMut) -> Option<(usize, usize)> {
23    let left_idx = src.iter().position(|&b| b == b':');
24    let right_idx = src.iter().position(|&b| b == b',');
25
26    match (left_idx, right_idx) {
27        (Some(l), Some(r)) => Some((l, r)),
28        _ => None
29    }
30}
31
32pub struct NetStringCodec {}
33
34impl NetStringCodec {
35    pub fn extract_frame(&self, src: &mut BytesMut) -> Result<Option<String>, Box<dyn std::error::Error>> {
36        if let Some((lhs, rhs)) = extract_frameborders(src) {
37            let data = src.split_to(rhs + 1); // <- modify src
38            check(lhs, rhs, &data[..])?;
39            let data = String::from_utf8(Vec::from(&data[lhs + 1..rhs]))?;
40            return Ok(Some(data));
41        }
42        Ok(None)
43    }
44}
45
46impl Decoder for NetStringCodec {
47    type Item = String;
48    type Error = Error;
49
50    fn decode(&mut self, src: &mut BytesMut) -> Result<Option<Self::Item>, Self::Error> {
51        println!("Trying to decode {:?}", src);
52
53        match self.extract_frame(src) {
54            Ok(frame) => Ok(frame),
55            Err(_) => Err(Error::new(ErrorKind::InvalidData, "Could not parse data"))
56        }
57    }
58}
59
60
61impl Encoder<String> for NetStringCodec {
62    type Error = std::io::Error;
63
64    fn encode(&mut self, item: String, dst: &mut BytesMut) -> Result<(), Self::Error> {
65        println!("Encoding {:?}", item);
66        let data = format!("{}:{},", item.len(), item);
67        println!("Encoded data: {:?}", data);
68        dst.put(data.as_bytes());
69        Ok(())
70    }
71}
72