tokio-rs/tokio
1
0
Fork 0
mirror of https://github.com/tokio-rs/tokio.git synced 2026-08-26 00:00:16 +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/net/mod.rs
+1 -1
View File
@@ -8,4 +8,4 @@ mod udp;
pub use self::tcp::{TcpStream, TcpStreamNew};
pub use self::tcp::{TcpListener, Incoming};
pub use self::udp::UdpSocket;
pub use self::udp::{UdpSocket, UdpCodec, UdpFramed};
src/net/udp/frame.rs
+154
View File
@@ -0,0 +1,154 @@
use std::io;
use std::net::{SocketAddr, Ipv4Addr, SocketAddrV4};
use futures::{Async, Poll, Stream, Sink, StartSend, AsyncSink};
use net::UdpSocket;
/// Encoding of frames via buffers.
///
/// This trait is used when constructing an instance of `UdpFramed` and provides
/// the `In` and `Out` types which are decoded and encoded from the socket,
/// respectively.
///
/// Because UDP is a connectionless protocol, the `decode` method receives the
/// address where data came from and the `encode` method is also 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 UdpCodec {
/// The type of decoded frames.
type In;
/// The type of frames to be encoded.
type Out;
/// Attempts to decode a frame from the provided buffer of bytes.
///
/// This method is called by `UdpFramed` on a single datagram which has been
/// read from a socket. The `buf` argument contains the data that was
/// received from the remote address, and `src` is the address the data came
/// from. Note that typically this method should require the entire contents
/// of `buf` to be valid or otherwise return an error with trailing data.
///
/// 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]) -> io::Result<Self::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;
}
/// A unified `Stream` and `Sink` interface to an underlying `UdpSocket`, using
/// the `UdpCodec` trait to encode and decode frames.
///
/// You can acquire a `UdpFramed` instance by using the `UdpSocket::framed`
/// adapter.
pub struct UdpFramed<C> {
socket: UdpSocket,
codec: C,
rd: Vec<u8>,
wr: Vec<u8>,
out_addr: SocketAddr,
}
impl<C: UdpCodec> Stream for UdpFramed<C> {
type Item = C::In;
type Error = io::Error;
fn poll(&mut self) -> Poll<Option<C::In>, io::Error> {
let (n, addr) = try_nb!(self.socket.recv_from(&mut self.rd));
trace!("received {} bytes, decoding", n);
let frame = try!(self.codec.decode(&addr, &self.rd[..n]));
trace!("frame decoded from buffer");
Ok(Async::Ready(Some(frame)))
}
}
impl<C: UdpCodec> Sink for UdpFramed<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 = 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() {
return Ok(Async::Ready(()))
}
trace!("writing; remaining={}", self.wr.len());
let n = try_nb!(self.socket.send_to(&self.wr, &self.out_addr));
trace!("written {}", n);
let wrote_all = n == self.wr.len();
self.wr.clear();
if wrote_all {
Ok(Async::Ready(()))
} else {
Err(io::Error::new(io::ErrorKind::Other,
"failed to write entire datagram to socket"))
}
}
}
pub fn new<C: UdpCodec>(socket: UdpSocket, codec: C) -> UdpFramed<C> {
UdpFramed {
socket: socket,
codec: codec,
out_addr: SocketAddr::V4(SocketAddrV4::new(Ipv4Addr::new(0, 0, 0, 0), 0)),
rd: vec![0; 64 * 1024],
wr: Vec::with_capacity(8 * 1024),
}
}
impl<C> UdpFramed<C> {
/// 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
}
}
src/net/udp.rs → src/net/udp/mod.rs
+123 -41
View File
@@ -1,10 +1,10 @@
use std::io;
use std::mem;
use std::net::{self, SocketAddr, Ipv4Addr, Ipv6Addr};
use std::fmt;
use io::{FramedUdp, framed_udp};
use futures::{Async, Future, Poll};
use mio;
use std::mem;
use reactor::{Handle, PollEvented};
@@ -13,6 +13,9 @@ pub struct UdpSocket {
io: PollEvented<mio::udp::UdpSocket>,
}
mod frame;
pub use self::frame::{UdpFramed, UdpCodec};
impl UdpSocket {
/// Create a new UDP socket bound to the specified address.
///
@@ -43,11 +46,27 @@ impl UdpSocket {
UdpSocket::new(udp, handle)
}
/// 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<C>(self, codec : C) -> FramedUdp<C> {
framed_udp(self, codec)
/// Provides a `Stream` and `Sink` interface for reading and writing to this
/// `UdpSocket` object, using the provided `UdpCodec` to read and write the
/// raw data.
///
/// Raw UDP sockets work with datagrams, but higher-level code usually
/// wants to batch these into meaningful chunks, called "frames". This
/// method layers framing on top of this socket by using the `UdpCodec`
/// trait to handle encoding and decoding of messages frames. Note that
/// the incoming and outgoing frame types may be distinct.
///
/// This function returns a *single* object that is both `Stream` and
/// `Sink`; grouping this into a single object is often useful for layering
/// things which require both read and write access to the underlying
/// object.
///
/// If you want to work more directly with the streams and sink, consider
/// calling `split` on the `UdpFramed` returned by this method, which will
/// break them into separate objects, allowing them to interact more
/// easily.
pub fn framed<C: UdpCodec>(self, codec: C) -> UdpFramed<C> {
frame::new(self, codec)
}
/// Returns the local address that this stream is bound to.
@@ -93,27 +112,27 @@ impl UdpSocket {
Err(e) => Err(e),
}
}
/// Creates a future that will write the entire contents of the buffer `buf` to
/// the stream `a` provided.
/// Creates a future that will write the entire contents of the buffer
/// `buf` provided as a datagram to this socket.
///
/// The returned future will return after data has been written to the outbound
/// socket.
/// The future will resolve to the stream as well as the buffer (for reuse if
/// needed).
/// The returned future will return after data has been written to the
/// outbound socket. The future will resolve to the stream as well as the
/// buffer (for reuse if needed).
///
/// Any error which happens during writing will cause both the stream and the
/// buffer to get destroyed.
/// Any error which happens during writing will cause both the stream and
/// the buffer to get destroyed. Note that failure to write the entire
/// buffer is considered an error for the purposes of sending a datagram.
///
/// The `buf` parameter here only requires the `AsRef<[u8]>` trait, which should
/// be broadly applicable to accepting data which can be converted to a slice.
/// The `Window` struct is also available in this crate to provide a different
/// window into a slice if necessary.
pub fn send_dgram<T>(self, buf: T, addr : SocketAddr) -> SendDGram<T>
/// The `buf` parameter here only requires the `AsRef<[u8]>` trait, which
/// should be broadly applicable to accepting data which can be converted
/// to a slice. The `Window` struct is also available in this crate to
/// provide a different window into a slice if necessary.
pub fn send_dgram<T>(self, buf: T, addr: SocketAddr) -> SendDgram<T>
where T: AsRef<[u8]>,
{
SendDGram {
state: UdpState::Writing {
SendDgram {
state: SendState::Writing {
sock: self,
addr: addr,
buf: buf,
@@ -137,6 +156,31 @@ impl UdpSocket {
}
}
/// Creates a future that receive a datagram to be written to the buffer
/// provided.
///
/// The returned future will return after a datagram has been received on
/// this socket. The future will resolve to the socket, the buffer, the
/// amount of data read, and the address the data was received from.
///
/// An error during reading will cause the socket and buffer to get
/// destroyed and the socket will be returned.
///
/// The `buf` parameter here only requires the `AsMut<[u8]>` trait, which
/// should be broadly applicable to accepting data which can be converted
/// to a slice. The `Window` struct is also available in this crate to
/// provide a different window into a slice if necessary.
pub fn recv_dgram<T>(self, buf: T) -> RecvDgram<T>
where T: AsMut<[u8]>,
{
RecvDgram {
state: RecvState::Reading {
sock: self,
buf: buf,
},
}
}
/// Gets the value of the `SO_BROADCAST` option for this socket.
///
/// For more information about this option, see
@@ -284,16 +328,14 @@ impl fmt::Debug for UdpSocket {
}
}
/// A future used to write the entire contents of some data to a stream.
/// A future used to write the entire contents of some data to a UDP socket.
///
/// This is created by the [`write_all`] top-level method.
///
/// [`write_all`]: fn.write_all.html
pub struct SendDGram<T> {
state: UdpState<T>,
/// This is created by the `UdpSocket::send_dgram` method.
pub struct SendDgram<T> {
state: SendState<T>,
}
enum UdpState<T> {
enum SendState<T> {
Writing {
sock: UdpSocket,
buf: T,
@@ -302,11 +344,11 @@ enum UdpState<T> {
Empty,
}
fn incomplete_write(reason : &str) -> io::Error {
fn incomplete_write(reason: &str) -> io::Error {
io::Error::new(io::ErrorKind::Other, reason)
}
impl<T> Future for SendDGram<T>
impl<T> Future for SendDgram<T>
where T: AsRef<[u8]>,
{
type Item = (UdpSocket, T);
@@ -314,19 +356,59 @@ impl<T> Future for SendDGram<T>
fn poll(&mut self) -> Poll<(UdpSocket, T), io::Error> {
match self.state {
UdpState::Writing { ref sock, ref buf, ref addr} => {
let buf = buf.as_ref();
let n = try_nb!(sock.send_to(&buf, addr));
if n != buf.len() {
return Err(incomplete_write("Failed to send entire message in datagram"))
SendState::Writing { ref sock, ref buf, ref addr } => {
let n = try_nb!(sock.send_to(buf.as_ref(), addr));
if n != buf.as_ref().len() {
return Err(incomplete_write("failed to send entire message \
in datagram"))
}
}
UdpState::Empty => panic!("poll a SendDGram after it's done"),
SendState::Empty => panic!("poll a SendDgram after it's done"),
}
match mem::replace(&mut self.state, UdpState::Empty) {
UdpState::Writing { sock, buf, .. } => Ok(Async::Ready((sock, (buf).into()))),
UdpState::Empty => panic!(),
match mem::replace(&mut self.state, SendState::Empty) {
SendState::Writing { sock, buf, addr: _ } => {
Ok(Async::Ready((sock, buf)))
}
SendState::Empty => panic!(),
}
}
}
/// A future used to receive a datagram from a UDP socket.
///
/// This is created by the `UdpSocket::recv_dgram` method.
pub struct RecvDgram<T> {
state: RecvState<T>,
}
enum RecvState<T> {
Reading {
sock: UdpSocket,
buf: T,
},
Empty,
}
impl<T> Future for RecvDgram<T>
where T: AsMut<[u8]>,
{
type Item = (UdpSocket, T, usize, SocketAddr);
type Error = io::Error;
fn poll(&mut self) -> Poll<Self::Item, io::Error> {
let (n, addr) = match self.state {
RecvState::Reading { ref sock, ref mut buf } => {
try_nb!(sock.recv_from(buf.as_mut()))
}
RecvState::Empty => panic!("poll a RecvDgram after it's done"),
};
match mem::replace(&mut self.state, RecvState::Empty) {
RecvState::Reading { sock, buf } => {
Ok(Async::Ready((sock, buf, n, addr)))
}
RecvState::Empty => panic!(),
}
}
}