dfir_rs/util/
udp.rs

Help
1#![cfg(not(target_arch = "wasm32"))]
2
3use std::net::SocketAddr;
4
5use bytes::Bytes;
6use futures::stream::{SplitSink, SplitStream};
7use tokio::net::UdpSocket;
8use tokio_util::codec::length_delimited::LengthDelimitedCodec;
9use tokio_util::codec::{BytesCodec, Decoder, Encoder, LinesCodec};
10use tokio_util::udp::UdpFramed;
11
12/// A framed UDP `Sink` (sending).
13pub type UdpFramedSink<Codec, Item> = SplitSink<UdpFramed<Codec>, (Item, SocketAddr)>;
14/// A framed UDP `Stream` (receiving).
15pub type UdpFramedStream<Codec> = SplitStream<UdpFramed<Codec>>;
16/// Returns a UDP `Stream`, `Sink`, and address for the given socket, using the given `Codec` to
17/// handle delineation between inputs/outputs.
18pub fn udp_framed<Codec, Item>(
19    socket: UdpSocket,
20    codec: Codec,
21) -> (
22    UdpFramedSink<Codec, Item>,
23    UdpFramedStream<Codec>,
24    SocketAddr,
25)
26where
27    Codec: Encoder<Item> + Decoder,
28{
29    let framed = UdpFramed::new(socket, codec);
30    let addr = framed.get_ref().local_addr().unwrap();
31    let split = futures::stream::StreamExt::split(framed);
32    (split.0, split.1, addr)
33}
34
35/// A UDP length-delimited frame `Sink` (sending).
36pub type UdpSink = UdpFramedSink<LengthDelimitedCodec, Bytes>;
37/// A UDP length-delimited frame `Stream` (receiving).
38pub type UdpStream = UdpFramedStream<LengthDelimitedCodec>;
39/// Helper creates a UDP `Stream` and `Sink` for `Bytes` strings where each string is
40/// length-delimited.
41pub fn udp_bytes(socket: UdpSocket) -> (UdpSink, UdpStream, SocketAddr) {
42    udp_framed(socket, LengthDelimitedCodec::new())
43}
44
45/// A UDP undelimited bytes `Sink` (sending).
46pub type UdpBytesSink = UdpFramedSink<BytesCodec, Bytes>;
47/// A UDP undelimited bytes `Stream` (receiving).
48pub type UdpBytesStream = UdpFramedStream<BytesCodec>;
49/// Helper creates a UDP `Stream` and `Sink` for undelimited streams of `Bytes`.
50pub fn udp_bytestream(socket: UdpSocket) -> (UdpBytesSink, UdpBytesStream, SocketAddr) {
51    udp_framed(socket, BytesCodec::new())
52}
53
54/// A UDP newline-delimited `String` `Sink` (sending).
55pub type UdpLinesSink = UdpFramedSink<LinesCodec, String>;
56/// A UDP newline-delimited `String` `Stream` (receivng).
57pub type UdpLinesStream = UdpFramedStream<LinesCodec>;
58/// Helper creates a UDP `Stream` and `Sink` for `String`s delimited by newlines.
59pub fn udp_lines(
60    socket: UdpSocket,
61) -> (
62    UdpFramedSink<LinesCodec, String>,
63    UdpFramedStream<LinesCodec>,
64    SocketAddr,
65) {
66    udp_framed(socket, LinesCodec::new())
67}
dfir_rs/util/udp.rs

dfir_rs/util/
udp.rs
