-
-
Notifications
You must be signed in to change notification settings - Fork 2.3k
/
split.rs
147 lines (124 loc) · 4.24 KB
/
split.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
//! `UnixStream` and `UnixDatagram` split support.
//!
//! ## UnixStream
//!
//! A `UnixStream` can be split into a read half and a write half with
//! `UnixStream::split`. The read half implements `AsyncRead` while the write
//! half implements `AsyncWrite`.
//!
//! Compared to the generic split of `AsyncRead + AsyncWrite`, this specialized
//! split has no associated overhead and enforces all invariants at the type
//! level.
//!
//! ## UnixDatagram
//!
//! A `UnixDatagram` can be split into a receive half and a send half with
//! `UnixDatagram::split`. The send half implements `send` and `send_to` and
//! the receiving one implements `recv` and `recv_from`.
//!
//! This split method has no overhead and enforces all invariants at the type
//! level.
use crate::io::{AsyncRead, AsyncWrite};
use crate::net::{UnixDatagram, UnixStream};
use crate::future::poll_fn;
use std::io;
use std::mem::MaybeUninit;
use std::net::Shutdown;
use std::os::unix::net::SocketAddr;
use std::path::Path;
use std::pin::Pin;
use std::task::{Context, Poll};
/// Read half of a `UnixStream`.
#[derive(Debug)]
pub struct ReadHalf<'a>(&'a UnixStream);
/// Write half of a `UnixStream`.
#[derive(Debug)]
pub struct WriteHalf<'a>(&'a UnixStream);
/// Receiving half of a `UnixDatagram`.
#[derive(Debug)]
pub struct RecvHalf<'a>(&'a UnixDatagram);
/// Sending half of a `UnixDatagram`.
#[derive(Debug)]
pub struct SendHalf<'a>(&'a UnixDatagram);
pub(crate) fn split_stream(stream: &mut UnixStream) -> (ReadHalf<'_>, WriteHalf<'_>) {
(ReadHalf(stream), WriteHalf(stream))
}
pub(crate) fn split_dgram(dgram: &mut UnixDatagram) -> (RecvHalf<'_>, SendHalf<'_>) {
(RecvHalf(dgram), SendHalf(dgram))
}
impl AsyncRead for ReadHalf<'_> {
unsafe fn prepare_uninitialized_buffer(&self, _: &mut [MaybeUninit<u8>]) -> bool {
false
}
fn poll_read(
self: Pin<&mut Self>,
cx: &mut Context<'_>,
buf: &mut [u8],
) -> Poll<io::Result<usize>> {
self.0.poll_read_priv(cx, buf)
}
}
impl AsyncWrite for WriteHalf<'_> {
fn poll_write(
self: Pin<&mut Self>,
cx: &mut Context<'_>,
buf: &[u8],
) -> Poll<io::Result<usize>> {
self.0.poll_write_priv(cx, buf)
}
fn poll_flush(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<io::Result<()>> {
Poll::Ready(Ok(()))
}
fn poll_shutdown(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<io::Result<()>> {
self.0.shutdown(Shutdown::Write).into()
}
}
impl AsRef<UnixStream> for ReadHalf<'_> {
fn as_ref(&self) -> &UnixStream {
self.0
}
}
impl AsRef<UnixStream> for WriteHalf<'_> {
fn as_ref(&self) -> &UnixStream {
self.0
}
}
impl RecvHalf<'_> {
/// Receives a datagram from the socket.
pub async fn recv(&mut self, buf: &mut [u8]) -> io::Result<usize> {
poll_fn(|cx| self.0.poll_recv_priv(cx, buf)).await
}
/// Try to receive a datagram from the peer without waiting.
pub fn try_recv(&mut self, buf: &mut [u8]) -> io::Result<usize> {
self.0.try_recv(buf)
}
/// Receives a datagram with the source address from the socket.
pub async fn recv_from(&mut self, buf: &mut [u8]) -> io::Result<(usize, SocketAddr)> {
poll_fn(|cx| self.0.poll_recv_from_priv(cx, buf)).await
}
/// Try to receive data from the socket without waiting.
pub fn try_recv_from(&mut self, buf: &mut [u8]) -> io::Result<(usize, SocketAddr)> {
self.0.try_recv_from(buf)
}
}
impl SendHalf<'_> {
/// Sends a datagram to the socket's peer.
pub async fn send(&mut self, buf: &[u8]) -> io::Result<usize> {
poll_fn(|cx| self.0.poll_send_priv(cx, buf)).await
}
/// Try to send a datagram to the peer without waiting.
pub fn try_send(&mut self, buf: &[u8]) -> io::Result<usize> {
self.0.try_send(buf)
}
/// Sends a datagram to the specified address.
pub async fn send_to<P>(&mut self, buf: &[u8], target: P) -> io::Result<usize>
where
P: AsRef<Path> + Unpin,
{
poll_fn(|cx| self.0.poll_send_to_priv(cx, buf, target.as_ref())).await
}
/// Try to send a datagram to the peer without waiting.
pub fn try_send_to<P>(&mut self, buf: &[u8], target: P) -> io::Result<usize> {
self.0.try_send_to(buf, target)
}
}