diff --git a/src/net/udp/mod.rs b/src/net/udp/mod.rs index c49e91b74..d5590fa6c 100644 --- a/src/net/udp/mod.rs +++ b/src/net/udp/mod.rs @@ -74,6 +74,46 @@ impl UdpSocket { self.io.get_ref().local_addr() } + /// Connects the UDP socket setting the default destination for send() and + /// limiting packets that are read via recv from the address specified in addr. + pub fn connect(&self, addr: SocketAddr) -> io::Result<()> { + self.io.get_ref().connect(addr) + } + + /// Sends data on the socket to the address previously bound via connect(). + /// On success, returns the number of bytes written. + pub fn send(&self, buf: &[u8]) -> io::Result { + if let Async::NotReady = self.io.poll_write() { + return Err(io::ErrorKind::WouldBlock.into()) + } + match self.io.get_ref().send(buf) { + Ok(n) => Ok(n), + Err(e) => { + if e.kind() == io::ErrorKind::WouldBlock { + self.io.need_write(); + } + Err(e) + } + } + } + + /// Receives data from the socket previously bound with connect(). + /// On success, returns the number of bytes read. + pub fn recv(&self, buf: &mut [u8]) -> io::Result { + if let Async::NotReady = self.io.poll_read() { + return Err(io::ErrorKind::WouldBlock.into()) + } + match self.io.get_ref().recv(buf) { + Ok(n) => Ok(n), + Err(e) => { + if e.kind() == io::ErrorKind::WouldBlock { + self.io.need_read(); + } + Err(e) + } + } + } + /// Test whether this socket is ready to be read or not. /// /// If the socket is *not* readable then the current task is scheduled to diff --git a/tests/udp.rs b/tests/udp.rs index f22266d3a..fe14bc36a 100644 --- a/tests/udp.rs +++ b/tests/udp.rs @@ -16,52 +16,85 @@ macro_rules! t { }) } -#[test] -fn send_messages() { +fn send_messages(send: S, recv: R) { 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 a = t!(UdpSocket::bind(&([127, 0, 0, 1], 0).into(), &l.handle())); + let mut b = t!(UdpSocket::bind(&([127, 0, 0, 1], 0).into(), &l.handle())); let a_addr = t!(a.local_addr()); let b_addr = t!(b.local_addr()); { - let send = SendMessage::new(a, b_addr, b"1234"); - let recv = RecvMessage::new(b, a_addr, b"1234"); + let send = SendMessage::new(a, send.clone(), b_addr, b"1234"); + let recv = RecvMessage::new(b, recv.clone(), 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""); + let send = SendMessage::new(a, send, b_addr, b""); + let recv = RecvMessage::new(b, recv, a_addr, b""); t!(l.run(send.join(recv))); } } -struct SendMessage { +#[test] +fn send_to_and_recv_from() { + send_messages(SendTo {}, RecvFrom {}); +} + +#[test] +fn send_and_recv() { + send_messages(Send {}, Recv {}); +} + +trait SendFn { + fn send(&self, &UdpSocket, &[u8], &SocketAddr) -> Result; +} + +#[derive(Debug, Clone)] +struct SendTo {} + +impl SendFn for SendTo { + fn send(&self, socket: &UdpSocket, buf: &[u8], addr: &SocketAddr) -> Result { + socket.send_to(buf, addr) + } +} + +#[derive(Debug, Clone)] +struct Send {} + +impl SendFn for Send { + fn send(&self, socket: &UdpSocket, buf: &[u8], addr: &SocketAddr) -> Result { + socket.connect(*addr).expect("could not connect"); + socket.send(buf) + } +} + +struct SendMessage { socket: Option, + send: S, addr: SocketAddr, data: &'static [u8], } -impl SendMessage { - fn new(socket: UdpSocket, addr: SocketAddr, data: &'static [u8]) -> SendMessage { +impl SendMessage { + fn new(socket: UdpSocket, send: S, addr: SocketAddr, data: &'static [u8]) -> SendMessage { SendMessage { socket: Some(socket), - addr: addr, + send: send, + addr: addr, data: data, } } } -impl Future for SendMessage { +impl Future for SendMessage { type Item = UdpSocket; type Error = io::Error; fn poll(&mut self) -> Poll { - let n = try_nb!(self.socket.as_ref().unwrap() - .send_to(&self.data[..], &self.addr)); + let n = try_nb!(self.send.send(self.socket.as_ref().unwrap(), &self.data[..], &self.addr)); assert_eq!(n, self.data.len()); @@ -69,36 +102,62 @@ impl Future for SendMessage { } } -struct RecvMessage { - socket: Option, - addr: SocketAddr, - data: &'static [u8], +trait RecvFn { + fn recv(&self, &UdpSocket, &mut [u8], &SocketAddr) -> Result; } -impl RecvMessage { - fn new(socket: UdpSocket, expected_addr: SocketAddr, - expected_data: &'static [u8]) -> RecvMessage - { +#[derive(Debug, Clone)] +struct RecvFrom {} + +impl RecvFn for RecvFrom { + fn recv(&self, socket: &UdpSocket, buf: &mut [u8], + expected_addr: &SocketAddr) -> Result { + socket.recv_from(buf).map(|(s, addr)| { + assert_eq!(addr, *expected_addr); + s + }) + } +} + +#[derive(Debug, Clone)] +struct Recv {} + +impl RecvFn for Recv { + fn recv(&self, socket: &UdpSocket, buf: &mut [u8], _: &SocketAddr) -> Result { + socket.recv(buf) + } +} + +struct RecvMessage { + socket: Option, + recv: R, + expected_addr: SocketAddr, + expected_data: &'static [u8], +} + +impl RecvMessage { + fn new(socket: UdpSocket, recv: R, expected_addr: SocketAddr, + expected_data: &'static [u8]) -> RecvMessage { RecvMessage { socket: Some(socket), - addr: expected_addr, - data: expected_data, + recv: recv, + expected_addr: expected_addr, + expected_data: expected_data, } } } -impl Future for RecvMessage { +impl Future for RecvMessage { type Item = UdpSocket; type Error = io::Error; fn poll(&mut self) -> Poll { - let mut buf = vec![0u8; 10 + self.data.len() * 10]; - let (n, addr) = try_nb!(self.socket.as_ref().unwrap() - .recv_from(&mut buf[..])); + let mut buf = vec![0u8; 10 + self.expected_data.len() * 10]; + let n = try_nb!(self.recv.recv(&self.socket.as_ref().unwrap(), &mut buf[..], + &self.expected_addr)); - assert_eq!(n, self.data.len()); - assert_eq!(&buf[..self.data.len()], &self.data[..]); - assert_eq!(addr, self.addr); + assert_eq!(n, self.expected_data.len()); + assert_eq!(&buf[..self.expected_data.len()], &self.expected_data[..]); Ok(self.socket.take().unwrap().into()) } @@ -185,3 +244,4 @@ fn send_framed() { assert_eq!(received.0, Some(())); } } +