tokio bencode string decoder

decode.rs

use bytes::Buf;
use tokio_util::codec;

use super::BeEncode;

pub struct StringDecoder;

impl codec::Decoder for StringDecoder {
    type Item = BeEncode;

    type Error = std::io::Error;

    fn decode(&mut self, buf: &mut bytes::BytesMut) -> Result<Option<Self::Item>, Self::Error> {
        let colon_idx = buf.iter().position(|&b| b == b':');
        if colon_idx.is_none() {
            return Ok(None);
        }

        let colon_idx = colon_idx.unwrap();
        let length_str = std::str::from_utf8(&buf[..colon_idx]).unwrap();
        let length = length_str.parse::<usize>().unwrap();

        if buf.len() < colon_idx + 1 + length {
            return Ok(None);
        }

        let content = std::str::from_utf8(&buf[colon_idx + 1..colon_idx + 1 + length])
            .unwrap()
            .to_string();

        buf.advance(colon_idx + 1 + length);
        return Ok(Some(BeEncode::String(content)));
    }
}

#[cfg(test)]
mod tests {
    use std::collections::VecDeque;

    use futures::StreamExt;
    use tokio_util::codec::FramedRead;

    use crate::bencode::BeEncode;

    use super::StringDecoder;

    macro_rules! mock {
        ($($x:expr,)*) => {{
            let mut v = VecDeque::new();
            v.extend(vec![$($x),*]);
            Mock { calls: v }
        }};
    }

    struct Mock {
        calls: VecDeque<tokio::io::Result<Vec<u8>>>,
    }

    impl tokio::io::AsyncRead for Mock {
        fn poll_read(
            mut self: std::pin::Pin<&mut Self>,
            _cx: &mut std::task::Context<'_>,
            buf: &mut tokio::io::ReadBuf<'_>,
        ) -> futures::task::Poll<std::io::Result<()>> {
            use futures::task::Poll::{Pending, Ready};
            use tokio::io::ErrorKind::WouldBlock;

            match self.calls.pop_front() {
                Some(Ok(data)) => {
                    debug_assert!(buf.remaining() >= data.len());
                    buf.put_slice(&data);
                    Ready(Ok(()))
                }
                Some(Err(ref e)) if e.kind() == WouldBlock => Pending,
                Some(Err(e)) => Ready(Err(e)),
                None => Ready(Ok(())),
            }
        }
    }

    #[tokio::test]
    async fn test_decode_simple_string() {
        let mock = mock! {
            Ok(b"5:hello".to_vec()),
            Ok(b"5:".to_vec()),
            Ok(b"world".to_vec()),
        };

        let decoder = StringDecoder;
        let mut reader = FramedRead::new(mock, decoder);
        assert_eq!(
            reader.next().await.unwrap().expect("1"),
            BeEncode::String("hello".into())
        );
        assert_eq!(
            reader.next().await.unwrap().expect("2"),
            BeEncode::String("world".into())
        );
    }
}