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

#![cfg(not(target_arch = "wasm32"))]

use std::net::SocketAddr;

use bytes::Bytes;
use futures::stream::{SplitSink, SplitStream};
use tokio::net::UdpSocket;
use tokio_util::codec::length_delimited::LengthDelimitedCodec;
use tokio_util::codec::{BytesCodec, Decoder, Encoder, LinesCodec};
use tokio_util::udp::UdpFramed;

/// A framed UDP `Sink` (sending).
pub type UdpFramedSink<Codec, Item> = SplitSink<UdpFramed<Codec>, (Item, SocketAddr)>;
/// A framed UDP `Stream` (receiving).
pub type UdpFramedStream<Codec> = SplitStream<UdpFramed<Codec>>;
/// Returns a UDP `Stream`, `Sink`, and address for the given socket, using the given `Codec` to
/// handle delineation between inputs/outputs.
pub fn udp_framed<Codec, Item>(
    socket: UdpSocket,
    codec: Codec,
) -> (
    UdpFramedSink<Codec, Item>,
    UdpFramedStream<Codec>,
    SocketAddr,
)
where
    Codec: Encoder<Item> + Decoder,
{
    let framed = UdpFramed::new(socket, codec);
    let addr = framed.get_ref().local_addr().unwrap();
    let split = futures::stream::StreamExt::split(framed);
    (split.0, split.1, addr)
}

/// A UDP length-delimited frame `Sink` (sending).
pub type UdpSink = UdpFramedSink<LengthDelimitedCodec, Bytes>;
/// A UDP length-delimited frame `Stream` (receiving).
pub type UdpStream = UdpFramedStream<LengthDelimitedCodec>;
/// Helper creates a UDP `Stream` and `Sink` for `Bytes` strings where each string is
/// length-delimited.
pub fn udp_bytes(socket: UdpSocket) -> (UdpSink, UdpStream, SocketAddr) {
    udp_framed(socket, LengthDelimitedCodec::new())
}

/// A UDP undelimited bytes `Sink` (sending).
pub type UdpBytesSink = UdpFramedSink<BytesCodec, Bytes>;
/// A UDP undelimited bytes `Stream` (receiving).
pub type UdpBytesStream = UdpFramedStream<BytesCodec>;
/// Helper creates a UDP `Stream` and `Sink` for undelimited streams of `Bytes`.
pub fn udp_bytestream(socket: UdpSocket) -> (UdpBytesSink, UdpBytesStream, SocketAddr) {
    udp_framed(socket, BytesCodec::new())
}

/// A UDP newline-delimited `String` `Sink` (sending).
pub type UdpLinesSink = UdpFramedSink<LinesCodec, String>;
/// A UDP newline-delimited `String` `Stream` (receivng).
pub type UdpLinesStream = UdpFramedStream<LinesCodec>;
/// Helper creates a UDP `Stream` and `Sink` for `String`s delimited by newlines.
pub fn udp_lines(
    socket: UdpSocket,
) -> (
    UdpFramedSink<LinesCodec, String>,
    UdpFramedStream<LinesCodec>,
    SocketAddr,
) {
    udp_framed(socket, LinesCodec::new())
}