UDP 0-length datagrams

Resolves #248, preventing UdpFramed from sending 0-length datagrams.
Also, adds 0-length tests for UdpSocket.
This commit is contained in:
Henning Ottesen
2017-09-11 17:32:37 +02:00
parent 7b94cf307d
commit bda6ed9241
2 changed files with 161 additions and 31 deletions
+16 -8
View File
@@ -61,6 +61,7 @@ pub struct UdpFramed<C> {
rd: Vec<u8>, rd: Vec<u8>,
wr: Vec<u8>, wr: Vec<u8>,
out_addr: SocketAddr, out_addr: SocketAddr,
flushed: bool,
} }
impl<C: UdpCodec> Stream for UdpFramed<C> { impl<C: UdpCodec> Stream for UdpFramed<C> {
@@ -81,29 +82,35 @@ impl<C: UdpCodec> Sink for UdpFramed<C> {
type SinkError = io::Error; type SinkError = io::Error;
fn start_send(&mut self, item: C::Out) -> StartSend<C::Out, io::Error> { fn start_send(&mut self, item: C::Out) -> StartSend<C::Out, io::Error> {
if self.wr.len() > 0 { trace!("sending frame");
try!(self.poll_complete());
if self.wr.len() > 0 { if !self.flushed {
return Ok(AsyncSink::NotReady(item)); match try!(self.poll_complete()) {
Async::Ready(()) => {},
Async::NotReady => return Ok(AsyncSink::NotReady(item)),
} }
} }
self.out_addr = self.codec.encode(item, &mut self.wr); self.out_addr = self.codec.encode(item, &mut self.wr);
self.flushed = false;
trace!("frame encoded; length={}", self.wr.len());
Ok(AsyncSink::Ready) Ok(AsyncSink::Ready)
} }
fn poll_complete(&mut self) -> Poll<(), io::Error> { fn poll_complete(&mut self) -> Poll<(), io::Error> {
trace!("flushing framed transport"); if self.flushed {
if self.wr.is_empty() {
return Ok(Async::Ready(())) return Ok(Async::Ready(()))
} }
trace!("writing; remaining={}", self.wr.len()); trace!("flushing frame; length={}", self.wr.len());
let n = try_nb!(self.socket.send_to(&self.wr, &self.out_addr)); let n = try_nb!(self.socket.send_to(&self.wr, &self.out_addr));
trace!("written {}", n); trace!("written {}", n);
let wrote_all = n == self.wr.len(); let wrote_all = n == self.wr.len();
self.wr.clear(); self.wr.clear();
self.flushed = true;
if wrote_all { if wrote_all {
Ok(Async::Ready(())) Ok(Async::Ready(()))
} else { } else {
@@ -125,6 +132,7 @@ pub fn new<C: UdpCodec>(socket: UdpSocket, codec: C) -> UdpFramed<C> {
out_addr: SocketAddr::V4(SocketAddrV4::new(Ipv4Addr::new(0, 0, 0, 0), 0)), out_addr: SocketAddr::V4(SocketAddrV4::new(Ipv4Addr::new(0, 0, 0, 0), 0)),
rd: vec![0; 64 * 1024], rd: vec![0; 64 * 1024],
wr: Vec::with_capacity(8 * 1024), wr: Vec::with_capacity(8 * 1024),
flushed: true,
} }
} }
+145 -23
View File
@@ -5,8 +5,8 @@ extern crate tokio_core;
use std::io; use std::io;
use std::net::SocketAddr; use std::net::SocketAddr;
use futures::{Future, Poll}; use futures::{Future, Poll, Stream, Sink};
use tokio_core::net::UdpSocket; use tokio_core::net::{UdpSocket, UdpCodec};
use tokio_core::reactor::Core; use tokio_core::reactor::Core;
macro_rules! t { macro_rules! t {
@@ -19,47 +19,169 @@ macro_rules! t {
#[test] #[test]
fn send_messages() { fn send_messages() {
let mut l = t!(Core::new()); let mut l = t!(Core::new());
let a = t!(UdpSocket::bind(&t!("127.0.0.1:0".parse()), &l.handle())); let mut a = t!(UdpSocket::bind(&t!("127.0.0.1:0".parse()), &l.handle()));
let b = t!(UdpSocket::bind(&t!("127.0.0.1:0".parse()), &l.handle())); let mut b = t!(UdpSocket::bind(&t!("127.0.0.1:0".parse()), &l.handle()));
let a_addr = t!(a.local_addr()); let a_addr = t!(a.local_addr());
let b_addr = t!(b.local_addr()); let b_addr = t!(b.local_addr());
let send = SendMessage { socket: a, addr: b_addr }; {
let recv = RecvMessage { socket: b, expected_addr: a_addr }; let send = SendMessage::new(a, b_addr, b"1234");
t!(l.run(send.join(recv))); let recv = RecvMessage::new(b, a_addr, b"1234");
let (sendt, received) = t!(l.run(send.join(recv)));
a = sendt;
b = received;
}
{
let send = SendMessage::new(a, b_addr, b"");
let recv = RecvMessage::new(b, a_addr, b"");
t!(l.run(send.join(recv)));
}
} }
struct SendMessage { struct SendMessage {
socket: UdpSocket, socket: Option<UdpSocket>,
addr: SocketAddr, addr: SocketAddr,
data: &'static [u8],
}
impl SendMessage {
fn new(socket: UdpSocket, addr: SocketAddr, data: &'static [u8]) -> SendMessage {
SendMessage {
socket: Some(socket),
addr: addr,
data: data,
}
}
} }
impl Future for SendMessage { impl Future for SendMessage {
type Item = (); type Item = UdpSocket;
type Error = io::Error; type Error = io::Error;
fn poll(&mut self) -> Poll<(), io::Error> { fn poll(&mut self) -> Poll<UdpSocket, io::Error> {
let n = try_nb!(self.socket.send_to(b"1234", &self.addr)); let n = try_nb!(self.socket.as_ref().unwrap()
assert_eq!(n, 4); .send_to(&self.data[..], &self.addr));
Ok(().into())
assert_eq!(n, self.data.len());
Ok(self.socket.take().unwrap().into())
} }
} }
struct RecvMessage { struct RecvMessage {
socket: UdpSocket, socket: Option<UdpSocket>,
expected_addr: SocketAddr, addr: SocketAddr,
data: &'static [u8],
}
impl RecvMessage {
fn new(socket: UdpSocket, expected_addr: SocketAddr,
expected_data: &'static [u8]) -> RecvMessage
{
RecvMessage {
socket: Some(socket),
addr: expected_addr,
data: expected_data,
}
}
} }
impl Future for RecvMessage { impl Future for RecvMessage {
type Item = (); type Item = UdpSocket;
type Error = io::Error; type Error = io::Error;
fn poll(&mut self) -> Poll<(), io::Error> { fn poll(&mut self) -> Poll<UdpSocket, io::Error> {
let mut buf = [0; 32]; let mut buf = vec![0u8; 10 + self.data.len() * 10];
let (n, addr) = try_nb!(self.socket.recv_from(&mut buf)); let (n, addr) = try_nb!(self.socket.as_ref().unwrap()
assert_eq!(n, 4); .recv_from(&mut buf[..]));
assert_eq!(&buf[..4], b"1234");
assert_eq!(addr, self.expected_addr); assert_eq!(n, self.data.len());
Ok(().into()) assert_eq!(&buf[..self.data.len()], &self.data[..]);
assert_eq!(addr, self.addr);
Ok(self.socket.take().unwrap().into())
}
}
#[test]
fn send_dgrams() {
let mut l = t!(Core::new());
let mut a = t!(UdpSocket::bind(&t!("127.0.0.1:0".parse()), &l.handle()));
let mut b = t!(UdpSocket::bind(&t!("127.0.0.1:0".parse()), &l.handle()));
let mut buf = [0u8; 50];
let b_addr = t!(b.local_addr());
{
let send = a.send_dgram(&b"4321"[..], b_addr);
let recv = b.recv_dgram(&mut buf[..]);
let (sendt, received) = t!(l.run(send.join(recv)));
assert_eq!(received.2, 4);
assert_eq!(&received.1[..4], b"4321");
a = sendt.0;
b = received.0;
}
{
let send = a.send_dgram(&b""[..], b_addr);
let recv = b.recv_dgram(&mut buf[..]);
let received = t!(l.run(send.join(recv))).1;
assert_eq!(received.2, 0);
}
}
#[derive(Debug, Clone)]
struct Codec {
data: &'static [u8],
from: SocketAddr,
to: SocketAddr,
}
impl UdpCodec for Codec {
type In = ();
type Out = &'static [u8];
fn decode(&mut self, src: &SocketAddr, buf: &[u8]) -> io::Result<Self::In> {
assert_eq!(src, &self.from);
assert_eq!(buf, self.data);
Ok(())
}
fn encode(&mut self, msg: Self::Out, buf: &mut Vec<u8>) -> SocketAddr {
assert_eq!(msg, self.data);
buf.extend_from_slice(msg);
self.to
}
}
#[test]
fn send_framed() {
let mut l = t!(Core::new());
let mut a_soc = t!(UdpSocket::bind(&t!("127.0.0.1:0".parse()), &l.handle()));
let mut b_soc = t!(UdpSocket::bind(&t!("127.0.0.1:0".parse()), &l.handle()));
let a_addr = t!(a_soc.local_addr());
let b_addr = t!(b_soc.local_addr());
{
let a = a_soc.framed(Codec { data: &b"4567"[..], from: a_addr, to: b_addr});
let b = b_soc.framed(Codec { data: &b"4567"[..], from: a_addr, to: b_addr});
let send = a.send(&b"4567"[..]);
let recv = b.into_future().map_err(|e| e.0);
let (sendt, received) = t!(l.run(send.join(recv)));
assert_eq!(received.0, Some(()));
a_soc = sendt.into_inner();
b_soc = received.1.into_inner();
}
{
let a = a_soc.framed(Codec { data: &b""[..], from: a_addr, to: b_addr});
let b = b_soc.framed(Codec { data: &b""[..], from: a_addr, to: b_addr});
let send = a.send(&b""[..]);
let recv = b.into_future().map_err(|e| e.0);
let received = t!(l.run(send.join(recv))).1;
assert_eq!(received.0, Some(()));
} }
} }