Improved OTP listener, call signout correctly in standalone mode

snx-rs/src/main.rs

@@ -2,27 +2,24 @@ use std::{collections::VecDeque, future::Future, sync::Arc};
 
 use clap::Parser;
 use futures::pin_mut;
-use tokio::{
-    signal::unix,
-    sync::{mpsc, oneshot},
-};
+use tokio::{signal::unix, sync::mpsc};
 use tracing::{debug, metadata::LevelFilter, warn};
 
+use crate::cmdline::CmdlineParams;
+use snxcore::model::params::TunnelType;
 use snxcore::{
-    browser::run_otp_listener,
+    browser::spawn_otp_listener,
     ccc::CccHttpClient,
     model::{
         params::{OperationMode, TunnelParams},
         MfaType, SessionState,
     },
     platform,
-    prompt::{SecurePrompt, TtyPrompt, OTP_TIMEOUT},
+    prompt::{SecurePrompt, TtyPrompt},
     server::CommandServer,
     server_info, tunnel,
 };
 
-use crate::cmdline::CmdlineParams;
-
 mod cmdline;
 
 fn is_root() -> bool {
@@ -161,9 +158,8 @@ async fn main_standalone(params: TunnelParams) -> anyhow::Result<()> {
             MfaType::SamlSso => {
                 println!("For SAML authentication open the following URL in your browser:");
                 println!("{}", challenge.prompt);
-                let (tx, rx) = oneshot::channel();
-                tokio::spawn(run_otp_listener(tx));
-                let otp = tokio::time::timeout(OTP_TIMEOUT, rx).await??;
+                let receiver = spawn_otp_listener();
+                let otp = receiver.await??;
                 session = connector.challenge_code(session, &otp).await?;
             }
             MfaType::UserNameInput => {
@@ -173,7 +169,7 @@ async fn main_standalone(params: TunnelParams) -> anyhow::Result<()> {
         }
     }
 
-    let tunnel = connector.create_tunnel(session, command_sender).await?;
+    let tunnel = connector.create_tunnel(session.clone(), command_sender).await?;
 
     if let Err(e) = platform::start_network_state_monitoring().await {
         warn!("Unable to start network monitoring: {}", e);
@@ -198,6 +194,11 @@ async fn main_standalone(params: TunnelParams) -> anyhow::Result<()> {
                 }
             }
             result = &mut tunnel_fut => {
+                if params.tunnel_type == TunnelType::Ssl || !params.ike_persist {
+                    debug!("Signing out");
+                    let client = CccHttpClient::new(params.clone(), Some(session));
+                    let _ = client.signout().await;
+                }
                 break result;
             }
         }

snxcore/src/browser.rs

@@ -1,3 +1,5 @@
+use std::time::Duration;
+
 use anyhow::anyhow;
 use once_cell::sync::Lazy;
 use regex::Regex;
@@ -7,6 +9,8 @@ use tokio::{
     sync::oneshot,
 };
 
+const OTP_TIMEOUT: Duration = Duration::from_secs(120);
+
 pub trait BrowserController {
     fn open(&self, url: &str) -> anyhow::Result<()>;
     fn close(&self);
@@ -22,30 +26,41 @@ impl BrowserController for SystemBrowser {
     fn close(&self) {}
 }
 
-pub async fn run_otp_listener(sender: oneshot::Sender<String>) -> anyhow::Result<()> {
+pub fn spawn_otp_listener() -> oneshot::Receiver<anyhow::Result<String>> {
     static OTP_RE: Lazy<Regex> = Lazy::new(|| Regex::new(r"^GET /(?<otp>[0-9a-f]{60}|[0-9A-F]{60}).*").unwrap());
 
-    let tcp = TcpListener::bind("127.0.0.1:7779").await?;
-    let (mut stream, _) = tcp.accept().await?;
+    let (sender, receiver) = oneshot::channel();
 
-    let mut buf = [0u8; 65];
-    stream.read_exact(&mut buf).await?;
+    let fut = async move {
+        let tcp = TcpListener::bind("127.0.0.1:7779").await?;
+        let (mut stream, _) = tcp.accept().await?;
 
-    let mut data = String::from_utf8_lossy(&buf).into_owned();
+        let mut buf = [0u8; 65];
+        stream.read_exact(&mut buf).await?;
 
-    while stream.read(&mut buf[0..1]).await.is_ok() && buf[0] != b'\n' && buf[0] != b'\r' {
-        data.push(buf[0].into());
-    }
+        let mut data = String::from_utf8_lossy(&buf).into_owned();
+
+        while stream.read(&mut buf[0..1]).await.is_ok() && buf[0] != b'\n' && buf[0] != b'\r' {
+            data.push(buf[0].into());
+        }
 
-    let _ = stream.shutdown().await;
-    drop(stream);
-    drop(tcp);
+        let _ = stream.shutdown().await;
 
-    if let Some(captures) = OTP_RE.captures(&data) {
-        if let Some(otp) = captures.name("otp") {
-            let _ = sender.send(otp.as_str().to_owned());
-            return Ok(());
+        if let Some(captures) = OTP_RE.captures(&data) {
+            if let Some(otp) = captures.name("otp") {
+                return Ok(otp.as_str().to_owned());
+            }
         }
-    }
-    Err(anyhow!("No OTP acquired!"))
+        Err(anyhow!("Invalid OTP reply"))
+    };
+
+    tokio::spawn(async move {
+        let result = tokio::time::timeout(OTP_TIMEOUT, fut)
+            .await
+            .unwrap_or_else(|e| Err(e.into()));
+
+        let _ = sender.send(result);
+    });
+
+    receiver
 }

snxcore/src/controller.rs

@@ -1,17 +1,16 @@
 use std::{collections::VecDeque, str::FromStr, sync::Arc, time::Duration};
 
 use anyhow::anyhow;
-use tokio::sync::oneshot;
 use tracing::warn;
 
 use crate::{
-    browser::{run_otp_listener, BrowserController},
+    browser::{spawn_otp_listener, BrowserController},
     ccc::CccHttpClient,
     model::{
         params::TunnelParams, ConnectionStatus, MfaChallenge, MfaType, TunnelServiceRequest, TunnelServiceResponse,
     },
     platform::{self, UdpSocketExt},
-    prompt::{SecurePrompt, OTP_TIMEOUT},
+    prompt::SecurePrompt,
     server_info,
 };
 
@@ -148,18 +147,17 @@ where
                 result
             }
             MfaType::SamlSso => {
-                let (tx, rx) = oneshot::channel();
-                tokio::spawn(run_otp_listener(tx));
+                let receiver = spawn_otp_listener();
 
                 self.browser_controller.open(&mfa.prompt)?;
 
-                match tokio::time::timeout(OTP_TIMEOUT, rx).await {
+                match receiver.await {
                     Ok(Ok(otp)) => {
                         self.browser_controller.close();
                         Ok(otp)
                     }
-                    _ => {
-                        warn!("Unable to acquire OTP from the browser");
+                    other => {
+                        warn!("Unable to acquire OTP from the browser: {:?}", other);
                         Err(anyhow!("Unable to acquire OTP from the browser!"))
                     }
                 }

snxcore/src/prompt.rs

@@ -1,11 +1,6 @@
 use anyhow::anyhow;
 use std::io::Write;
-use std::{
-    io::{stderr, stdin, IsTerminal},
-    time::Duration,
-};
-
-pub const OTP_TIMEOUT: Duration = Duration::from_secs(120);
+use std::io::{stderr, stdin, IsTerminal};
 
 pub trait SecurePrompt {
     fn get_secure_input(&self, prompt: &str) -> anyhow::Result<String>;