sunset_async/
async_sunset.rs

1#[allow(unused_imports)]
2pub use log::{debug, error, info, log, trace, warn};
3
4use core::future::{Future, poll_fn};
5use core::pin::pin;
6use core::sync::atomic::AtomicBool;
7use core::sync::atomic::Ordering::{AcqRel, Acquire, Relaxed};
8use core::task::{Context, Poll, Poll::Pending, Poll::Ready};
9
10// thumbv6m has no atomic usize add/sub.
11use portable_atomic::AtomicUsize;
12
13use embassy_futures::join;
14use embassy_futures::select::select;
15#[allow(unused_imports)]
16use embassy_sync::blocking_mutex::raw::{CriticalSectionRawMutex, NoopRawMutex};
17use embassy_sync::mutex::{Mutex, MutexGuard};
18use embassy_sync::signal::Signal;
19use embedded_io_async::{Read, Write};
20
21use crate::async_channel::ChanIO;
22use sunset::ChanData::{Normal, Stderr};
23use sunset::config::MAX_CHANNELS;
24use sunset::error::TrapBug;
25use sunset::event::Event;
26use sunset::{ChanData, ChanHandle, ChanNum, CliServ, Error, Result, Runner, error};
27
28/// A raw mutex
29///
30/// This is the [`RawMutex`](embassy_sync::blocking_mutex::raw::RawMutex)
31/// type used internally by `sunset-async`.
32/// When `multi-thread` feature for is enabled it will use
33/// `embassy-sync`'s [`CriticalSectionRawMutex`], otherwise it will
34/// use [`NoopRawMutex`] (no locking is required for single threaded).
35#[cfg(feature = "multi-thread")]
36pub type SunsetRawMutex = CriticalSectionRawMutex;
37
38/// A raw mutex
39///
40/// This is the [`RawMutex`](embassy_sync::blocking_mutex::raw::RawMutex)
41/// type used internally by `sunset-async`.
42/// When `multi-thread` feature for is enabled it will use
43/// `embassy-sync`'s [`CriticalSectionRawMutex`], otherwise it will
44/// use [`NoopRawMutex`] (no locking is required for single threaded).
45///
46/// Applications may use this for their own `embassy-sync` data structures.
47#[cfg(not(feature = "multi-thread"))]
48pub type SunsetRawMutex = NoopRawMutex;
49
50/// An async mutex
51///
52/// This is the [`Mutex`](embassy_sync::mutex::Mutex) type used internally
53/// by `sunset-async`.
54/// When `multi-thread` feature is enabled it will use
55/// `embassy-sync`'s [`CriticalSectionRawMutex`], otherwise it will
56/// use [`NoopRawMutex`] (no locking is required for single threaded).
57pub type SunsetMutex<T> = Mutex<SunsetRawMutex, T>;
58
59struct Inner<'a, CS: CliServ> {
60    runner: Runner<'a, CS>,
61
62    // May only be safely modified when the corresponding
63    // `chan_refcounts` is zero.
64    chan_handles: [Option<ChanHandle>; MAX_CHANNELS],
65}
66
67impl<'a, CS: CliServ> Inner<'a, CS> {
68    /// Helper to lookup the corresponding ChanHandle
69    ///
70    /// Returns split references that will be required by many callers
71    fn fetch(&mut self, num: ChanNum) -> Result<(&mut Runner<'a, CS>, &ChanHandle)> {
72        let ch = self
73            .chan_handles
74            .get(num.0 as usize)
75            .ok_or(Error::BadChannel { num })?
76            .as_ref()
77            .trap()?;
78        Ok((&mut self.runner, ch))
79    }
80}
81
82/// A handle used for storage from a [`SSHClient::progress()`](crate::SSHClient::progress)
83/// or [`SSHServer::progress()`](crate::SSHServer::progress) call.
84pub struct ProgressHolder<'g, 'a, CS: CliServ> {
85    guard: Option<MutexGuard<'g, SunsetRawMutex, Inner<'a, CS>>>,
86}
87
88impl<'g, 'a, CS: CliServ> ProgressHolder<'g, 'a, CS> {
89    pub fn new() -> Self {
90        Self { guard: None }
91    }
92}
93
94impl<CS: CliServ> Default for ProgressHolder<'_, '_, CS> {
95    fn default() -> Self {
96        Self::new()
97    }
98}
99
100/// Provides an async wrapper for Sunset core
101///
102/// A [`ChanHandle`] provided by sunset core must be added with [`add_channel()`] before
103/// a method can be called with the equivalent ChanNum.
104///
105/// Applications use `async_sunset::{Client,Server}`.
106pub(crate) struct AsyncSunset<'a, CS: CliServ> {
107    inner: SunsetMutex<Inner<'a, CS>>,
108
109    progress_notify: Signal<SunsetRawMutex, ()>,
110    last_progress_idled: AtomicBool,
111
112    // wake_progress() should be called after modifying these atomics, to
113    // trigger the progress loop to handle state changes
114
115    // When draining the last events
116    moribund: AtomicBool,
117
118    // Refcount for `Inner::chan_handles`. Must be non-async so it can be
119    // decremented on `ChanIn::drop()` etc.
120    // The pending chan_refcount=0 handling occurs in the `progress()` loop.
121    chan_refcounts: [AtomicUsize; MAX_CHANNELS],
122
123    /// Refcount for Normal ChanIn or ChanInOut.
124    ///
125    /// Used to discard incoming data when none are remaining.
126    chan_norm_readcounts: [AtomicUsize; MAX_CHANNELS],
127    /// Refcount for Stderr ChanIn or ChanInOut.
128    chan_stderr_readcounts: [AtomicUsize; MAX_CHANNELS],
129}
130
131impl<CS: CliServ> core::fmt::Debug for AsyncSunset<'_, CS> {
132    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
133        let mut d = f.debug_struct("AsyncSunset");
134        if let Ok(i) = self.inner.try_lock() {
135            d.field("runner", &i.runner);
136        } else {
137            d.field("inner", &"(locked)");
138        }
139        d.finish_non_exhaustive()
140    }
141}
142
143impl<'a, CS: CliServ> AsyncSunset<'a, CS> {
144    pub fn new(runner: Runner<'a, CS>) -> Self {
145        let inner = Inner { runner, chan_handles: Default::default() };
146        let inner = Mutex::new(inner);
147
148        let progress_notify = Signal::new();
149
150        Self {
151            inner,
152            moribund: AtomicBool::new(false),
153            progress_notify,
154            chan_refcounts: Default::default(),
155            chan_norm_readcounts: Default::default(),
156            chan_stderr_readcounts: Default::default(),
157            last_progress_idled: AtomicBool::new(false),
158        }
159    }
160
161    /// Runs the session to completion
162    pub async fn run(
163        &self,
164        rsock: &mut impl Read,
165        wsock: &mut impl Write,
166    ) -> Result<()> {
167        // Some loops need to terminate other loops on completion.
168        // prog finish -> stop rx
169        // rx finish -> stop tx
170        let tx_stop = Signal::<SunsetRawMutex, ()>::new();
171        let rx_stop = Signal::<SunsetRawMutex, ()>::new();
172
173        let tx = async {
174            self.output_loop(wsock).await.inspect(|r| warn!("tx complete {r:?}"))
175        };
176        let tx = select(tx, tx_stop.wait());
177
178        // rxbuf outside the async block avoids an extraneous copy somehow
179        let mut rxbuf = [0; 1024];
180        let rx = async {
181            loop {
182                // TODO: make sunset read directly from socket, no intermediate buffer.
183                let l = match rsock.read(&mut rxbuf).await {
184                    Ok(0) => {
185                        debug!("net EOF");
186                        self.with_runner(|r| r.close_input()).await;
187                        self.moribund.store(true, Relaxed);
188                        self.wake_progress();
189                        break Ok(());
190                    }
191                    Ok(l) => l,
192                    Err(_) => {
193                        info!("socket read error");
194                        self.with_runner(|r| r.close_input()).await;
195                        break Err(Error::ChannelEOF);
196                    }
197                };
198                let mut rxbuf = &rxbuf[..l];
199                while !rxbuf.is_empty() {
200                    let n = self.input(rxbuf).await?;
201                    self.wake_progress();
202                    rxbuf = &rxbuf[n..];
203                }
204            }
205            .inspect(|r| warn!("rx complete {r:?}"))
206        };
207
208        // TODO: if RX fails (bad decrypt etc) it doesn't cancel prog, so gets stuck
209        let rx = async {
210            let r = select(rx, rx_stop.wait()).await;
211            tx_stop.signal(());
212            r
213        };
214
215        // TODO: we might want to let `prog` run until buffers are drained
216        // in case a disconnect message was received.
217        // TODO Is there a nice way than this?
218        let f = join::join(rx, tx).await;
219        let (_frx, _ftx) = f;
220
221        // debug!("frx {_frx:?}");
222        // debug!("ftx {_ftx:?}");
223
224        // TODO: is this a good way to do cancellation...?
225        // self.with_runner(|runner| runner.close()).await;
226        // // Wake any channels that were awoken after the runner closed
227        // let mut inner = self.inner.lock().await;
228        // self.wake_channels(&mut inner)?;
229        Ok(())
230    }
231
232    fn wake_progress(&self) {
233        trace!("wake_progress");
234        self.progress_notify.signal(())
235    }
236
237    fn discard_channels(&self, inner: &mut Inner<CS>) -> Result<()> {
238        if let Some((num, dt, _len)) = inner.runner.read_channel_ready() {
239            if self.chan_readcount(num, dt).load(Acquire) == 0 {
240                // There are no live ChanIn or ChanInOut for the num/dt,
241                // so nothing will read the channel.
242                // Discard the data so it doesn't block forever.
243                let ch = inner.chan_handles[num.0 as usize].as_ref().trap()?;
244                inner.runner.discard_read_channel(ch)?;
245            }
246        }
247        Ok(())
248    }
249
250    /// Check for channels that have reached zero refcount
251    ///
252    /// When a ChanIO is dropped the refcount may reach 0, but
253    /// without "async Drop" it isn't possible to take the `inner` lock during
254    /// `drop()`.
255    /// Instead this runs periodically from an async context to release channels.
256    fn clear_refcounts(&self, inner: &mut Inner<CS>) -> Result<()> {
257        for (ch, count) in
258            inner.chan_handles.iter_mut().zip(self.chan_refcounts.iter())
259        {
260            let count = count.load(Acquire);
261            if count > 0 {
262                debug_assert!(ch.is_some());
263                continue;
264            }
265            if let Some(ch) = ch.take() {
266                // done with the channel
267                inner.runner.channel_done(ch)?;
268            }
269        }
270        Ok(())
271    }
272
273    /// Returns an `Event`.
274    ///
275    /// The returned `Event` borrows from the mutex locked in `ph`.
276    pub(crate) async fn progress<'g, 'f>(
277        &'g self,
278        ph: &'f mut ProgressHolder<'g, 'a, CS>,
279    ) -> Result<Event<'f, 'a>> {
280        // In case a ProgressHolder was reused, release any guard.
281        *ph = ProgressHolder::default();
282
283        // Ideally we would .wait() after calling .progress() below when
284        // Event::None is returned, but the borrow checker won't allow that.
285        // Instead we wait at the start of the next progress() call,
286        // but will return immediately if something external
287        // has woken the progress_notify in the interim.
288        //
289        // TODO: rework once rustc's polonius is stable.
290        // https://github.com/rust-lang/rust/issues/54663
291        //
292        // This is a non-atomic swap since thumbv6m won't support it.
293        // Only one task should be calling progress(), so that's OK.
294        let need_wait = self.last_progress_idled.load(Relaxed);
295        if need_wait {
296            self.last_progress_idled.store(false, Relaxed);
297            self.progress_notify.wait().await;
298        }
299
300        // The returned event borrows from a guard inside ProgressHolder
301        let inner = ph.guard.insert(self.inner.lock().await);
302
303        // Drop deferred finished channels
304        self.clear_refcounts(inner)?;
305        // Discard unhandled input
306        self.discard_channels(inner)?;
307
308        if self.moribund.load(Relaxed) {
309            // if we're flushing, we exit once there is no progress
310            debug!("All data flushed")
311            // TODO make this do something!
312        }
313
314        let ev = inner.runner.progress();
315        if matches!(ev, Ok(Event::None)) {
316            // nothing happened, will progress_notify.wait() next progress() call, see above.
317            self.last_progress_idled.store(true, Relaxed);
318        }
319        ev
320    }
321
322    pub(crate) async fn with_runner<F, R>(&self, f: F) -> R
323    where
324        F: FnOnce(&mut Runner<CS>) -> R,
325    {
326        let mut inner = self.inner.lock().await;
327        f(&mut inner.runner)
328    }
329
330    /// Fetch the relevant atomic counter
331    fn chan_readcount(&self, num: ChanNum, dt: ChanData) -> &AtomicUsize {
332        let counts = match dt {
333            Normal => &self.chan_norm_readcounts,
334            Stderr => &self.chan_stderr_readcounts,
335        };
336        &counts[num.0 as usize]
337    }
338
339    /// helper to perform a function on the `inner`, returning a `Poll` value
340    async fn poll_inner<F, T>(&self, mut f: F) -> T
341    where
342        F: FnMut(&mut Inner<CS>, &mut Context) -> Poll<T>,
343    {
344        poll_fn(|cx| {
345            // Attempt to lock .inner
346            let i = self.inner.lock();
347            let i = pin!(i);
348            match i.poll(cx) {
349                Poll::Ready(mut inner) => f(&mut inner, cx),
350                Poll::Pending => {
351                    // .inner lock is busy
352                    Poll::Pending
353                }
354            }
355        })
356        .await
357    }
358
359    pub async fn output_loop(&self, wsock: &mut impl Write) -> Result<()> {
360        poll_fn(|cx| {
361            // Attempt to lock .inner
362            let i = self.inner.lock();
363            let i = pin!(i);
364            let Ready(mut inner) = i.poll(cx) else {
365                return Pending;
366            };
367
368            loop {
369                let buf = inner.runner.output_buf();
370                if buf.is_empty() {
371                    // no output ready
372                    inner.runner.set_output_waker(cx.waker());
373                    return Pending;
374                }
375
376                let res = {
377                    let w = wsock.write(buf);
378                    let w = pin!(w);
379                    w.poll(cx)
380                };
381
382                let r = match res {
383                    Pending => {
384                        // wsock has set a waker
385                        Pending
386                    }
387                    Ready(Ok(0)) => {
388                        info!("socket EOF");
389                        inner.runner.close_output();
390                        Ready(error::ChannelEOF.fail())
391                    }
392                    Ready(Ok(write_len)) => {
393                        let buf_len = buf.len();
394                        inner.runner.consume_output(write_len);
395                        if write_len < buf_len {
396                            // Must keep going until either wsock
397                            // or output_buf returns Pending and
398                            // registers a waker.
399                            continue;
400                        }
401                        inner.runner.set_output_waker(cx.waker());
402                        if !inner.runner.is_output_pending() {
403                            // All output was sent. Wake progress
404                            // in case window adjustments etc need to be sent
405                            // now that there is available space.
406                            self.wake_progress();
407                        }
408                        Pending
409                    }
410                    Ready(Err(_e)) => {
411                        info!("socket write error");
412                        inner.runner.close_output();
413                        Ready(error::ChannelEOF.fail())
414                    }
415                };
416                return r;
417            }
418        })
419        .await
420    }
421
422    pub async fn input(&self, buf: &[u8]) -> Result<usize> {
423        let res = self
424            .poll_inner(|inner, cx| {
425                if inner.runner.is_input_ready() {
426                    match inner.runner.input(buf) {
427                        Ok(0) => {
428                            inner.runner.set_input_waker(cx.waker());
429                            Poll::Pending
430                        }
431                        Ok(n) => Poll::Ready(Ok(n)),
432                        Err(e) => Poll::Ready(Err(e)),
433                    }
434                } else {
435                    inner.runner.set_input_waker(cx.waker());
436                    Poll::Pending
437                }
438            })
439            .await;
440        self.wake_progress();
441        res
442    }
443
444    /// Adds a new channel handle provided by sunset core.
445    ///
446    /// AsyncSunset will take ownership of the handle.
447    ///
448    /// The channel will have an initial refcount of 1 for the
449    /// returned ChanIO.
450    /// chan_norm_readcounts and chan_stderr_readcounts are initially
451    /// 0, will be set by ChanIn or ChanInOut.
452    ///
453    /// ChanIO will take care of `inc_chan()` on clone, `dec_chan()` on drop.
454    pub(crate) async fn add_channel(
455        &self,
456        handle: ChanHandle,
457    ) -> Result<ChanIO<'_>> {
458        let mut inner = self.inner.lock().await;
459        let num = handle.num();
460        let idx = num.0 as usize;
461        if inner.chan_handles[idx].is_some() {
462            return error::Bug.fail();
463        }
464        inner.chan_handles[idx] = Some(handle);
465
466        debug_assert_eq!(self.chan_refcounts[idx].load(Relaxed), 0);
467        self.chan_refcounts[idx].store(1, Relaxed);
468        Ok(ChanIO::new_normal(num, self))
469    }
470}
471
472// necessary for the &dyn ChanCore
473#[cfg(feature = "multi-thread")]
474pub(crate) trait MaybeSend: Sync {}
475#[cfg(not(feature = "multi-thread"))]
476pub(crate) trait MaybeSend {}
477
478impl<'a, CS: CliServ> MaybeSend for AsyncSunset<'a, CS> {}
479
480// Ideally the poll_...() methods would be async, but that isn't
481// dyn compatible at present. Instead run poll_fn in the ChanIO caller.
482pub(crate) trait ChanCore: MaybeSend {
483    fn inc_chan(&self, num: ChanNum);
484    fn dec_chan(&self, num: ChanNum);
485    fn inc_read_chan(&self, num: ChanNum, dt: ChanData);
486    fn dec_read_chan(&self, num: ChanNum, dt: ChanData);
487
488    fn poll_until_channel_closed(
489        &self,
490        cx: &mut Context,
491        num: ChanNum,
492    ) -> Poll<Result<()>>;
493
494    fn poll_read_channel(
495        &self,
496        cx: &mut Context,
497        num: ChanNum,
498        dt: ChanData,
499        buf: &mut [u8],
500    ) -> Poll<Result<usize>>;
501
502    fn poll_write_channel(
503        &self,
504        cx: &mut Context,
505        num: ChanNum,
506        dt: ChanData,
507        buf: &[u8],
508    ) -> Poll<Result<usize>>;
509
510    // Client only
511    fn poll_term_window_change(
512        &self,
513        cx: &mut Context,
514        num: ChanNum,
515        winch: &sunset::packets::WinChange,
516    ) -> Poll<Result<()>>;
517}
518
519impl<'a, CS: CliServ> ChanCore for AsyncSunset<'a, CS> {
520    /// Counts live ChanIO instances
521    fn inc_chan(&self, num: ChanNum) {
522        // Relaxed is OK, doesn't perform any action until later decrement.
523        let c = self.chan_refcounts[num.0 as usize].fetch_add(1, Relaxed);
524        debug_assert_ne!(c, 0);
525        // overflow shouldn't be possible unless ChanIn etc is leaking
526        debug_assert_ne!(c, usize::MAX);
527    }
528
529    /// Counts live ChanIO instances
530    fn dec_chan(&self, num: ChanNum) {
531        // refcounts that hit zero will be cleaned up later in clear_refcounts()
532        let c = self.chan_refcounts[num.0 as usize].fetch_sub(1, AcqRel);
533        debug_assert_ne!(c, 0);
534        if c == 1 {
535            // refcount hit zero, progress() will clean it up
536            // in an async context
537            self.wake_progress();
538        }
539    }
540
541    /// Counts live ChanIn or ChanInOut instances
542    fn inc_read_chan(&self, num: ChanNum, dt: ChanData) {
543        let c = self.chan_readcount(num, dt).fetch_add(1, AcqRel);
544        debug_assert_ne!(c, usize::MAX);
545    }
546
547    /// Counts live ChanIn or ChanInOut instances
548    fn dec_read_chan(&self, num: ChanNum, dt: ChanData) {
549        let c = self.chan_readcount(num, dt).fetch_sub(1, AcqRel);
550        debug_assert_ne!(c, 0);
551        if c == 1 {
552            // refcount hit zero, wake progress so that any data already
553            // pending will get discarded (by wake_channels()).
554            self.wake_progress();
555        }
556    }
557
558    fn poll_until_channel_closed(
559        &self,
560        cx: &mut Context,
561        num: ChanNum,
562    ) -> Poll<Result<()>> {
563        // Attempt to lock .inner
564        let i = self.inner.lock();
565        let i = pin!(i);
566        let Ready(mut inner) = i.poll(cx) else {
567            return Pending;
568        };
569
570        let (runner, h) = inner.fetch(num)?;
571        if runner.is_channel_closed(h) {
572            Poll::Ready(Ok(()))
573        } else {
574            // read Normal is arbitrary, any read or write should get woken on close
575            runner.set_channel_read_waker(h, Normal, cx.waker());
576            Poll::Pending
577        }
578    }
579
580    /// Reads channel data.
581    fn poll_read_channel(
582        &self,
583        cx: &mut Context,
584        num: ChanNum,
585        dt: ChanData,
586        buf: &mut [u8],
587    ) -> Poll<Result<usize>> {
588        // Attempt to lock .inner
589        let i = self.inner.lock();
590        let i = pin!(i);
591        let Ready(mut inner) = i.poll(cx) else {
592            return Pending;
593        };
594
595        let (runner, h) = inner.fetch(num)?;
596        let i = match runner.read_channel(h, dt, buf) {
597            Ok(0) => {
598                // 0 bytes read, pending
599                trace!("read ch {num:?} dt {dt:?} pending");
600                runner.set_channel_read_waker(h, dt, cx.waker());
601                Poll::Pending
602            }
603            Err(Error::ChannelEOF) => Poll::Ready(Ok(0)),
604            r => {
605                trace!("read ready ch {num:?} dt {dt:?} {r:?}");
606                Poll::Ready(r)
607            }
608        };
609        if matches!(i, Poll::Ready(_)) {
610            self.wake_progress()
611        }
612        i
613    }
614
615    fn poll_write_channel(
616        &self,
617        cx: &mut Context,
618        num: ChanNum,
619        dt: ChanData,
620        buf: &[u8],
621    ) -> Poll<Result<usize>> {
622        if buf.is_empty() {
623            return Poll::Ready(Ok(0));
624        }
625
626        // Attempt to lock .inner
627        let i = self.inner.lock();
628        let i = pin!(i);
629        let Ready(mut inner) = i.poll(cx) else {
630            return Pending;
631        };
632
633        let (runner, h) = inner.fetch(num)?;
634        let l = runner.write_channel(h, dt, buf);
635        if let Ok(0) = l {
636            // 0 bytes written, pending
637            trace!("write ch {num:?} dt {dt:?} pending");
638            runner.set_channel_write_waker(h, dt, cx.waker());
639            Poll::Pending
640        } else {
641            trace!("write ready ch {num:?} dt {dt:?} {l:?}");
642            self.wake_progress();
643            Poll::Ready(l)
644        }
645    }
646
647    fn poll_term_window_change(
648        &self,
649        cx: &mut Context,
650        num: ChanNum,
651        winch: &sunset::packets::WinChange,
652    ) -> Poll<Result<()>> {
653        // Attempt to lock .inner
654        let i = self.inner.lock();
655        let i = pin!(i);
656        let Ready(mut inner) = i.poll(cx) else {
657            return Pending;
658        };
659        let (runner, h) = inner.fetch(num)?;
660        Poll::Ready(runner.term_window_change(h, winch))
661    }
662}
sunset_async/async_sunset.rs

sunset_async/
async_sunset.rs
