From 4e67182bcbcea837aff0f358be35ef85e7bcd704 Mon Sep 17 00:00:00 2001
From: Josh Robson Chase <josh@robsonchase.com>
Date: Thu, 17 Aug 2023 11:15:17 -0400
Subject: [PATCH] ngrok: parse error responses and surface the msg and code

---
 ngrok/src/internals/proto.rs       | 61 ++++++++++++++++++++++++++++++
 ngrok/src/internals/raw_session.rs | 24 ++++++++++--
 ngrok/src/lib.rs                   |  1 +
 ngrok/src/session.rs               | 22 +++++++++--
 4 files changed, 102 insertions(+), 6 deletions(-)

diff --git a/ngrok/src/internals/proto.rs b/ngrok/src/internals/proto.rs
index 516fadc..2412241 100644
--- a/ngrok/src/internals/proto.rs
+++ b/ngrok/src/internals/proto.rs
@@ -1,5 +1,6 @@
 use std::{
     collections::HashMap,
+    error,
     fmt,
     io,
     ops::{
@@ -38,6 +39,66 @@ pub const SRV_INFO_REQ: StreamType = StreamType::clamp(8);
 
 pub const VERSION: &str = "2";
 
+/// An error that may have an ngrok error code.
+/// All ngrok error codes are documented at https://ngrok.com/docs/errors
+pub trait NgrokError: error::Error {
+    /// Return the ngrok error code, if one exists for this error.
+    fn error_code(&self) -> Option<&str> {
+        None
+    }
+    /// Return the error message minus the ngrok error code.
+    /// If this error has no error code, this is equivalent to
+    /// `format!("{error}")`.
+    fn msg(&self) -> String {
+        format!("{self}")
+    }
+}
+
+#[derive(Serialize, Deserialize, Debug, Clone, Default)]
+pub struct ErrResp {
+    pub msg: String,
+    pub error_code: Option<String>,
+}
+
+impl<'a> From<&'a str> for ErrResp {
+    fn from(value: &'a str) -> Self {
+        let mut error_code = None;
+        let mut msg_lines = vec![];
+        for line in value.lines().filter(|l| !l.is_empty()) {
+            if line.starts_with("ERR_NGROK_") {
+                error_code = line.split('_').nth(2).map(String::from);
+            } else {
+                msg_lines.push(line);
+            }
+        }
+        ErrResp {
+            error_code,
+            msg: msg_lines.join("\n"),
+        }
+    }
+}
+
+impl error::Error for ErrResp {}
+
+impl fmt::Display for ErrResp {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        self.msg.fmt(f)?;
+        if let Some(code) = &self.error_code {
+            write!(f, "\n\nERR_NGROK_{code}")?;
+        }
+        Ok(())
+    }
+}
+
+impl NgrokError for ErrResp {
+    fn error_code(&self) -> Option<&str> {
+        self.error_code.as_deref()
+    }
+    fn msg(&self) -> String {
+        self.msg.clone()
+    }
+}
+
 #[derive(Serialize, Deserialize, Debug, Clone, Default)]
 #[serde(rename_all = "PascalCase")]
 pub struct Auth {
diff --git a/ngrok/src/internals/raw_session.rs b/ngrok/src/internals/raw_session.rs
index fa388f1..c1edc98 100644
--- a/ngrok/src/internals/raw_session.rs
+++ b/ngrok/src/internals/raw_session.rs
@@ -57,6 +57,8 @@ use super::{
         BindOpts,
         BindResp,
         CommandResp,
+        ErrResp,
+        NgrokError,
         ProxyHeader,
         ReadHeaderError,
         Restart,
@@ -95,8 +97,24 @@ pub enum RpcError {
     #[error("failed to deserialize rpc response")]
     InvalidResponse(#[from] serde_json::Error),
     /// There was an error in the RPC response.
-    #[error("rpc error response: {0}")]
-    Response(String),
+    #[error("rpc error response:\n{0}")]
+    Response(ErrResp),
+}
+
+impl NgrokError for RpcError {
+    fn error_code(&self) -> Option<&str> {
+        match self {
+            RpcError::Response(resp) => resp.error_code(),
+            _ => None,
+        }
+    }
+
+    fn msg(&self) -> String {
+        match self {
+            RpcError::Response(resp) => resp.msg(),
+            _ => format!("{self}"),
+        }
+    }
 }
 
 #[derive(Error, Debug)]
@@ -258,7 +276,7 @@ impl RpcClient {
         if let Ok(err) = err_resp {
             if !err.error.is_empty() {
                 debug!(?err, "decoded rpc error response");
-                return Err(RpcError::Response(err.error));
+                return Err(RpcError::Response(err.error.as_str().into()));
             }
         }
 
diff --git a/ngrok/src/lib.rs b/ngrok/src/lib.rs
index 34e33a3..5277e88 100644
--- a/ngrok/src/lib.rs
+++ b/ngrok/src/lib.rs
@@ -55,6 +55,7 @@ pub mod prelude {
     #[doc(inline)]
     pub use crate::{
         config::TunnelBuilder,
+        internals::proto::NgrokError,
         tunnel::{
             LabelsTunnel,
             ProtoTunnel,
diff --git a/ngrok/src/session.rs b/ngrok/src/session.rs
index a9eeb7a..efee417 100644
--- a/ngrok/src/session.rs
+++ b/ngrok/src/session.rs
@@ -88,6 +88,7 @@ use crate::{
             AuthExtra,
             BindExtra,
             BindOpts,
+            NgrokError,
             SecretString,
         },
         raw_session::{
@@ -277,13 +278,13 @@ pub enum ConnectError {
     /// This might occur when there's a protocol mismatch interfering with the
     /// heartbeat routine.
     #[error("failed to start ngrok session")]
-    Start(StartSessionError),
+    Start(#[source] StartSessionError),
     /// An error occurred when attempting to authenticate.
     #[error("authentication failure")]
-    Auth(RpcError),
+    Auth(#[source] RpcError),
     /// An error occurred when rebinding tunnels during a reconnect
     #[error("error rebinding tunnel after reconnect")]
-    Rebind(RpcError),
+    Rebind(#[source] RpcError),
     /// The (re)connect function gave up.
     ///
     /// This will never be returned by the default connect function, and is
@@ -292,6 +293,21 @@ pub enum ConnectError {
     Canceled,
 }
 
+impl NgrokError for ConnectError {
+    fn error_code(&self) -> Option<&str> {
+        match self {
+            ConnectError::Auth(resp) | ConnectError::Rebind(resp) => resp.error_code(),
+            _ => None,
+        }
+    }
+    fn msg(&self) -> String {
+        match self {
+            ConnectError::Auth(resp) | ConnectError::Rebind(resp) => resp.msg(),
+            _ => format!("{self}"),
+        }
+    }
+}
+
 impl Default for SessionBuilder {
     fn default() -> Self {
         SessionBuilder {