From 0b7cfba124e59c064968fc45fe704a188277d704 Mon Sep 17 00:00:00 2001
From: Cliff Dyer <cliff.dyer@deepgram.com>
Date: Wed, 24 Jul 2024 16:34:07 -0400
Subject: [PATCH] Make Deepgram constructors fallible when taking a custom URL

---
 CHANGELOG.md            |  1 +
 src/lib.rs              | 57 +++++++++++++++++++++++++++++++----------
 src/listen/rest.rs      |  2 +-
 src/listen/websocket.rs |  3 ++-
 4 files changed, 48 insertions(+), 15 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 058e93d0..bb04da79 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -47,6 +47,7 @@ Some Enums have changed and may need to be updated
 - Add support for pre-recorded features when streaming
 - Add Speech to Text
 - Reorganize Code
+- Made Deepgram constructors fallible when taking a custom URL
 
 ### Streaming Features
 - endpointing
diff --git a/src/lib.rs b/src/lib.rs
index da064e42..dda02133 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -117,6 +117,10 @@ pub struct Deepgram {
 #[derive(Debug, Error)]
 #[non_exhaustive]
 pub enum DeepgramError {
+    /// The provided URL could not be used as a base_url.
+    #[error("The provided base_url is not valid. Please provide a URL starting with http:// or https://")]
+    InvalidBaseUrl,
+
     /// No source was provided to the request builder.
     #[error("No source was provided to the request builder.")]
     NoSource,
@@ -170,7 +174,9 @@ impl Deepgram {
     /// Panics under the same conditions as [`reqwest::Client::new`].
     pub fn new<K: AsRef<str>>(api_key: K) -> Self {
         let api_key = Some(api_key.as_ref().to_owned());
-        Self::inner_constructor(DEEPGRAM_BASE_URL.try_into().unwrap(), api_key)
+        // Unwrap: `DEEPGRAM_BASE_URL` is a valid URL.
+        let base_url = parse_base_url(DEEPGRAM_BASE_URL).unwrap();
+        Self::inner_constructor(base_url, api_key)
     }
 
     /// Construct a new Deepgram client with the specified base URL.
@@ -200,17 +206,19 @@ impl Deepgram {
     /// );
     /// ```
     ///
+    /// # Errors
+    ///
+    /// Returns a [`DeepgramError::InvalidBaseUrl`] if base_url is not a valid URL.
+    ///
     /// # Panics
     ///
-    /// Panics under the same conditions as [`reqwest::Client::new`], or if `base_url`
-    /// is not a valid URL.
-    pub fn with_base_url<U>(base_url: U) -> Self
+    /// Panics under the same conditions as [`reqwest::Client::new`].
+    pub fn with_base_url<U>(base_url: U) -> Result<Self>
     where
         U: TryInto<Url>,
-        U::Error: std::fmt::Debug,
     {
-        let base_url = base_url.try_into().expect("base_url must be a valid Url");
-        Self::inner_constructor(base_url, None)
+        let base_url = parse_base_url(base_url)?;
+        Ok(Self::inner_constructor(base_url, None))
     }
 
     /// Construct a new Deepgram client with the specified base URL and
@@ -236,18 +244,23 @@ impl Deepgram {
     /// );
     /// ```
     ///
+    /// # Errors
+    ///
+    /// Returns a [`DeepgramError::InvalidBaseUrl`] if base_url is not a valid URL.
+    ///
     /// # Panics
     ///
-    /// Panics under the same conditions as [`reqwest::Client::new`], or if `base_url`
-    /// is not a valid URL.
-    pub fn with_base_url_and_api_key<U, K>(base_url: U, api_key: K) -> Self
+    /// Panics under the same conditions as [`reqwest::Client::new`].
+    pub fn with_base_url_and_api_key<U, K>(base_url: U, api_key: K) -> Result<Self>
     where
         U: TryInto<Url>,
-        U::Error: std::fmt::Debug,
         K: AsRef<str>,
     {
-        let base_url = base_url.try_into().expect("base_url must be a valid Url");
-        Self::inner_constructor(base_url, Some(api_key.as_ref().to_owned()))
+        let base_url = parse_base_url(base_url)?;
+        Ok(Self::inner_constructor(
+            base_url,
+            Some(api_key.as_ref().to_owned()),
+        ))
     }
 
     fn inner_constructor(base_url: Url, api_key: Option<String>) -> Self {
@@ -282,6 +295,24 @@ impl Deepgram {
     }
 }
 
+fn parse_base_url<U>(base_url: U) -> Result<Url>
+where
+    U: TryInto<Url>,
+{
+    let Ok(base_url) = base_url.try_into() else {
+        return Err(DeepgramError::InvalidBaseUrl);
+    };
+    if base_url.cannot_be_a_base() {
+        return Err(DeepgramError::InvalidBaseUrl);
+    }
+
+    if !matches!(base_url.scheme(), "http" | "https" | "ws" | "wss") {
+        return Err(DeepgramError::InvalidBaseUrl);
+    }
+
+    Ok(base_url)
+}
+
 /// Sends the request and checks the response for an error.
 ///
 /// If there is an error, it translates it into a [`DeepgramError::DeepgramApiError`].
diff --git a/src/listen/rest.rs b/src/listen/rest.rs
index d9e174b8..9b924d92 100644
--- a/src/listen/rest.rs
+++ b/src/listen/rest.rs
@@ -284,7 +284,7 @@ mod tests {
 
     #[test]
     fn listen_url_custom_host() {
-        let dg = Deepgram::with_base_url("http://localhost:8888/abc/");
+        let dg = Deepgram::with_base_url("http://localhost:8888/abc/").unwrap();
         assert_eq!(
             &dg.transcription().listen_url().to_string(),
             "http://localhost:8888/abc/v1/listen"
diff --git a/src/listen/websocket.rs b/src/listen/websocket.rs
index b058a499..df4b3d6f 100644
--- a/src/listen/websocket.rs
+++ b/src/listen/websocket.rs
@@ -412,7 +412,8 @@ mod tests {
 
     #[test]
     fn test_stream_url_custom_host() {
-        let dg = crate::Deepgram::with_base_url_and_api_key("http://localhost:8080", "token");
+        let dg =
+            crate::Deepgram::with_base_url_and_api_key("http://localhost:8080", "token").unwrap();
         assert_eq!(
             dg.transcription().listen_stream_url().to_string(),
             "ws://localhost:8080/v1/listen",