feat: gpauth support macos

crates/auth/src/webview/auth_messenger.rs

@@ -0,0 +1,229 @@
+use anyhow::bail;
+use gpapi::{auth::SamlAuthData, error::AuthDataParseError};
+use log::{error, info};
+use regex::Regex;
+use tokio::sync::{mpsc, RwLock};
+use tokio_util::sync::CancellationToken;
+
+#[derive(Debug)]
+pub(crate) enum AuthDataLocation {
+  #[cfg(not(target_os = "macos"))]
+  Headers,
+  Body,
+}
+
+#[derive(Debug)]
+pub(crate) enum AuthError {
+  /// Failed to load page due to TLS error
+  #[cfg(not(target_os = "macos"))]
+  TlsError,
+  /// 1. Found auth data in headers/body but it's invalid
+  /// 2. Loaded an empty page, failed to load page. etc.
+  Invalid(anyhow::Error, AuthDataLocation),
+  /// No auth data found in headers/body
+  NotFound(AuthDataLocation),
+}
+
+impl AuthError {
+  pub fn invalid_from_body(err: anyhow::Error) -> Self {
+    Self::Invalid(err, AuthDataLocation::Body)
+  }
+
+  pub fn not_found_in_body() -> Self {
+    Self::NotFound(AuthDataLocation::Body)
+  }
+}
+
+#[cfg(not(target_os = "macos"))]
+impl AuthError {
+  pub fn not_found_in_headers() -> Self {
+    Self::NotFound(AuthDataLocation::Headers)
+  }
+}
+
+pub(crate) enum AuthEvent {
+  Data(SamlAuthData, AuthDataLocation),
+  Error(AuthError),
+  RaiseWindow,
+  Close,
+}
+
+pub struct AuthMessenger {
+  tx: mpsc::UnboundedSender<AuthEvent>,
+  rx: RwLock<mpsc::UnboundedReceiver<AuthEvent>>,
+  raise_window_cancel_token: RwLock<Option<CancellationToken>>,
+}
+
+impl AuthMessenger {
+  pub fn new() -> Self {
+    let (tx, rx) = mpsc::unbounded_channel();
+
+    Self {
+      tx,
+      rx: RwLock::new(rx),
+      raise_window_cancel_token: Default::default(),
+    }
+  }
+
+  pub async fn subscribe(&self) -> anyhow::Result<AuthEvent> {
+    let mut rx = self.rx.write().await;
+    if let Some(event) = rx.recv().await {
+      return Ok(event);
+    }
+    bail!("Failed to receive auth event");
+  }
+
+  pub fn send_auth_event(&self, event: AuthEvent) {
+    if let Err(event) = self.tx.send(event) {
+      error!("Failed to send auth event: {}", event);
+    }
+  }
+
+  pub fn send_auth_error(&self, err: AuthError) {
+    self.send_auth_event(AuthEvent::Error(err));
+  }
+
+  fn send_auth_data(&self, data: SamlAuthData, location: AuthDataLocation) {
+    self.send_auth_event(AuthEvent::Data(data, location));
+  }
+
+  pub fn schedule_raise_window(&self, delay: u64) {
+    let Ok(mut guard) = self.raise_window_cancel_token.try_write() else {
+      return;
+    };
+
+    // Return if the previous raise window task is still running
+    if let Some(token) = guard.as_ref() {
+      if !token.is_cancelled() {
+        info!("Raise window task is still running, skipping...");
+        return;
+      }
+    }
+
+    let cancel_token = CancellationToken::new();
+    let cancel_token_clone = cancel_token.clone();
+
+    *guard = Some(cancel_token_clone);
+
+    let tx = self.tx.clone();
+    tokio::spawn(async move {
+      info!("Displaying the window in {} second(s)...", delay);
+
+      tokio::select! {
+        _ = tokio::time::sleep(tokio::time::Duration::from_secs(delay)) => {
+          cancel_token.cancel();
+
+          if let Err(err) = tx.send(AuthEvent::RaiseWindow) {
+            error!("Failed to send raise window event: {}", err);
+          }
+        }
+        _ = cancel_token.cancelled() => {
+          info!("Cancelled raise window task");
+        }
+      }
+    });
+  }
+
+  pub fn cancel_raise_window(&self) {
+    if let Ok(mut cancel_token) = self.raise_window_cancel_token.try_write() {
+      if let Some(token) = cancel_token.take() {
+        token.cancel();
+      }
+    }
+  }
+
+  pub fn read_from_html(&self, html: &str) {
+    if html.contains("Temporarily Unavailable") {
+      return self.send_auth_error(AuthError::invalid_from_body(anyhow::anyhow!("Temporarily Unavailable")));
+    }
+
+    let auth_result = SamlAuthData::from_html(html).or_else(|err| {
+      info!("Read auth data from html failed: {}, extracting gpcallback...", err);
+
+      if let Some(gpcallback) = extract_gpcallback(html) {
+        info!("Found gpcallback from html...");
+        SamlAuthData::from_gpcallback(&gpcallback)
+      } else {
+        Err(err)
+      }
+    });
+
+    match auth_result {
+      Ok(data) => self.send_auth_data(data, AuthDataLocation::Body),
+      Err(AuthDataParseError::Invalid(err)) => self.send_auth_error(AuthError::invalid_from_body(err)),
+      Err(AuthDataParseError::NotFound) => self.send_auth_error(AuthError::not_found_in_body()),
+    }
+  }
+
+  #[cfg(not(target_os = "macos"))]
+  pub fn read_from_response(&self, auth_response: &impl super::webview_auth::GetHeader) {
+    use log::warn;
+
+    let Some(status) = auth_response.get_header("saml-auth-status") else {
+      return self.send_auth_error(AuthError::not_found_in_headers());
+    };
+
+    // Do not send auth error when reading from headers, as the html body may contain the auth data
+    if status != "1" {
+      warn!("Found invalid saml-auth-status in headers: {}", status);
+      return;
+    }
+
+    let username = auth_response.get_header("saml-username");
+    let prelogin_cookie = auth_response.get_header("prelogin-cookie");
+    let portal_userauthcookie = auth_response.get_header("portal-userauthcookie");
+
+    match SamlAuthData::new(username, prelogin_cookie, portal_userauthcookie) {
+      Ok(auth_data) => self.send_auth_data(auth_data, AuthDataLocation::Headers),
+      Err(err) => {
+        warn!("Failed to read auth data from headers: {}", err);
+      }
+    }
+  }
+}
+
+fn extract_gpcallback(html: &str) -> Option<String> {
+  let re = Regex::new(r#"globalprotectcallback:[^"]+"#).unwrap();
+  re.captures(html)
+    .and_then(|captures| captures.get(0))
+    .map(|m| html_escape::decode_html_entities(m.as_str()).to_string())
+}
+
+#[cfg(test)]
+mod tests {
+  use super::*;
+
+  #[test]
+  fn extract_gpcallback_some() {
+    let html = r#"
+      <meta http-equiv="refresh" content="0; URL=globalprotectcallback:PGh0bWw+PCEtLSA8c">
+      <meta http-equiv="refresh" content="0; URL=globalprotectcallback:PGh0bWw+PCEtLSA8c">
+    "#;
+
+    assert_eq!(
+      extract_gpcallback(html).as_deref(),
+      Some("globalprotectcallback:PGh0bWw+PCEtLSA8c")
+    );
+  }
+
+  #[test]
+  fn extract_gpcallback_cas() {
+    let html = r#"
+      <meta http-equiv="refresh" content="0; URL=globalprotectcallback:cas-as=1&amp;un=xyz@email.com&amp;token=very_long_string">
+    "#;
+
+    assert_eq!(
+      extract_gpcallback(html).as_deref(),
+      Some("globalprotectcallback:cas-as=1&un=xyz@email.com&token=very_long_string")
+    );
+  }
+
+  #[test]
+  fn extract_gpcallback_none() {
+    let html = r#"
+      <meta http-equiv="refresh" content="0; URL=PGh0bWw+PCEtLSA8c">
+    "#;
+
+    assert_eq!(extract_gpcallback(html), None);
+  }
+}