From 71d8672aab2b6c4712942783920e01db578eac9c Mon Sep 17 00:00:00 2001 From: Rick Richardson Date: Sun, 20 Nov 2016 09:08:03 -0800 Subject: [PATCH] implemented moste of udp frames test --- src/io/mod.rs | 2 + src/io/udp_frame.rs | 128 ++++++++++++++++++--------------------- src/net/udp.rs | 18 +++--- tests/udp-line-frames.rs | 74 ++++++++++++++++++++++ 4 files changed, 144 insertions(+), 78 deletions(-) create mode 100644 tests/udp-line-frames.rs diff --git a/src/io/mod.rs b/src/io/mod.rs index 15aba3b74..694174892 100644 --- a/src/io/mod.rs +++ b/src/io/mod.rs @@ -33,6 +33,7 @@ macro_rules! try_nb { mod copy; mod frame; +mod udp_frame; mod flush; mod read_exact; mod read_to_end; @@ -43,6 +44,7 @@ mod window; mod write_all; pub use self::copy::{copy, Copy}; pub use self::frame::{EasyBuf, EasyBufMut, FramedRead, FramedWrite, Framed, Codec}; +pub use self::udp_frame::{FramedUdp, framed_udp, FramedUdpRead, FramedUdpWrite, CodecUdp}; pub use self::flush::{flush, Flush}; pub use self::read_exact::{read_exact, ReadExact}; pub use self::read_to_end::{read_to_end, ReadToEnd}; diff --git a/src/io/udp_frame.rs b/src/io/udp_frame.rs index 9c1eeb958..a1b64db09 100644 --- a/src/io/udp_frame.rs +++ b/src/io/udp_frame.rs @@ -1,12 +1,9 @@ use std::io; -use std::ops::{Deref, DerefMut}; -use std::sync::Arc; -use net::udp::UdpSocket +use std::net::SocketAddr; +use net::UdpSocket; use futures::{Async, Poll, Stream, Sink, StartSend, AsyncSink}; use futures::sync::BiLock; -use io::Io; - /// Encoding of frames via buffers. /// /// This trait is used when constructing an instance of `FramedUdp`. It provides @@ -20,10 +17,13 @@ use io::Io; /// or encoding, which is particularly useful for streaming parsers. In many /// cases, though, this type will simply be a unit struct (e.g. `struct /// HttpCodec`). -pub trait EncodeUdp { +pub trait CodecUdp { /// The type of frames to be encoded. type Out; + + /// The type of decoded frames. + type In; /// Encodes a frame into the buffer provided. @@ -32,87 +32,64 @@ pub trait EncodeUdp { /// The `buf` provided is an internal buffer of the `Framed` instance and /// will be written out when possible. /// - /// The codec also determines the destination to which the buffer should + /// The encode method also determines the destination to which the buffer should /// be directed, which will be returned as a SocketAddr; fn encode(&mut self, msg: Self::Out, buf: &mut Vec) -> SocketAddr; -} - -/// Decoding of frames via buffers. -/// -/// This trait is used when constructing an instance of `FramedUdp`. It provides -/// one type: `In` for decoding incoming frames from a Datagram -/// -/// Because UDP is a connectionless protocol, the decode method will also be -/// supplied with a SocketAddr of the remote host which sent the datagram -/// -/// The trait itself is implemented on a type that can track state for decoding -/// or encoding, which is particularly useful for streaming parsers. In many -/// cases, though, this type will simply be a unit struct (e.g. `struct -/// HttpCodec`). -pub trait DecodeUdp { - /// The type of decoded frames. - type In; - + /// Attempts to decode a frame from the provided buffer of bytes. /// /// This method is called by `FramedUdp` on a single datagram which has been /// read from a socket. /// - /// It is required that the Decoder empty the read buffer in every call to + /// It is required that the decode method empty the read buffer in every call to /// decode, as the next poll_read that occurs will write the next datagram /// into the buffer, without regard for what is already there. /// /// If the bytes look valid, but a frame isn't fully available yet, then /// `Ok(None)` is returned. This indicates to the `Framed` instance that /// it needs to read some more bytes before calling this method again. - /// In such a case, it is the decoder's responsibility to copy the data + /// In such a case, it is decode's responsibility to copy the data /// into their own internal buffer for future use. /// /// Finally, if the bytes in the buffer are malformed then an error is /// returned indicating why. This informs `Framed` that the stream is now /// corrupt and should be terminated. /// - /// When dealing with connectionless streams, there will likely be some sort - /// of state machine. fn decode(&mut self, src: &SocketAddr, buf: &mut Vec) -> Result, io::Error>; } /// A unified `Stream` and `Sink` interface to an underlying `Io` object, using -/// the `Encode` and `Decode` traits to encode and decode frames. +/// the `CodecUdp` trait to encode and decode frames. /// /// You can acquire a `Framed` instance by using the `Io::framed` adapter. -pub struct FramedUdp { +pub struct FramedUdp { socket: UdpSocket, - encoder: E, - decoder: D, + codec: C, out_addr : Option, rd: Vec, wr: Vec, } -impl Stream for Framed { - type Item = D::In; +impl Stream for FramedUdp { + type Item = C::In; type Error = io::Error; fn poll(&mut self) -> Poll, io::Error> { loop { - let before = self.rd.len(); - let ret = self.socket.recv_from(self.rd.mut_bytes(), &mut inaddr); + let ret = self.socket.recv_from(self.rd.as_mut_slice()); match ret { Ok((n, addr)) => { trace!("read {} bytes", n); trace!("attempting to decode a frame"); - if let Some(frame) = try!(self.decoder.decode(&addr, &mut self.rd)) { + if let Some(frame) = try!(self.codec.decode(&addr, &mut self.rd)) { trace!("frame decoded from buffer"); self.rd.clear(); return Ok(Async::Ready(Some(frame))); } } Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => { - if self.rd.len() == before { - return Ok(Async::NotReady) - } + return Ok(Async::NotReady) } Err(e) => return Err(e), } @@ -120,11 +97,11 @@ impl Stream for Framed { } } -impl Sink for Framed { - type SinkItem = E::Out; +impl Sink for FramedUdp { + type SinkItem = C::Out; type SinkError = io::Error; - fn start_send(&mut self, item: C::Out) -> StartSend { + fn start_send(&mut self, item: C::Out) -> StartSend { if self.wr.len() > 0 { try!(self.poll_complete()); if self.wr.len() > 0 { @@ -140,9 +117,9 @@ impl Sink for Framed { trace!("flushing framed transport"); while !self.wr.is_empty() { - if let Some(outaddr) = self.out_addr.ref() { + if let Some(outaddr) = self.out_addr { trace!("writing; remaining={}", self.wr.len()); - let n = try_nb!(self.socket.send_to(&self.wr, outaddr)); + let n = try_nb!(self.socket.send_to(&self.wr, &outaddr)); self.wr.clear(); self.out_addr = None; if n != self.wr.len() { @@ -160,22 +137,37 @@ impl Sink for Framed { } } -pub fn framed_udp(socket : UdpSocket, decoder : D, encoder : E) -> Framed { - Framed { - socket: socket, - encoder: encoder, - decoder: decoder, - rd: Vec::with_capacity(64 * 1024), - wr: Vec::with_capacity(64 * 1024) - } +/// Helper function that Creates a new FramedUdp object. It moves the supplied socket, codec +/// into the resulting FramedUdp +pub fn framed_udp(socket : UdpSocket, codec : C) -> FramedUdp { + FramedUdp::new( + socket, + codec, + Vec::with_capacity(64 * 1024), + Vec::with_capacity(64 * 1024) + ) } -impl FramedUdp { +impl FramedUdp { + + /// Creates a new FramedUdp object. It moves the supplied socket, codec + /// supplied vecs. + pub fn new(sock : UdpSocket, codec : C, rd_buf : Vec, wr_buf : Vec) -> FramedUdp { + FramedUdp { + socket: sock, + codec : codec, + out_addr: None, + rd: rd_buf, + wr: wr_buf + } + } + + /// Splits this `Stream + Sink` object into separate `Stream` and `Sink` /// objects, which can be useful when you want to split ownership between /// tasks, or allow direct interaction between the two objects (e.g. via /// `Sink::send_all`). - pub fn split(self) -> (FramedRead, FramedWrite) { + pub fn split(self) -> (FramedUdpRead, FramedUdpWrite) { let (a, b) = BiLock::new(self); let read = FramedUdpRead { framed: a }; let write = FramedUdpWrite { framed: b }; @@ -210,17 +202,17 @@ impl FramedUdp { self.socket } } -/// A `Stream` interface to an underlying `Io` object, using the `Decode` trait +/// A `Stream` interface to an underlying `Io` object, using the `CodecUdp` trait /// to decode frames. -pub struct FramedRead { - framed: BiLock>, +pub struct FramedUdpRead { + framed: BiLock>, } -impl Stream for FramedRead { - type Item = D::In; +impl Stream for FramedUdpRead { + type Item = C::In; type Error = io::Error; - fn poll(&mut self) -> Poll, io::Error> { + fn poll(&mut self) -> Poll, io::Error> { if let Async::Ready(mut guard) = self.framed.poll_lock() { guard.poll() } else { @@ -229,17 +221,17 @@ impl Stream for FramedRead { } } -/// A `Sink` interface to an underlying `Io` object, using the `Encode` trait +/// A `Sink` interface to an underlying `Io` object, using the `CodecUdp` trait /// to encode frames. -pub struct FramedWrite { - framed: BiLock>, +pub struct FramedUdpWrite { + framed: BiLock>, } -impl Sink for FramedWrite { - type SinkItem = E::Out; +impl Sink for FramedUdpWrite { + type SinkItem = C::Out; type SinkError = io::Error; - fn start_send(&mut self, item: E::Out) -> StartSend { + fn start_send(&mut self, item: C::Out) -> StartSend { if let Async::Ready(mut guard) = self.framed.poll_lock() { guard.start_send(item) } else { diff --git a/src/net/udp.rs b/src/net/udp.rs index 2a2e0a03a..34de052c3 100644 --- a/src/net/udp.rs +++ b/src/net/udp.rs @@ -1,7 +1,7 @@ use std::io; use std::net::{self, SocketAddr, Ipv4Addr, Ipv6Addr}; use std::fmt; - +use io::FramedUdp; use futures::Async; use mio; @@ -45,15 +45,13 @@ impl UdpSocket { /// Creates a FramedUdp object, which leverages a supplied `EncodeUdp` /// and `DecodeUdp` to implement `Stream` and `Sink` /// This moves the socket into the newly created FramedUdp object - pub fn framed(self, decoder : D, encoder : E) -> Framed { - FramedUdp { - socket: self, - encoder: encoder, - decoder: decoder, - is_readable: false, - rd: Vec::with_capacity(64 * 1024); - wr: Vec::with_capacity(64 * 1024), - } + pub fn framed(self, codec : C) -> FramedUdp { + FramedUdp::new( + self, + codec, + Vec::with_capacity(64 * 1024), + Vec::with_capacity(64 * 1024) + ) } /// Returns the local address that this stream is bound to. diff --git a/tests/udp-line-frames.rs b/tests/udp-line-frames.rs new file mode 100644 index 000000000..90e8968d6 --- /dev/null +++ b/tests/udp-line-frames.rs @@ -0,0 +1,74 @@ +extern crate tokio_core; +extern crate env_logger; +extern crate futures; + +use std::io; +use std::net::{SocketAddr}; +use futures::{future, Future, Stream, Sink, IntoFuture}; +use tokio_core::io::{write_all, read, FramedUdp, CodecUdp, Io}; +use tokio_core::net::{UdpSocket}; +use tokio_core::reactor::{Core, Timeout}; +use std::time::Duration; +use std::str; + +pub struct LineCodec { + addr : Option +} + +impl CodecUdp for LineCodec { + type In = Vec; + type Out = Vec; + + fn decode(&mut self, addr : &SocketAddr, buf: &mut Vec) -> Result, io::Error> { + self.addr = Some(*addr); + match buf.as_slice().iter().position(|&b| b == b'\n') { + Some(i) => Ok(Some(buf[.. i + 1].into())), + None => Ok(None), + } + } + + fn encode(&mut self, item: Vec, into: &mut Vec) -> SocketAddr { + into.extend_from_slice(item.as_slice()); + into.push('\n' as u8); + + self.addr.unwrap() + } +} + +#[test] +fn echo() { + drop(env_logger::init()); + + let mut core = Core::new().unwrap(); + let handle = core.handle(); + + let srvcodec = LineCodec { addr : None }; + let clicodec = LineCodec { addr : None }; + + let srvaddr : SocketAddr = "127.0.0.1:31999".parse().unwrap(); + let clientaddr : SocketAddr = "127.0.0.1:32000".parse().unwrap(); + + let server = UdpSocket::bind(&srvaddr, &handle).unwrap(); + let client = UdpSocket::bind(&clientaddr, &handle).unwrap(); + + let job = client.send_to(b"PING", &srvaddr); + let _ = core.run(job.into_future()).unwrap(); + + let (srvstream, srvsink) = server.framed(srvcodec).split(); + let srvloop = srvstream.for_each(move |buf| { + println!("{}", str::from_utf8(buf.as_slice()).unwrap()); + srvsink.send(b"PONG".to_vec()).map(|_| ()).wait() + }); + + let (clistream, clisink) = client.framed(clicodec).split(); + let cliloop = clistream.for_each(move |buf| { + println!("{}", str::from_utf8(buf.as_slice()).unwrap()); + clisink.send(b"PING".to_vec()).map(|_| ()).wait() + }); + + let timeout = Timeout::new(Duration::from_millis(500), &handle).unwrap(); + + let wait = future::select_all(vec![timeout.boxed(), srvloop.boxed(), cliloop.boxed()]); + core.run(wait); + +}