rustev_remote_http/
wire.rs

1//! What left the process, and bounded reading of what came back.
2//!
3//! [`SentTracker`] knows whether any request byte may have left the process.
4//! It is a one-way state machine shared between the connection and the
5//! attempt: the connection must win [`SentTracker::begin_write`] before its
6//! first request byte, and the attempt wins [`SentTracker::try_abort`] only
7//! if no write ever began, after which none can. That is what lets an
8//! adapter acknowledge `stopped` for an unsent request without a race
9//! (spec 009, 3.4.3). Handshake bytes (TLS) are below the tracked layer and
10//! are not request bytes.
11
12use std::io;
13use std::pin::Pin;
14use std::sync::atomic::{AtomicU8, AtomicU64, Ordering};
15use std::sync::{Arc, Mutex};
16use std::task::{Context, Poll};
17
18use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};
19
20use crate::phases::PhaseClock;
21
22const IDLE: u8 = 0;
23const WRITING: u8 = 1;
24const ABORTED: u8 = 2;
25
26/// Whether request bytes may have left the process. Clones share state.
27#[derive(Debug, Clone, Default)]
28pub struct SentTracker(Arc<Tracker>);
29
30#[derive(Debug, Default)]
31struct Tracker {
32    state: AtomicU8,
33    bytes: AtomicU64,
34}
35
36impl SentTracker {
37    pub fn new() -> Self {
38        Self::default()
39    }
40
41    /// Called before a request byte is written. False once aborted: the
42    /// write must not happen.
43    pub fn begin_write(&self) -> bool {
44        match self
45            .0
46            .state
47            .compare_exchange(IDLE, WRITING, Ordering::SeqCst, Ordering::SeqCst)
48        {
49            Ok(_) => true,
50            Err(s) => s == WRITING,
51        }
52    }
53
54    /// Forbid every future write. True if no write ever began, so no request
55    /// byte left the process and none will; false if one may have.
56    pub fn try_abort(&self) -> bool {
57        match self
58            .0
59            .state
60            .compare_exchange(IDLE, ABORTED, Ordering::SeqCst, Ordering::SeqCst)
61        {
62            Ok(_) => true,
63            Err(s) => s == ABORTED,
64        }
65    }
66
67    /// True once [`Self::try_abort`] won: no request byte ever left.
68    pub fn is_aborted(&self) -> bool {
69        self.0.state.load(Ordering::SeqCst) == ABORTED
70    }
71
72    /// True once a write began: request bytes may have left the process.
73    /// Conservative: a write that began counts even if it failed.
74    pub fn may_have_sent(&self) -> bool {
75        self.0.state.load(Ordering::SeqCst) == WRITING
76    }
77
78    /// Bytes the tracked layer accepted for sending.
79    pub fn bytes(&self) -> u64 {
80        self.0.bytes.load(Ordering::SeqCst)
81    }
82
83    fn add(&self, n: usize) {
84        self.0.bytes.fetch_add(n as u64, Ordering::SeqCst);
85    }
86}
87
88/// What a tracked stream reports to: the exchange's [`SentTracker`] and,
89/// under an enabled phase policy, its phase clock (spec 020, 3.4). A pooled
90/// connection (spec 020, 3.5) carries several exchanges, one at a time, so
91/// its hooks are swapped for each.
92#[derive(Debug, Clone, Default)]
93pub struct IoHooks {
94    pub sent: SentTracker,
95    pub clock: Option<PhaseClock>,
96}
97
98/// Hooks shared by a stream and the exchanges that use it.
99pub type SharedHooks = Arc<Mutex<IoHooks>>;
100
101/// A stream whose writes are tracked by a [`SentTracker`]. Wrap the layer
102/// that carries request bytes: plain TCP, or the plaintext side of TLS.
103pub struct TrackedIo<S> {
104    inner: S,
105    hooks: SharedHooks,
106}
107
108fn hooks_of(h: &SharedHooks) -> IoHooks {
109    h.lock().unwrap_or_else(|e| e.into_inner()).clone()
110}
111
112impl<S> TrackedIo<S> {
113    pub fn new(inner: S, sent: SentTracker) -> Self {
114        TrackedIo {
115            inner,
116            hooks: Arc::new(Mutex::new(IoHooks { sent, clock: None })),
117        }
118    }
119
120    /// A stream reporting to hooks its owner may swap between exchanges.
121    pub fn shared(inner: S, hooks: SharedHooks) -> Self {
122        TrackedIo { inner, hooks }
123    }
124
125    /// Report completed flushes and arriving response bytes to `clock`.
126    pub fn with_clock(self, clock: Option<PhaseClock>) -> Self {
127        self.hooks.lock().unwrap_or_else(|e| e.into_inner()).clock = clock;
128        self
129    }
130
131    /// The tracker of the exchange now using the stream.
132    pub fn tracker(&self) -> SentTracker {
133        hooks_of(&self.hooks).sent
134    }
135}
136
137fn aborted() -> io::Error {
138    io::Error::new(
139        io::ErrorKind::ConnectionAborted,
140        "the attempt stopped before its request was sent",
141    )
142}
143
144impl<S: AsyncRead + Unpin> AsyncRead for TrackedIo<S> {
145    fn poll_read(
146        mut self: Pin<&mut Self>,
147        cx: &mut Context<'_>,
148        buf: &mut ReadBuf<'_>,
149    ) -> Poll<io::Result<()>> {
150        let before = buf.filled().len();
151        let r = Pin::new(&mut self.inner).poll_read(cx, buf);
152        if let Poll::Ready(Ok(())) = &r
153            && buf.filled().len() > before
154            && let Some(c) = hooks_of(&self.hooks).clock
155        {
156            c.response_byte();
157        }
158        r
159    }
160}
161
162impl<S: AsyncWrite + Unpin> AsyncWrite for TrackedIo<S> {
163    fn poll_write(
164        mut self: Pin<&mut Self>,
165        cx: &mut Context<'_>,
166        buf: &[u8],
167    ) -> Poll<io::Result<usize>> {
168        if buf.is_empty() {
169            return Poll::Ready(Ok(0));
170        }
171        let sent = hooks_of(&self.hooks).sent;
172        if !sent.begin_write() {
173            return Poll::Ready(Err(aborted()));
174        }
175        let r = Pin::new(&mut self.inner).poll_write(cx, buf);
176        if let Poll::Ready(Ok(n)) = &r {
177            sent.add(*n);
178        }
179        r
180    }
181
182    fn poll_write_vectored(
183        mut self: Pin<&mut Self>,
184        cx: &mut Context<'_>,
185        bufs: &[io::IoSlice<'_>],
186    ) -> Poll<io::Result<usize>> {
187        if bufs.iter().all(|b| b.is_empty()) {
188            return Poll::Ready(Ok(0));
189        }
190        let sent = hooks_of(&self.hooks).sent;
191        if !sent.begin_write() {
192            return Poll::Ready(Err(aborted()));
193        }
194        let r = Pin::new(&mut self.inner).poll_write_vectored(cx, bufs);
195        if let Poll::Ready(Ok(n)) = &r {
196            sent.add(*n);
197        }
198        r
199    }
200
201    fn is_write_vectored(&self) -> bool {
202        self.inner.is_write_vectored()
203    }
204
205    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
206        let r = Pin::new(&mut self.inner).poll_flush(cx);
207        if let Poll::Ready(Ok(())) = &r
208            && let Some(c) = hooks_of(&self.hooks).clock
209        {
210            c.flushed();
211        }
212        r
213    }
214
215    fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
216        Pin::new(&mut self.inner).poll_shutdown(cx)
217    }
218}
219
220/// `sha256:` and the lowercase hex digest of `bytes`, for exchange records.
221pub fn digest_bytes(bytes: &[u8]) -> String {
222    use sha2::{Digest, Sha256};
223    let d = Sha256::digest(bytes);
224    let mut s = String::with_capacity(71);
225    s.push_str("sha256:");
226    for b in d {
227        s.push_str(&format!("{b:02x}"));
228    }
229    s
230}
231
232#[cfg(test)]
233mod tests {
234    use super::*;
235    use tokio::io::AsyncWriteExt;
236
237    #[test]
238    fn an_abort_before_any_write_forbids_every_write() {
239        let t = SentTracker::new();
240        assert!(!t.may_have_sent());
241        assert!(t.try_abort());
242        assert!(t.try_abort(), "idempotent");
243        assert!(t.is_aborted());
244        assert!(!t.begin_write());
245        assert!(!t.may_have_sent());
246    }
247
248    #[test]
249    fn a_write_that_began_cannot_be_unsent() {
250        let t = SentTracker::new();
251        assert!(t.begin_write());
252        assert!(t.begin_write());
253        assert!(!t.try_abort());
254        assert!(t.may_have_sent());
255    }
256
257    #[tokio::test(flavor = "current_thread")]
258    async fn tracked_writes_are_counted_and_refused_after_abort() {
259        let t = SentTracker::new();
260        let mut io = TrackedIo::new(Vec::<u8>::new(), t.clone());
261        io.write_all(b"hello").await.unwrap();
262        assert_eq!(t.bytes(), 5);
263        assert!(t.may_have_sent());
264        let t2 = SentTracker::new();
265        assert!(t2.try_abort());
266        let mut io = TrackedIo::new(Vec::<u8>::new(), t2.clone());
267        assert!(io.write_all(b"x").await.is_err());
268        assert_eq!(t2.bytes(), 0);
269    }
270
271    #[test]
272    fn digests_are_tagged_hex() {
273        let d = digest_bytes(b"");
274        assert_eq!(
275            d,
276            "sha256:e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855"
277        );
278    }
279}
rustev_remote_http/wire.rs

rustev_remote_http/
wire.rs
