base_endpoint/tls/rustls/rustls_connection/
duplex.rs

1// Released under the MIT License.
2// Copyright, 2026, by Samuel Williams.
3
4use super::RustlsConnection;
5use crate::stream::{ReadableStream, WritableStream};
6use base_scheduler::{
7    Cancellation, OperationError, OperationResult, SocketClose, SocketStreamReader,
8    SocketStreamReading, SocketStreamWriter, SocketStreamWriting,
9};
10use event_listener::Event;
11use futures_util::{future::join, task::AtomicWaker};
12use std::{
13    convert::Infallible,
14    future::{Future, poll_fn},
15    io::{self, BufRead, Read, Write},
16    pin::Pin,
17    sync::{Arc, Mutex, MutexGuard, PoisonError},
18    task::{Context, Poll, Wake, Waker},
19};
20
21const CAPACITY: usize = 16 * 1024;
22
23#[cfg(test)]
24mod tests;
25
26fn lock<T>(value: &Mutex<T>) -> MutexGuard<'_, T> {
27    value.lock().unwrap_or_else(PoisonError::into_inner)
28}
29
30#[derive(Clone)]
31enum Failure {
32    Cancelled,
33    Io(Arc<io::Error>),
34}
35impl Failure {
36    fn from_error(error: OperationError) -> Self {
37        match error {
38            OperationError::Cancelled(_) => Self::Cancelled,
39            OperationError::Io(error) => Self::Io(Arc::new(error)),
40        }
41    }
42    fn error(&self) -> OperationError {
43        match self {
44            Self::Cancelled => base_scheduler::Cancelled.into(),
45            Self::Io(error) => io::Error::new(error.kind(), error.clone()).into(),
46        }
47    }
48}
49
50struct State {
51    tls: ::rustls::Connection,
52    plaintext: usize,
53    issued: u64,
54    delivered: u64,
55    failure: Option<Failure>,
56    finished: bool,
57    output_closed: bool,
58}
59struct Shared {
60    state: Mutex<State>,
61    writable: Event,
62    transmit: AtomicWaker,
63}
64
65// One retained transport future, with a stable socket waker. Application I/O
66// can make immediate progress without handing off to the supervisor. A blocked
67// socket wakes both possible pollers; neither can replace the other's waker.
68struct ProgressWake {
69    foreground: AtomicWaker,
70    background: AtomicWaker,
71}
72impl Wake for ProgressWake {
73    fn wake(self: Arc<Self>) {
74        self.wake_by_ref();
75    }
76    fn wake_by_ref(self: &Arc<Self>) {
77        self.foreground.wake();
78        self.background.wake();
79    }
80}
81struct ProgressState<'scope> {
82    work: Option<Pin<Box<dyn Future<Output = io::Result<()>> + Send + 'scope>>>,
83    result: Option<Result<(), Arc<io::Error>>>,
84}
85struct IoProgress<'scope> {
86    progress: Mutex<ProgressState<'scope>>,
87    wake: Arc<ProgressWake>,
88    waker: Waker,
89}
90impl<'scope> IoProgress<'scope> {
91    fn new(work: impl Future<Output = io::Result<()>> + Send + 'scope) -> Self {
92        let wake = Arc::new(ProgressWake {
93            foreground: AtomicWaker::new(),
94            background: AtomicWaker::new(),
95        });
96        Self {
97            progress: Mutex::new(ProgressState {
98                work: Some(Box::pin(work)),
99                result: None,
100            }),
101            waker: Waker::from(wake.clone()),
102            wake,
103        }
104    }
105    fn poll(&self, cx: &mut Context<'_>, foreground: bool) -> Poll<io::Result<()>> {
106        let mut progress = lock(&self.progress);
107        if let Some(work) = &mut progress.work {
108            if foreground {
109                self.wake.foreground.register(cx.waker());
110            } else {
111                self.wake.background.register(cx.waker());
112            }
113            let mut cx = Context::from_waker(&self.waker);
114            if let Poll::Ready(result) = work.as_mut().poll(&mut cx) {
115                progress.work = None;
116                progress.result = Some(result.map_err(Arc::new));
117                self.wake.wake_by_ref();
118            }
119        }
120        match &progress.result {
121            Some(Ok(())) => Poll::Ready(Ok(())),
122            Some(Err(error)) => Poll::Ready(Err(io::Error::new(error.kind(), error.clone()))),
123            None => Poll::Pending,
124        }
125    }
126}
127impl Shared {
128    fn notify_all(&self) {
129        self.writable.notify(usize::MAX);
130        self.transmit.wake();
131    }
132    fn fail(&self, error: OperationError) {
133        lock(&self.state)
134            .failure
135            .get_or_insert_with(|| Failure::from_error(error));
136        self.notify_all();
137    }
138    async fn delivered(
139        &self,
140        transmitter: &IoProgress<'_>,
141        ticket: u64,
142        cancellation: &Cancellation,
143    ) -> OperationResult<()> {
144        loop {
145            event_listener::listener!(self.writable => listener);
146            cancellation.check()?;
147            {
148                let state = lock(&self.state);
149                if let Some(error) = &state.failure {
150                    return Err(error.error());
151                }
152                if state.delivered >= ticket {
153                    return Ok(());
154                }
155                if state.finished {
156                    return Err(closed());
157                }
158            }
159            let mut waiting = std::pin::pin!(wait(listener, cancellation));
160            poll_fn(|cx| {
161                if let Poll::Ready(Err(error)) = transmitter.poll(cx, true) {
162                    return Poll::Ready(Err(error.into()));
163                }
164                waiting.as_mut().poll(cx)
165            })
166            .await?;
167        }
168    }
169}
170
171fn closed() -> OperationError {
172    io::Error::new(io::ErrorKind::BrokenPipe, "TLS direction is closed").into()
173}
174
175// Register for cancellation wakeups, reporting a requested cancellation. Lanes
176// check the signal first, so there readiness means a concurrent request.
177fn poll_cancelled<T>(
178    cancelled: Pin<&mut impl Future<Output = ()>>,
179    cx: &mut Context<'_>,
180) -> Poll<OperationResult<T>> {
181    if cancelled.poll(cx).is_ready() {
182        Poll::Ready(Err(base_scheduler::Cancelled.into()))
183    } else {
184        Poll::Pending
185    }
186}
187
188async fn wait(
189    listener: impl Future<Output = ()> + Send,
190    cancellation: &Cancellation,
191) -> OperationResult<()> {
192    let mut listener = std::pin::pin!(listener);
193    let mut cancelled = std::pin::pin!(cancellation.cancelled());
194    std::future::poll_fn(|cx| {
195        if let Poll::Ready(result) = poll_cancelled(cancelled.as_mut(), cx) {
196            return Poll::Ready(result);
197        }
198        listener.as_mut().poll(cx).map(Ok)
199    })
200    .await
201}
202
203/// Plaintext input which drives its retained TLS receive future. No lock spans
204/// an asynchronous wait; paused input does not prevent application output.
205pub struct EncryptedSocketStreamReader<'scope> {
206    shared: Arc<Shared>,
207    receiver: Arc<IoProgress<'scope>>,
208    closed: bool,
209}
210impl Shared {
211    fn read(&self, buffer: &mut [u8]) -> Option<OperationResult<usize>> {
212        let mut state = lock(&self.state);
213        if state.finished {
214            return Some(Err(closed()));
215        }
216        if buffer.is_empty() {
217            return Some(Ok(0));
218        }
219        match state.tls.reader().read(buffer) {
220            Ok(count) => {
221                state.plaintext = state.plaintext.saturating_sub(count);
222                Some(Ok(count))
223            }
224            Err(error) if error.kind() == io::ErrorKind::WouldBlock => {
225                state.failure.as_ref().map(|failure| Err(failure.error()))
226            }
227            Err(error) => Some(Err(error.into())),
228        }
229    }
230}
231impl ReadableStream for EncryptedSocketStreamReader<'_> {
232    async fn read_into(
233        &mut self,
234        buffer: &mut [u8],
235        cancellation: &Cancellation,
236    ) -> OperationResult<usize> {
237        let mut cancelled = std::pin::pin!(cancellation.cancelled());
238        poll_fn(|cx| {
239            cancellation.check()?;
240            if self.closed {
241                return Poll::Ready(Ok(0));
242            }
243            if let Some(result) = self.shared.read(buffer) {
244                return Poll::Ready(result);
245            }
246            let progress = self.receiver.poll(cx, true);
247            if let Some(result) = self.shared.read(buffer) {
248                return Poll::Ready(result);
249            }
250            if let Poll::Ready(Err(error)) = progress {
251                return Poll::Ready(Err(error.into()));
252            }
253            poll_cancelled(cancelled.as_mut(), cx)
254        })
255        .await
256    }
257    async fn close(&mut self) -> io::Result<()> {
258        self.closed = true;
259        Ok(())
260    }
261}
262
263/// Ordered plaintext output. Writes accept a bounded prefix into TLS; flush
264/// waits for its ciphertext to reach the socket. No transport lock spans await.
265pub struct EncryptedSocketStreamWriter<'scope> {
266    shared: Arc<Shared>,
267    transmitter: Arc<IoProgress<'scope>>,
268    closed: bool,
269}
270impl WritableStream for EncryptedSocketStreamWriter<'_> {
271    async fn write(
272        &mut self,
273        buffer: &[u8],
274        cancellation: &Cancellation,
275    ) -> OperationResult<usize> {
276        loop {
277            event_listener::listener!(self.shared.writable => listener);
278            cancellation.check()?;
279            if self.closed {
280                return Err(closed());
281            }
282            let result = {
283                let mut state = lock(&self.shared.state);
284                if let Some(error) = &state.failure {
285                    return Err(error.error());
286                }
287                if state.finished || state.output_closed {
288                    return Err(closed());
289                }
290                // A full rustls buffer accepts zero bytes. Rustls reports only
291                // the connection's fatal error, never `WouldBlock`, from writes.
292                match state.tls.writer().write(buffer) {
293                    Ok(0) if !buffer.is_empty() => None,
294                    Ok(count) => {
295                        if count != 0 {
296                            state.issued += 1;
297                        }
298                        Some(Ok(count))
299                    }
300                    Err(error) => Some(Err(error.into())),
301                }
302            };
303            if let Some(result) = result {
304                poll_fn(|cx| {
305                    let _ = self.transmitter.poll(cx, true);
306                    Poll::Ready(())
307                })
308                .await;
309                return result;
310            }
311            let mut waiting = std::pin::pin!(wait(listener, cancellation));
312            poll_fn(|cx| {
313                if let Poll::Ready(Err(error)) = self.transmitter.poll(cx, true) {
314                    return Poll::Ready(Err(error.into()));
315                }
316                waiting.as_mut().poll(cx)
317            })
318            .await?;
319        }
320    }
321    async fn flush(&mut self, cancellation: &Cancellation) -> OperationResult<()> {
322        if self.closed {
323            return Err(closed());
324        }
325        // Delivery checks cancellation, failure and completed output in that
326        // order before waiting, so an already delivered flush needs no wait.
327        let ticket = lock(&self.shared.state).issued;
328        self.shared
329            .delivered(&self.transmitter, ticket, cancellation)
330            .await
331    }
332    async fn close(&mut self) -> io::Result<()> {
333        self.closed = true;
334        Ok(())
335    }
336}
337
338impl<C: Send + 'static> RustlsConnection<C> {
339    /// Run an application with independent plaintext directions. Reads and
340    /// writes drive retained I/O futures directly; this future supervises
341    /// control output and joins persistent scheduler lanes. No task is spawned.
342    ///
343    /// On successful application completion, send close_notify. On failure or
344    /// cancellation, abort pending I/O. Every path joins lane retirement and
345    /// consumes the socket before returning; dropping this future is not a join.
346    pub async fn with_duplex<'scope, S, F, Fut, T, E>(
347        self,
348        scheduler: &'scope S,
349        cancellation: &Cancellation,
350        application: F,
351    ) -> Result<T, E>
352    where
353        S: SocketStreamReading<Socket = C> + SocketStreamWriting + SocketClose + ?Sized,
354        F: FnOnce(EncryptedSocketStreamReader<'scope>, EncryptedSocketStreamWriter<'scope>) -> Fut
355            + Send,
356        Fut: Future<Output = Result<T, E>> + Send,
357        T: Send,
358        E: From<OperationError> + From<io::Error> + Send,
359    {
360        let Self {
361            socket,
362            mut state,
363            incoming,
364            incoming_start,
365            incoming_end,
366            outgoing,
367            outgoing_start,
368            ..
369        } = self;
370        let result = async {
371            let mut reader = scheduler.socket_stream_reader(&socket)?;
372            let mut writer = match scheduler.socket_stream_writer(&socket) {
373                Ok(writer) => writer,
374                Err(error) => {
375                    let _ = SocketStreamReader::close(&mut reader).await;
376                    return Err(E::from(error));
377                }
378            };
379            state.set_buffer_limit(Some(CAPACITY));
380            let plaintext = state
381                .reader()
382                .fill_buf()
383                .map(|bytes| bytes.len())
384                .unwrap_or(0);
385            let shared = Arc::new(Shared {
386                state: Mutex::new(State {
387                    tls: state,
388                    plaintext,
389                    issued: 0,
390                    delivered: 0,
391                    failure: None,
392                    finished: false,
393                    output_closed: false,
394                }),
395                writable: Event::new(),
396                transmit: AtomicWaker::new(),
397            });
398            let stop = cancellation.child();
399            let transmit_shared = shared.clone();
400            let transmit_stop = stop.clone();
401            let transmit_lifetime = cancellation.clone();
402            let transmitter = Arc::new(IoProgress::new(async move {
403                // Output only stops by failing, including cancellation.
404                let Err(error) = transmit(
405                    &transmit_shared,
406                    &mut writer,
407                    outgoing,
408                    outgoing_start,
409                    &transmit_stop,
410                )
411                .await;
412                if !transmit_stop.is_cancelled() || transmit_lifetime.is_cancelled() {
413                    transmit_shared.fail(error);
414                    transmit_stop.cancel();
415                }
416                SocketStreamWriter::close(&mut writer).await
417            }));
418            let receive_shared = shared.clone();
419            let receive_stop = stop.clone();
420            let receive_lifetime = cancellation.clone();
421            let receiver = Arc::new(IoProgress::new(async move {
422                let result = receive(
423                    &receive_shared,
424                    &mut reader,
425                    &incoming[incoming_start..incoming_end],
426                    &receive_stop,
427                )
428                .await;
429                if let Err(error) = result {
430                    if !receive_stop.is_cancelled() || receive_lifetime.is_cancelled() {
431                        receive_shared.fail(error);
432                        receive_stop.cancel();
433                    }
434                }
435                SocketStreamReader::close(&mut reader).await
436            }));
437            let input = EncryptedSocketStreamReader {
438                shared: shared.clone(),
439                receiver: receiver.clone(),
440                closed: false,
441            };
442            let output = EncryptedSocketStreamWriter {
443                shared: shared.clone(),
444                transmitter: transmitter.clone(),
445                closed: false,
446            };
447            let operation = async {
448                let mut result = application(input, output).await;
449                lock(&shared.state).output_closed = true;
450                if result.is_ok() && !stop.is_cancelled() {
451                    let ticket = {
452                        let mut state = lock(&shared.state);
453                        state.tls.send_close_notify();
454                        state.issued += 1;
455                        state.issued
456                    };
457                    if let Err(error) = shared.delivered(&transmitter, ticket, &stop).await {
458                        // A lane failure cancels supervision; report its cause.
459                        let failure = lock(&shared.state).failure.clone();
460                        result = Err(E::from(failure.map_or(error, |failure| failure.error())));
461                    }
462                }
463                stop.cancel();
464                result
465            };
466            let receive = async {
467                // Application reads own ordinary receive progress. Teardown
468                // drives cancellation and joins the persistent socket lane.
469                stop.cancelled().await;
470                poll_fn(|cx| receiver.poll(cx, false)).await
471            };
472            let transmit = poll_fn(|cx| transmitter.poll(cx, false));
473            let (result, (read_closed, write_closed)) =
474                join(operation, join(receive, transmit)).await;
475            lock(&shared.state).finished = true;
476            shared.notify_all();
477            result.and_then(|value| {
478                read_closed?;
479                write_closed?;
480                Ok(value)
481            })
482        }
483        .await;
484        let closed = scheduler.socket_close(socket).await;
485        result.and_then(|value| {
486            closed?;
487            Ok(value)
488        })
489    }
490}
491
492async fn receive<R: SocketStreamReader>(
493    shared: &Shared,
494    reader: &mut R,
495    initial: &[u8],
496    cancellation: &Cancellation,
497) -> OperationResult<()> {
498    let mut buffer = vec![0; CAPACITY.max(initial.len())];
499    buffer[..initial.len()].copy_from_slice(initial);
500    let (mut position, mut end) = (0, initial.len());
501    let mut cancelled = std::pin::pin!(cancellation.cancelled());
502    loop {
503        // Once plaintext is available, return it to the consumer before
504        // reading the socket again. The reader polls us again after consuming
505        // that plaintext; teardown polls us after cancellation.
506        poll_fn(|cx| {
507            cancellation.check()?;
508            if lock(&shared.state).plaintext == 0 {
509                return Poll::Ready(Ok(()));
510            }
511            poll_cancelled(cancelled.as_mut(), cx)
512        })
513        .await?;
514        if position == end {
515            let count = reader.read_into(&mut buffer, cancellation).await?;
516            if count == 0 {
517                let mut state = lock(&shared.state);
518                state.tls.read_tls(&mut io::empty())?;
519                drop(state);
520                return Ok(());
521            }
522            position = 0;
523            end = count;
524        }
525        let peer_closed = {
526            let mut state = lock(&shared.state);
527            let count = state.tls.read_tls(&mut &buffer[position..end])?;
528            position += count;
529            let progress = state
530                .tls
531                .process_new_packets()
532                .map_err(|error| io::Error::new(io::ErrorKind::InvalidData, error))?;
533            state.plaintext = progress.plaintext_bytes_to_read();
534            if state.tls.wants_write() {
535                shared.transmit.wake();
536            }
537            progress.peer_has_closed()
538        };
539        if peer_closed {
540            return Ok(());
541        }
542    }
543}
544
545async fn transmit<W: SocketStreamWriter>(
546    shared: &Shared,
547    writer: &mut W,
548    mut buffer: Vec<u8>,
549    mut position: usize,
550    cancellation: &Cancellation,
551) -> OperationResult<Infallible> {
552    let mut end = buffer.len();
553    let mut dirty = position != end;
554    // Keep initialized storage between records. The valid ciphertext range is
555    // tracked separately so small writes do not zero the full capacity again.
556    buffer.resize(buffer.len().max(CAPACITY), 0);
557    let mut cancelled = std::pin::pin!(cancellation.cancelled());
558    loop {
559        cancellation.check()?;
560        if position != end {
561            let count = writer.write(&buffer[position..end], cancellation).await?;
562            if count == 0 {
563                return Err(io::Error::new(
564                    io::ErrorKind::WriteZero,
565                    "TLS output made no progress",
566                )
567                .into());
568            }
569            position += count;
570            dirty = true;
571            continue;
572        }
573        end = 0;
574        position = 0;
575        let ticket = {
576            let mut state = lock(&shared.state);
577            if state.tls.wants_write() {
578                // Rustls reports only destination errors; a memory cursor has none.
579                end = state
580                    .tls
581                    .write_tls(&mut io::Cursor::new(buffer.as_mut_slice()))
582                    .expect("writing TLS records to memory cannot fail");
583            }
584            if end == 0 && (dirty || state.delivered < state.issued) {
585                Some(state.issued)
586            } else {
587                None
588            }
589        };
590        if end != 0 {
591            shared.writable.notify(usize::MAX);
592            continue;
593        }
594        if let Some(ticket) = ticket {
595            writer.flush(cancellation).await?;
596            lock(&shared.state).delivered = ticket;
597            dirty = false;
598            shared.writable.notify(usize::MAX);
599        } else {
600            // Foreground output can poll this future while it is idle. Recheck
601            // demand before waiting, so ordinary writes need no notification.
602            poll_fn(|cx| {
603                shared.transmit.register(cx.waker());
604                cancellation.check()?;
605                let state = lock(&shared.state);
606                if state.tls.wants_write() || state.delivered < state.issued {
607                    return Poll::Ready(Ok(()));
608                }
609                drop(state);
610                poll_cancelled(cancelled.as_mut(), cx)
611            })
612            .await?;
613        }
614    }
615}
base_endpoint/tls/rustls/rustls_connection/duplex.rs

base_endpoint/tls/rustls/rustls_connection/
duplex.rs
