sunset_async/
async_sunset.rs

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

sunset_async/
async_sunset.rs
