Connect to public database URLs (#532)

* connect to public db url for postgres, redis, and mongo

* mysql connect command

* grammar

src/commands/connect.rs

@@ -1,6 +1,7 @@
 use anyhow::bail;
 use std::{collections::BTreeMap, fmt::Display};
 use tokio::process::Command;
+use url::Url;
 use which::which;
 
 use crate::controllers::{
@@ -120,44 +121,308 @@ fn get_connect_command(
     database_type: DatabaseType,
     variables: BTreeMap<String, String>,
 ) -> Result<(String, Vec<String>)> {
-    let pass_arg; // Hack to get ownership of formatted string outside match
-    let default = &"".to_string();
+    match &database_type {
+        DatabaseType::PostgreSQL => get_postgres_command(&variables),
+        DatabaseType::Redis => return get_redis_command(&variables),
+        DatabaseType::MongoDB => return get_mongo_command(&variables),
+        DatabaseType::MySQL => return get_mysql_command(&variables),
+    }
+}
 
-    let (cmd_name, args): (&str, Vec<&str>) = match &database_type {
-        DatabaseType::PostgreSQL => (
-            "psql",
-            vec![variables.get("DATABASE_URL").unwrap_or(default)],
-        ),
-        DatabaseType::Redis => (
-            "redis-cli",
-            vec!["-u", variables.get("REDIS_URL").unwrap_or(default)],
-        ),
-        DatabaseType::MongoDB => (
-            "mongosh",
-            vec![variables.get("MONGO_URL").unwrap_or(default).as_str()],
-        ),
-        DatabaseType::MySQL => {
-            // -p is a special case as it requires no whitespace between arg and value
-            pass_arg = format!("-p{}", variables.get("MYSQLPASSWORD").unwrap_or(default));
-            (
-                "mysql",
-                vec![
-                    "-h",
-                    variables.get("MYSQLHOST").unwrap_or(default),
-                    "-u",
-                    variables.get("MYSQLUSER").unwrap_or(default),
-                    "-P",
-                    variables.get("MYSQLPORT").unwrap_or(default),
-                    "-D",
-                    variables.get("MYSQLDATABASE").unwrap_or(default),
-                    pass_arg.as_str(),
-                ],
-            )
-        }
-    };
+fn host_is_tcp_proxy(connect_url: String) -> bool {
+    connect_url.contains("proxy.rlwy.net")
+}
+
+fn get_postgres_command(variables: &BTreeMap<String, String>) -> Result<(String, Vec<String>)> {
+    let connect_url = variables
+        .get("DATABASE_PUBLIC_URL")
+        .or_else(|| variables.get("DATABASE_URL"))
+        .map(|s| s.to_string())
+        .ok_or(RailwayError::ConnectionVariableNotFound(
+            "DATABASE_PUBLIC_URL".to_string(),
+        ))?;
+
+    if !host_is_tcp_proxy(connect_url.clone()) {
+        return Err(RailwayError::InvalidConnectionVariable.into());
+    }
+
+    Ok(("psql".to_string(), vec![connect_url]))
+}
+
+fn get_redis_command(variables: &BTreeMap<String, String>) -> Result<(String, Vec<String>)> {
+    let connect_url = variables
+        .get("REDIS_PUBLIC_URL")
+        .or_else(|| variables.get("REDIS_URL"))
+        .map(|s| s.to_string())
+        .ok_or(RailwayError::ConnectionVariableNotFound(
+            "REDIS_PUBLIC_URL".to_string(),
+        ))?;
+
+    if !host_is_tcp_proxy(connect_url.clone()) {
+        return Err(RailwayError::InvalidConnectionVariable.into());
+    }
+
+    Ok(("redis-cli".to_string(), vec!["-u".to_string(), connect_url]))
+}
+
+fn get_mongo_command(variables: &BTreeMap<String, String>) -> Result<(String, Vec<String>)> {
+    let connect_url = variables
+        .get("MONGO_PUBLIC_URL")
+        .or_else(|| variables.get("MONGO_URL"))
+        .map(|s| s.to_string())
+        .ok_or(RailwayError::ConnectionVariableNotFound(
+            "MONGO_PUBLIC_URL".to_string(),
+        ))?;
+
+    if !host_is_tcp_proxy(connect_url.clone()) {
+        return Err(RailwayError::InvalidConnectionVariable.into());
+    }
+
+    Ok(("mongosh".to_string(), vec![connect_url]))
+}
+
+fn get_mysql_command(variables: &BTreeMap<String, String>) -> Result<(String, Vec<String>)> {
+    let connect_url = variables
+        .get("MYSQL_PUBLIC_URL")
+        .or_else(|| variables.get("MYSQL_URL"))
+        .map(|s| s.to_string())
+        .ok_or(RailwayError::ConnectionVariableNotFound(
+            "MYSQL_PUBLIC_URL".to_string(),
+        ))?;
+
+    if !host_is_tcp_proxy(connect_url.clone()) {
+        return Err(RailwayError::InvalidConnectionVariable.into());
+    }
+
+    let parsed_url =
+        Url::parse(&connect_url).map_err(|_err| RailwayError::InvalidConnectionVariable)?;
+
+    let host = parsed_url.host_str().unwrap_or("");
+    let user = parsed_url.username();
+    let password = parsed_url.password().unwrap_or("");
+    let port = parsed_url.port().unwrap_or(3306);
+    let database = parsed_url.path().trim_start_matches('/');
+
+    let pass_arg = format!("-p{}", password);
 
     Ok((
-        cmd_name.to_string(),
-        args.iter().map(|s| s.to_string()).collect(),
+        "mysql".to_string(),
+        vec![
+            "-h".to_string(),
+            host.to_string(),
+            "-u".to_string(),
+            user.to_string(),
+            "-P".to_string(),
+            port.to_string(),
+            "-D".to_string(),
+            database.to_string(),
+            pass_arg,
+        ],
     ))
 }
+
+#[cfg(test)]
+mod test {
+    use super::*;
+
+    #[test]
+    fn test_is_tcp_proxy() {
+        assert!(host_is_tcp_proxy("roundhouse.proxy.rlwy.net".to_string()));
+        assert!(!host_is_tcp_proxy("localhost".to_string()));
+        assert!(!host_is_tcp_proxy("postgres.railway.interal".to_string()));
+    }
+
+    #[test]
+    fn test_gets_postgres_command() {
+        let private_postgres_url =
+            "postgresql://postgres:password@name.railway.internal:5432/railway".to_string();
+        let public_postgres_url =
+            "postgresql://postgres:password@roundhouse.proxy.rlwy.net:55555/railway".to_string();
+
+        // Valid DATABASE_PUBLIC_URL
+        {
+            let mut variables = BTreeMap::new();
+            variables.insert(
+                "DATABASE_PUBLIC_URL".to_string(),
+                public_postgres_url.clone(),
+            );
+            variables.insert("DATABASE_URL".to_string(), private_postgres_url.clone());
+
+            let (cmd, args) = get_postgres_command(&variables).unwrap();
+            assert_eq!(cmd, "psql");
+            assert_eq!(args, vec![public_postgres_url.clone()]);
+        }
+
+        // Valid DATABASE_URL
+        {
+            let mut variables = BTreeMap::new();
+            variables.insert("DATABASE_URL".to_string(), public_postgres_url.clone());
+
+            let (cmd, args) = get_postgres_command(&variables).unwrap();
+            assert_eq!(cmd, "psql");
+            assert_eq!(args, vec![public_postgres_url.clone()]);
+        }
+
+        {
+            let variables = BTreeMap::new();
+            let res = get_postgres_command(&variables);
+            assert!(res.is_err());
+            assert_eq!(
+                res.unwrap_err().to_string(),
+                RailwayError::ConnectionVariableNotFound("DATABASE_PUBLIC_URL".to_string())
+                    .to_string()
+            );
+        }
+
+        // Invalid DATABASE_URL
+        {
+            let mut variables = BTreeMap::new();
+            variables.insert("DATABASE_URL".to_string(), private_postgres_url.clone());
+
+            let res = get_postgres_command(&variables);
+            assert!(res.is_err());
+            assert_eq!(
+                res.unwrap_err().to_string(),
+                RailwayError::InvalidConnectionVariable.to_string()
+            );
+        }
+    }
+
+    #[test]
+    fn test_gets_redis_command() {
+        let private_redis_url = "redis://default:password@redis.railway.internal:6379".to_string();
+        let public_redis_url = "redis://default:password@monorail.proxy.rlwy.net:26137".to_string();
+
+        // Valid REDIS_PUBLIC_URL
+        {
+            let mut variables = BTreeMap::new();
+            variables.insert("REDIS_PUBLIC_URL".to_string(), public_redis_url.clone());
+            variables.insert("REDIS_URL".to_string(), private_redis_url.clone());
+
+            let (cmd, args) = get_redis_command(&variables).unwrap();
+            assert_eq!(cmd, "redis-cli");
+            assert_eq!(args, vec!["-u".to_string(), public_redis_url.clone()]);
+        }
+
+        // Valid REDIS_URL
+        {
+            let mut variables = BTreeMap::new();
+            variables.insert("REDIS_URL".to_string(), public_redis_url.clone());
+
+            let (cmd, args) = get_redis_command(&variables).unwrap();
+            assert_eq!(cmd, "redis-cli");
+            assert_eq!(args, vec!["-u".to_string(), public_redis_url.clone()]);
+        }
+
+        // No public Redis URL
+        {
+            let variables = BTreeMap::new();
+            let res = get_redis_command(&variables);
+            assert!(res.is_err());
+            assert_eq!(
+                res.unwrap_err().to_string(),
+                RailwayError::ConnectionVariableNotFound("REDIS_PUBLIC_URL".to_string())
+                    .to_string()
+            );
+        }
+    }
+
+    #[test]
+    fn test_gets_mongo_command() {
+        let private_mongo_url =
+            "mongodb://user:password@mongo.railway.internal:27017/railway".to_string();
+        let public_mongo_url =
+            "mongodb://user:password@roundhouse.proxy.rlwy.net:33333/railway".to_string();
+
+        // Valid MONGO_PUBLIC_URL
+        {
+            let mut variables = BTreeMap::new();
+            variables.insert("MONGO_PUBLIC_URL".to_string(), public_mongo_url.clone());
+            variables.insert("MONGO_URL".to_string(), private_mongo_url.clone());
+
+            let (cmd, args) = get_mongo_command(&variables).unwrap();
+            assert_eq!(cmd, "mongosh");
+            assert_eq!(args, vec![public_mongo_url.clone()]);
+        }
+
+        // Valid MONGO_URL
+        {
+            let mut variables = BTreeMap::new();
+            variables.insert("MONGO_URL".to_string(), public_mongo_url.clone());
+
+            let (cmd, args) = get_mongo_command(&variables).unwrap();
+            assert_eq!(cmd, "mongosh");
+            assert_eq!(args, vec![public_mongo_url.clone()]);
+        }
+
+        // No public Mongo URL
+        {
+            let variables = BTreeMap::new();
+            let res = get_mongo_command(&variables);
+            assert!(res.is_err());
+            assert_eq!(
+                res.unwrap_err().to_string(),
+                RailwayError::ConnectionVariableNotFound("MONGO_PUBLIC_URL".to_string())
+                    .to_string()
+            );
+        }
+
+        // Invalid MONGO_URL
+        {
+            let mut variables = BTreeMap::new();
+            variables.insert("MONGO_URL".to_string(), private_mongo_url.clone());
+
+            let res = get_mongo_command(&variables);
+            assert!(res.is_err());
+            assert_eq!(
+                res.unwrap_err().to_string(),
+                RailwayError::InvalidConnectionVariable.to_string()
+            );
+        }
+    }
+
+    #[test]
+    fn test_gets_mysql_command() {
+        let private_mysql_url =
+            "mysql://user:password@mysql.railway.internal:3306/railway".to_string();
+        let public_mysql_url =
+            "mysql://user:password@roundhouse.proxy.rlwy.net:12345/railway".to_string();
+
+        // Valid MYSQL_PUBLIC_URL
+        {
+            let mut variables = BTreeMap::new();
+            variables.insert("MYSQL_PUBLIC_URL".to_string(), public_mysql_url.clone());
+            variables.insert("MYSQL_URL".to_string(), private_mysql_url.clone());
+
+            let (cmd, args) = get_mysql_command(&variables).unwrap();
+            assert_eq!(cmd, "mysql");
+            assert_eq!(
+                args,
+                vec![
+                    "-h".to_string(),
+                    "roundhouse.proxy.rlwy.net".to_string(),
+                    "-u".to_string(),
+                    "user".to_string(),
+                    "-P".to_string(),
+                    "12345".to_string(),
+                    "-D".to_string(),
+                    "railway".to_string(),
+                    "-ppassword".to_string(),
+                ]
+            );
+        }
+
+        // Invalid URL format
+        {
+            let mut variables = BTreeMap::new();
+            variables.insert("MYSQL_URL".to_string(), "invalid_url".to_string());
+
+            let res = get_mysql_command(&variables);
+            assert!(res.is_err());
+            assert_eq!(
+                res.unwrap_err().to_string(),
+                RailwayError::InvalidConnectionVariable.to_string()
+            );
+        }
+    }
+}

src/errors.rs

@@ -55,4 +55,10 @@ pub enum RailwayError {
 
     #[error("2FA code is incorrect. Please try again.")]
     InvalidTwoFactorCode,
+
+    #[error("Could not find a variable to connect to the service with. Looking for \"{0}\".")]
+    ConnectionVariableNotFound(String),
+
+    #[error("Connection URL should point to the Railway TCP proxy")]
+    InvalidConnectionVariable,
 }