Framing packets from a UDP stream in Tokio

user_24545 0
Viewed 124

Consider this small server that listens to a UDP stream:

#[tokio::main]
async fn main() {
    let sock = UdpSocket::bind(address).await.unwrap();
    let mut frame = UdpFramed::new(sock, PacketCodec::new());
    loop {
        match frame.next().await {
            Some(Ok((packet, _addr))) => {
                println!("got packet: {}", packet);
            }
            _ => {}
        };
    }
}

PacketCodec attempts to convert the incoming byte stream to Packets via a custom logic that searches for a well-known preamble, and then chunks each packet according to a predefined size:

pub struct PacketCodec {
    sync: bool,
}

impl PacketCodec {
    pub fn new() -> Self {
        Self { sync: false }
    }
}

impl Decoder for PacketCodec {
    type Item = Packet;
    type Error = io::Error;

    fn decode(&mut self, buf: &mut BytesMut) -> Result<Option<Packet>, io::Error> {

        if !self.sync {
            if buf.len() > PACKET_LEN * 2 {
                let srchstr = String::from_utf8(buf).unwrap();
                match srchstr.find(PREAMBLE) {
                    Some(n) => {
                        let _junk = buf.split_to(n);
                        self.sync = true;
                    }
                    None => {
                        let _junk = buf.split_to(PACKET_LEN);
                        self.sync = false;
                    }
                };
            }
        }

        if self.sync {
            if buf.len() > PACKET_LEN {
                let bs = buf.split_to(PACKET_LEN);
                let packet = Packet::new(&bs);
                return Ok(Some(packet));
            }
        }

        buf.reserve(PACKET_LEN * 2);
        Ok(None)
    }
}

This code runs just fine but hangs after one iteration, when sync has already happened, it is waiting for the next packet, but the buf size never increases.

Since the decoder returns Ok(None) when the bytes stream is smaller than the packet size, I would expect the Tokio runtime to send more bytes the next iteration, yet it doesn't.

How should I handle the incoming bytes stream properly?

rust-tokio rust udp
1 Answers

Kudos to @eggyal for digging in and opening issue 4888

Using that information, I'm currently using the following workaround by overriding the default decode_eof method and ensuring that it never returns an error even if the buffer is empty:

fn decode_eof(&mut self, buf: &mut BytesMut) -> Result<Option<Self::Item>, Self::Error> {
    match self.decode(buf)? {
        Some(frame) => Ok(Some(frame)),
        None => Ok(None),
    }
}

Further details on whether this is actually an upstream bug will probably be discussed in the Github issue.

user_245450
Related