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
use super::StreamBase;
use async_lock::{Mutex, MutexGuard};
use std::io;
use x11rb::errors::ConnectionError;
use x11rb_protocol::RawFdContainer;
#[derive(Debug, Default)]
pub(super) struct WriteBuffer(Mutex<WriteBufferInner>);
#[derive(Debug)]
pub(super) struct WriteBufferGuard<'a>(MutexGuard<'a, WriteBufferInner>);
#[derive(Debug)]
pub(super) struct WriteBufferInner {
/// The buffer that is used for writing.
buffer: Vec<u8>,
/// The file descriptors that we are sending over.
fds: Vec<RawFdContainer>,
/// Whether the buffer has been corrupted.
///
/// A lock has to be explicitly unlock()d, otherwise the buffer is marked as corrupted.
/// This exists to detect futures that were not polled to completion and might have
/// written only a part of their data.
corrupted: bool,
}
impl Default for WriteBufferInner {
fn default() -> Self {
Self {
buffer: Vec::with_capacity(16384),
fds: vec![],
corrupted: false,
}
}
}
impl WriteBuffer {
/// Lock the write buffer for writing.
///
/// The returned guard must be unlocked with [`unlock()`] or else the write buffer will be
/// considered corrupted. This mechanism exists to catch futures being dropped without being
/// polled to completion. In this situation we cannot be sure how many bytes were already
/// written to the stream, so the complete connection is now broken.
pub(super) async fn lock(&self) -> Result<WriteBufferGuard<'_>, ConnectionError> {
let mut lock = self.0.lock().await;
if std::mem::replace(&mut lock.corrupted, true) {
return Err(ConnectionError::IoError(io::Error::new(
io::ErrorKind::Other,
"The write buffer was corrupted",
)));
}
Ok(WriteBufferGuard(lock))
}
}
impl WriteBufferGuard<'_> {
/// Unlock this guard.
pub(super) fn unlock(mut self) {
self.0.corrupted = false;
}
}
impl std::ops::Deref for WriteBufferGuard<'_> {
type Target = WriteBufferInner;
fn deref(&self) -> &Self::Target {
&self.0
}
}
impl std::ops::DerefMut for WriteBufferGuard<'_> {
fn deref_mut(&mut self) -> &mut Self::Target {
&mut self.0
}
}
impl WriteBufferInner {
/// Flush the write buffer.
pub(super) async fn flush<'b, S: StreamBase<'b>>(
&mut self,
stream: &'b S,
) -> Result<(), ConnectionError> {
// If we don't have any data to write, we are done.
if self.buffer.is_empty() && self.fds.is_empty() {
return Ok(());
}
// Write the entire buffer.
tracing::trace!(
"Flushing {} bytes of data and {} FDs",
self.buffer.len(),
self.fds.len()
);
let mut position = 0;
super::write_with(stream, {
let (buffer, fds) = (&mut self.buffer, &mut self.fds);
move |stream| {
while position < buffer.len() {
let n = stream.write(&buffer[position..], fds)?;
if n == 0 {
return Err(io::Error::new(
io::ErrorKind::WriteZero,
"failed to write whole buffer",
));
}
tracing::trace!("Flushing wrote {} bytes of data", n);
position += n;
}
Ok(())
}
})
.await?;
if !self.fds.is_empty() {
return Err(ConnectionError::IoError(io::Error::new(
io::ErrorKind::Other,
"failed to write all fds",
)));
}
// Reset the buffer.
self.buffer.clear();
Ok(())
}
/// Write a set of buffers to the stream.
pub(super) async fn write_all_vectored<'b, S: StreamBase<'b>>(
&mut self,
stream: &'b S,
mut bufs: &[io::IoSlice<'_>],
fds: &mut Vec<RawFdContainer>,
) -> Result<(), ConnectionError> {
// Get the total length of the buffers.
let mut total_len = bufs
.iter()
.fold(0usize, |sum, buf| sum.saturating_add(buf.len()));
tracing::trace!("Writing {} bytes of data and {} fds", total_len, fds.len());
// If our data doesn't fit, flush the buffer first.
if self.buffer.len() + total_len > self.buffer.capacity() {
self.flush(stream).await?;
}
// If our data fits now, write all of it.
if total_len < self.buffer.capacity() {
tracing::trace!("Data to write is appended to the buffer");
for buf in bufs {
self.buffer.extend_from_slice(buf);
}
self.fds.append(fds);
return Ok(());
}
debug_assert!(self.buffer.is_empty());
// Otherwise, write directly to the stream.
tracing::trace!("Data to write is written directly to the stream");
let mut partial: &[u8] = &[];
super::write_with(stream, |stream| {
while total_len > 0 || !partial.is_empty() {
// If the partial buffer is non-empty, write it.
if !partial.is_empty() {
let n = stream.write(partial, fds)?;
if n == 0 {
return Err(io::Error::from(io::ErrorKind::WriteZero));
}
partial = &partial[n..];
total_len -= n;
} else {
// Write the iov.
let mut n = stream.write_vectored(bufs, fds)?;
if n == 0 {
return Err(io::Error::from(io::ErrorKind::WriteZero));
}
// Calculate how much we have left to go.
total_len -= n;
while n > 0 {
if n >= bufs[0].len() {
n -= bufs[0].len();
bufs = &bufs[1..];
} else {
partial = &bufs[0][n..];
n = 0;
}
}
}
}
if !fds.is_empty() {
return Err(io::Error::new(
io::ErrorKind::Other,
"Left over FDs after sending the request",
));
}
Ok(())
})
.await?;
Ok(())
}
}