From 030037eeed8c98544baa6821e0681f318ef910ee Mon Sep 17 00:00:00 2001
From: schetterer <64965278+schetterer@users.noreply.github.com>
Date: Mon, 7 Dec 2020 14:28:16 +0800
Subject: [PATCH] Allow cluster to specify a custom http url (#55)

* Allow cluster to specify a custom http url
---
 common/src/client/mod.rs | 43 ++++++++++++++++++++++++++++++++++++----
 1 file changed, 39 insertions(+), 4 deletions(-)

diff --git a/common/src/client/mod.rs b/common/src/client/mod.rs
index 597b9a6..200a901 100644
--- a/common/src/client/mod.rs
+++ b/common/src/client/mod.rs
@@ -3,7 +3,7 @@ use std::str::FromStr;
 
 pub mod rpc;
 
-#[derive(Clone, Debug)]
+#[derive(Clone, Debug, Eq, PartialEq)]
 pub enum Cluster {
     Testnet,
     Mainnet,
@@ -11,12 +11,14 @@ pub enum Cluster {
     Devnet,
     Localnet,
     Debug,
+    Custom(String),
 }
 
 impl FromStr for Cluster {
     type Err = anyhow::Error;
     fn from_str(s: &str) -> Result<Cluster> {
         match s.to_lowercase().as_str() {
+            url if url.contains("http") => Ok(Cluster::Custom(s.to_owned())),
             "t" | "testnet" => Ok(Cluster::Testnet),
             "m" | "mainnet" => Ok(Cluster::Mainnet),
             "v" | "vipmainnet" => Ok(Cluster::VipMainnet),
@@ -24,7 +26,7 @@ impl FromStr for Cluster {
             "l" | "localnet" => Ok(Cluster::Localnet),
             "g" | "debug" => Ok(Cluster::Debug),
             _ => Err(anyhow::Error::msg(
-                "Cluster must be one of [testnet, mainnet, devnet]\n",
+                "Cluster must be one of [testnet, mainnet, devnet] or be an http or https url\n",
             )),
         }
     }
@@ -38,14 +40,15 @@ impl std::fmt::Display for Cluster {
             Cluster::VipMainnet => "vipmainnet",
             Cluster::Devnet => "devnet",
             Cluster::Localnet => "localnet",
-            Cluster::Debug => "debut",
+            Cluster::Debug => "debug",
+            Cluster::Custom(url) => url,
         };
         write!(f, "{}", clust_str)
     }
 }
 
 impl Cluster {
-    pub fn url(&self) -> &'static str {
+    pub fn url(&self) -> &str {
         match self {
             Cluster::Devnet => "https://devnet.solana.com",
             Cluster::Testnet => "https://testnet.solana.com",
@@ -53,6 +56,38 @@ impl Cluster {
             Cluster::VipMainnet => "https://vip-api.mainnet-beta.solana.com",
             Cluster::Localnet => "http://127.0.0.1:8899",
             Cluster::Debug => "http://34.90.18.145:8899",
+            Cluster::Custom(url) => url,
         }
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    fn test_cluster(name: &str, cluster: Cluster) {
+        assert_eq!(Cluster::from_str(name).unwrap(), cluster);
+    }
+
+    #[test]
+    fn test_cluster_parse() {
+        test_cluster("testnet", Cluster::Testnet);
+        test_cluster("mainnet", Cluster::Mainnet);
+        test_cluster("vipmainnet", Cluster::VipMainnet);
+        test_cluster("devnet", Cluster::Devnet);
+        test_cluster("localnet", Cluster::Localnet);
+        test_cluster("debug", Cluster::Debug);
+
+        let custom_http = "http://my_custom_url.test.net";
+        test_cluster(custom_http, Cluster::Custom(custom_http.into()));
+        let custom_https = "https://my_custom_url.test.net";
+        test_cluster(custom_https, Cluster::Custom(custom_https.into()));
+    }
+
+    #[test]
+    #[should_panic]
+    fn test_cluster_bad_parse() {
+        let bad_url = "httq://my_custom_url.test.net";
+        Cluster::from_str(bad_url).unwrap();
+    }
+}