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

base_endpoint/tls/rustls/rustls_connection/
duplex.rs
