tokio-rs/tokio
1
0
Fork 0
mirror of https://github.com/tokio-rs/tokio.git synced 2026-08-28 00:00:11 +02:00
Code Issues Packages Projects Releases Wiki Activity

Touch up codes for UDP

Browse Source 
* Move to `std::net` as it's all purely UDP related
* Rename to `UdpCodec` and `UdpFramed` to give a consistent `Udp` prefix
* Add `RecvDgram`, rename `SendDGram` to `SendDgram`
* Touch up some style here and there
This commit is contained in:
Alex Crichton
2016-11-22 11:48:09 -08:00
parent 0d10b0e05a
commit 9b62ade962
6 changed files with 293 additions and 322 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/io/mod.rs
-2
View File
@@ -33,7 +33,6 @@ macro_rules! try_nb {
mod copy;
mod frame;
mod udp_frame;
mod flush;
mod read_exact;
mod read_to_end;
@@ -44,7 +43,6 @@ mod window;
mod write_all;
pub use self::copy::{copy, Copy};
pub use self::frame::{EasyBuf, EasyBufMut, Framed, Codec};
pub use self::udp_frame::{FramedUdp, framed_udp, FramedUdpRead, FramedUdpWrite, CodecUdp, VecDGramCodec};
pub use self::flush::{flush, Flush};
pub use self::read_exact::{read_exact, ReadExact};
pub use self::read_to_end::{read_to_end, ReadToEnd};
src/io/udp_frame.rs
-262
View File
@@ -1,262 +0,0 @@
use std::io;
use std::net::SocketAddr;
use net::UdpSocket;
use futures::{Async, Poll, Stream, Sink, StartSend, AsyncSink};
use futures::sync::BiLock;
/// Encoding of frames via buffers.
///
/// This trait is used when constructing an instance of `FramedUdp`. It provides
/// one type: `Out` for encoding outgoing frames according to a protocol.
///
/// Because UDP is a connectionless protocol, the encode method will also be
/// responsible for determining the remote host to which the datagram should be
/// sent
///
/// 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 CodecUdp {
/// The type of frames to be encoded.
type Out;
/// The type of decoded frames.
type In;
/// Encodes a frame into the buffer provided.
///
/// This method will encode `msg` into the byte buffer provided by `buf`.
/// The `buf` provided is an internal buffer of the `Framed` instance and
/// will be written out when possible.
///
/// 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<u8>) -> SocketAddr;
/// 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 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.
///
/// 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.
///
fn decode(&mut self, src: &SocketAddr, buf: &[u8]) -> Result<Self::In, io::Error>;
}
/// A unified `Stream` and `Sink` interface to an underlying `Io` object, using
/// the `CodecUdp` trait to encode and decode frames.
///
/// You can acquire a `Framed` instance by using the `Io::framed` adapter.
pub struct FramedUdp<C> {
socket: UdpSocket,
codec: C,
out_addr : Option<SocketAddr>,
rd: Vec<u8>,
wr: Vec<u8>,
}
impl<C : CodecUdp> Stream for FramedUdp<C> {
type Item = C::In;
type Error = io::Error;
fn poll(&mut self) -> Poll<Option<C::In>, io::Error> {
loop {
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");
let frame = try!(self.codec.decode(&addr, & self.rd[.. n]));
trace!("frame decoded from buffer");
return Ok(Async::Ready(Some(frame)));
}
Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => {
return Ok(Async::NotReady)
}
Err(e) => return Err(e),
}
}
}
}
impl<C : CodecUdp> Sink for FramedUdp<C> {
type SinkItem = C::Out;
type SinkError = io::Error;
fn start_send(&mut self, item: C::Out) -> StartSend<C::Out, io::Error> {
if self.wr.len() > 0 {
try!(self.poll_complete());
if self.wr.len() > 0 {
return Ok(AsyncSink::NotReady(item));
}
}
self.out_addr = Some(self.codec.encode(&item, &mut self.wr));
Ok(AsyncSink::Ready)
}
fn poll_complete(&mut self) -> Poll<(), io::Error> {
trace!("flushing framed transport");
if !self.wr.is_empty() {
if let Some(outaddr) = self.out_addr {
let remaining = self.wr.len();
trace!("writing; remaining={}", remaining);
let n = try_nb!(self.socket.send_to(self.wr.as_slice(), &outaddr));
trace!("written {}", n);
self.wr.clear();
self.out_addr = None;
if n != remaining {
return Err(io::Error::new(io::ErrorKind::Other,
"failed to write entire datagram to socket"));
}
}
else {
panic!("outbound stream in invalid state: out_addr is not known");
}
}
return Ok(Async::Ready(()));
}
}
/// Helper function that Creates a new FramedUdp object. It moves the supplied socket, codec
/// into the resulting FramedUdp
pub fn framed_udp<C>(socket : UdpSocket, codec : C) -> FramedUdp<C> {
FramedUdp::new(
socket,
codec,
vec![0; 64 * 1024],
Vec::with_capacity(64 * 1024)
)
}
impl<C> FramedUdp<C> {
/// Creates a new FramedUdp object. It moves the supplied socket, codec
/// supplied vecs.
fn new(sock : UdpSocket, codec : C, rd_buf : Vec<u8>, wr_buf : Vec<u8>) -> FramedUdp<C> {
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) -> (FramedUdpRead<C>, FramedUdpWrite<C>) {
let (a, b) = BiLock::new(self);
let read = FramedUdpRead { framed: a };
let write = FramedUdpWrite { framed: b };
(read, write)
}
/// Returns a reference to the underlying I/O stream wrapped by `Framed`.
///
/// Note that care should be taken to not tamper with the underlying stream
/// of data coming in as it may corrupt the stream of frames otherwise being
/// worked with.
pub fn get_ref(&self) -> &UdpSocket {
&self.socket
}
/// Returns a mutable reference to the underlying I/O stream wrapped by
/// `Framed`.
///
/// Note that care should be taken to not tamper with the underlying stream
/// of data coming in as it may corrupt the stream of frames otherwise being
/// worked with.
pub fn get_mut(&mut self) -> &mut UdpSocket {
&mut self.socket
}
/// Consumes the `Framed`, returning its underlying I/O stream.
///
/// Note that care should be taken to not tamper with the underlying stream
/// of data coming in as it may corrupt the stream of frames otherwise being
/// worked with.
pub fn into_inner(self) -> UdpSocket {
self.socket
}
}
/// A `Stream` interface to an underlying `Io` object, using the `CodecUdp` trait
/// to decode frames.
pub struct FramedUdpRead<C> {
framed: BiLock<FramedUdp<C>>,
}
impl<C : CodecUdp> Stream for FramedUdpRead<C> {
type Item = C::In;
type Error = io::Error;
fn poll(&mut self) -> Poll<Option<C::In>, io::Error> {
if let Async::Ready(mut guard) = self.framed.poll_lock() {
guard.poll()
} else {
Ok(Async::NotReady)
}
}
}
/// A `Sink` interface to an underlying `Io` object, using the `CodecUdp` trait
/// to encode frames.
pub struct FramedUdpWrite<C> {
framed: BiLock<FramedUdp<C>>,
}
impl<C : CodecUdp> Sink for FramedUdpWrite<C> {
type SinkItem = C::Out;
type SinkError = io::Error;
fn start_send(&mut self, item: C::Out) -> StartSend<C::Out, io::Error> {
if let Async::Ready(mut guard) = self.framed.poll_lock() {
guard.start_send(item)
} else {
Ok(AsyncSink::NotReady(item))
}
}
fn poll_complete(&mut self) -> Poll<(), io::Error> {
if let Async::Ready(mut guard) = self.framed.poll_lock() {
guard.poll_complete()
} else {
Ok(Async::NotReady)
}
}
}
/// Default implementation of a DGram "parser"
/// This receives and produces a tuple of ('SocketAddr', `Vec`<u8>)
pub struct VecDGramCodec;
impl CodecUdp for VecDGramCodec {
type In = (SocketAddr, Vec<u8>);
type Out = (SocketAddr, Vec<u8>);
fn decode(&mut self, addr : &SocketAddr, buf: &[u8]) -> Result<Self::In, io::Error> {
Ok((*addr, buf.into()))
}
fn encode(&mut self, item: &Self::Out, into: &mut Vec<u8>) -> SocketAddr {
into.extend_from_slice(item.1.as_slice());
into.push('\n' as u8);
item.0
}
}