tokio_udp/frame.rs
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223
use std::io;
use std::net::{Ipv4Addr, SocketAddr, SocketAddrV4};
use futures::{Async, AsyncSink, Poll, Sink, StartSend, Stream};
use super::UdpSocket;
use bytes::{BufMut, BytesMut};
use tokio_codec::{Decoder, Encoder};
/// A unified `Stream` and `Sink` interface to an underlying `UdpSocket`, using
/// the `Encoder` and `Decoder` traits to encode and decode frames.
///
/// 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 `Encoder` and `Decoder` traits 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.
#[must_use = "sinks do nothing unless polled"]
#[derive(Debug)]
pub struct UdpFramed<C> {
socket: UdpSocket,
codec: C,
rd: BytesMut,
wr: BytesMut,
out_addr: SocketAddr,
flushed: bool,
is_readable: bool,
repeat_decode: bool,
current_addr: Option<SocketAddr>,
}
impl<C: Decoder> Stream for UdpFramed<C> {
type Item = (C::Item, SocketAddr);
type Error = C::Error;
fn poll(&mut self) -> Poll<Option<Self::Item>, Self::Error> {
self.rd.reserve(INITIAL_RD_CAPACITY);
if self.repeat_decode {
loop {
// Are there are still bytes left in the read buffer to decode?
if self.is_readable {
// Use deocde_eof since every datagram contains its own
// eof which is just the end of the datagram. This supports
// the lines use case where there may not be a terminating
// delimiter and thus you may never get the end of the frame.
// This is generally fine for most implementations of codec
// since by default this will defer to calling decode.
if let Some(frame) = self.codec.decode_eof(&mut self.rd)? {
trace!("frame decoded from buffer");
let current_addr = self
.current_addr
.expect("will always be set before this line is called");
return Ok(Async::Ready(Some((frame, current_addr))));
}
// if this line has been reached then decode has returned `None`.
self.is_readable = false;
self.rd.clear();
}
// We're out of data. Try and fetch more data to decode
let (n, addr) = unsafe {
// Read into the buffer without having to initialize the memory.
let (n, addr) = try_ready!(self.socket.poll_recv_from(self.rd.bytes_mut()));
self.rd.advance_mut(n);
(n, addr)
};
self.current_addr = Some(addr);
self.is_readable = true;
trace!("received {} bytes, decoding", n);
}
} else {
let (n, addr) = unsafe {
// Read into the buffer without having to initialize the memory.
let (n, addr) = try_ready!(self.socket.poll_recv_from(self.rd.bytes_mut()));
self.rd.advance_mut(n);
(n, addr)
};
trace!("received {} bytes, decoding", n);
let frame_res = self.codec.decode(&mut self.rd);
self.rd.clear();
let frame = frame_res?;
let result = frame.map(|frame| (frame, addr)); // frame -> (frame, addr)
trace!("frame decoded from buffer");
Ok(Async::Ready(result))
}
}
}
impl<C: Encoder> Sink for UdpFramed<C> {
type SinkItem = (C::Item, SocketAddr);
type SinkError = C::Error;
fn start_send(&mut self, item: Self::SinkItem) -> StartSend<Self::SinkItem, Self::SinkError> {
trace!("sending frame");
if !self.flushed {
match self.poll_complete()? {
Async::Ready(()) => {}
Async::NotReady => return Ok(AsyncSink::NotReady(item)),
}
}
let (frame, out_addr) = item;
self.codec.encode(frame, &mut self.wr)?;
self.out_addr = out_addr;
self.flushed = false;
trace!("frame encoded; length={}", self.wr.len());
Ok(AsyncSink::Ready)
}
fn poll_complete(&mut self) -> Poll<(), C::Error> {
if self.flushed {
return Ok(Async::Ready(()));
}
trace!("flushing frame; length={}", self.wr.len());
let n = try_ready!(self.socket.poll_send_to(&self.wr, &self.out_addr));
trace!("written {}", n);
let wrote_all = n == self.wr.len();
self.wr.clear();
self.flushed = true;
if wrote_all {
Ok(Async::Ready(()))
} else {
Err(io::Error::new(
io::ErrorKind::Other,
"failed to write entire datagram to socket",
)
.into())
}
}
fn close(&mut self) -> Poll<(), C::Error> {
try_ready!(self.poll_complete());
Ok(().into())
}
}
const INITIAL_RD_CAPACITY: usize = 64 * 1024;
const INITIAL_WR_CAPACITY: usize = 8 * 1024;
impl<C> UdpFramed<C> {
/// Create a new `UdpFramed` backed by the given socket and codec.
///
/// See struct level documentation for more details.
pub fn new(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: BytesMut::with_capacity(INITIAL_RD_CAPACITY),
wr: BytesMut::with_capacity(INITIAL_WR_CAPACITY),
flushed: true,
is_readable: false,
repeat_decode: false,
current_addr: None,
}
}
/// Create a new `UdpFramed` backed by the given socket and codec. That will
/// continue to call `decode_eof` until the decoder has cleared the entire buffer.
///
/// See struct level documentation for more details.
pub fn with_decode(socket: UdpSocket, codec: C, repeat_decode: bool) -> UdpFramed<C> {
UdpFramed {
socket: socket,
codec: codec,
out_addr: SocketAddr::V4(SocketAddrV4::new(Ipv4Addr::new(0, 0, 0, 0), 0)),
rd: BytesMut::with_capacity(INITIAL_RD_CAPACITY),
wr: BytesMut::with_capacity(INITIAL_WR_CAPACITY),
flushed: true,
is_readable: false,
repeat_decode,
current_addr: None,
}
}
/// Returns a reference to the underlying I/O stream wrapped by `Framed`.
///
/// # Note
///
/// 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
///
/// 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.
pub fn into_inner(self) -> UdpSocket {
self.socket
}
}