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

sunset_async/
async_sunset.rs
