From a0f6e5ffd01660d1c8b7507251cce32cfbeda748 Mon Sep 17 00:00:00 2001
From: Matt Johnston <matt@ucc.asn.au>
Date: Sun, 2 Jun 2024 22:19:22 +0800
Subject: [PATCH] Return exit codes from sunsetc

---
 async/examples/sunsetc.rs   | 34 ++++++++++++++--------------------
 async/src/cmdline_client.rs | 24 +++++++++++++-----------
 src/lib.rs                  |  2 +-
 3 files changed, 28 insertions(+), 32 deletions(-)

diff --git a/async/examples/sunsetc.rs b/async/examples/sunsetc.rs
index 8dd0a79..e7f8bf4 100644
--- a/async/examples/sunsetc.rs
+++ b/async/examples/sunsetc.rs
@@ -33,7 +33,8 @@ async fn real_main(tz: UtcOffset) -> Result<()> {
     // (or something wrapping std::sync::Mutex).
     // Need to figure how to make it configurable.
     let local = tokio::task::LocalSet::new();
-    local.run_until(run(args)).await
+    let exit_code = local.run_until(run(args)).await?;
+    std::process::exit(exit_code)
 }
 
 fn main() {
@@ -47,7 +48,7 @@ fn main() {
     }
 }
 
-async fn run(args: Args) -> Result<()> {
+async fn run(args: Args) -> Result<i32> {
 
     trace!("tracing sunsetc. args {:?}", args);
     debug!("verbose sunsetc");
@@ -109,27 +110,20 @@ async fn run(args: Args) -> Result<()> {
         let mut wsock = FromTokio::new(wsock);
 
         // SSH connection future
-        let ssh_fut = async {
-            let r = ssh.run(&mut rsock, &mut wsock).await;
-            trace!("ssh run finished {r:?}");
-            // TODO split
-            // hooks.lock().await.exited().await;
-            r
-        };
+        let ssh_fut = ssh.run(&mut rsock, &mut wsock);
 
         // Client session future
-        let session = async {
-            let r = app.run(&ssh).await;
-            trace!("client session run finished {r:?}");
-            r
+        let session = app.run(&ssh);
+
+        let exit_code = tokio::select! {
+            a = ssh_fut => {
+                a.map(|_| 1).context("SSH connection exited")?
+            }
+            a = session => {
+                a.context("client session exited")?
+            }
         };
-
-        let (res_ssh, res_session) = futures::future::join(ssh_fut, session).await;
-        debug!("res_ssh {res_ssh:?}");
-        debug!("res_session {res_session:?}");
-        res_ssh?;
-        res_session?;
-        Ok::<_, anyhow::Error>(())
+        Ok(exit_code)
     });
 
     ssh_task.await.context("Sunset task panicked")?
diff --git a/async/src/cmdline_client.rs b/async/src/cmdline_client.rs
index 7cb040f..8f9a540 100644
--- a/async/src/cmdline_client.rs
+++ b/async/src/cmdline_client.rs
@@ -229,7 +229,7 @@ impl CmdlineClient {
     ///
     /// Performs authentication, requests a shell or command, performs channel IO.
     /// Will return `Ok` after the session ends normally, or an error.
-    pub async fn run<'g: 'a, 'a>(&mut self, cli: &'g SSHClient<'a>) -> Result<()> {
+    pub async fn run<'g: 'a, 'a>(&mut self, cli: &'g SSHClient<'a>) -> Result<i32> {
 
         let mut winch_signal = self.want_pty
             .then(|| signal(SignalKind::window_change()))
@@ -245,6 +245,8 @@ impl CmdlineClient {
         let launch_chan: Channel::<SunsetRawMutex, (ChanInOut, Option<ChanIn>, Option<RawPtyGuard>), 1>
             = Channel::new();
 
+        let mut exit_code = 1i32;
+
         let prog_loop = async {
             loop {
                 let winch_fut = Fuse::terminated();
@@ -304,7 +306,14 @@ impl CmdlineClient {
                         launch_chan.send((io.clone().unwrap(), extin.clone(), self.pty_guard.take())).await;
                     }
                     CliEvent::SessionExit(ex) => {
-                        debug!("TODO handle session exit {ex:?}")
+                        trace!("session exit {ex:?}");
+                        if let sunset::CliSessionExit::Status(u) = ex {
+                            if u <= 255 {
+                                exit_code = i8::from_be_bytes([(u & 0xff) as u8]) as i32;
+                            } else {
+                                exit_code = 1;
+                            }
+                        }
                     }
                     CliEvent::Defunct => {
                         trace!("break defunct");
@@ -314,21 +323,14 @@ impl CmdlineClient {
             }
         };
 
-        let prog_loop = async {
-            let e = prog_loop.await;
-            debug!("loop done, {e:#?}");
-            e
-        };
-
         let chanio = async {
             let (io, extin, pty) = launch_chan.receive().await;
             Self::chan_run(io, extin, pty).await
         };
 
-        // embassy_futures::select::select(prog_loop, chanio).await;
-        let _ = embassy_futures::join::join(prog_loop, chanio).await;
+        embassy_futures::select::select(prog_loop, chanio).await;
 
-        Ok(())
+        Ok(exit_code)
     }
 
     /// Requests a PTY or non-PTY session
diff --git a/src/lib.rs b/src/lib.rs
index b3b2db6..9a944f5 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -58,7 +58,7 @@ pub use packets::{PubKey, Signature};
 pub use error::{Error,Result};
 pub use channel::{Pty, ChanOpened, SessionCommand};
 pub use sshnames::ChanFail;
-pub use channel::{ChanData, ChanNum};
+pub use channel::{ChanData, ChanNum, CliSessionExit};
 pub use auth::AuthSigMsg;
 
 pub use runner::Runner;