From 4e02b582ce9aa88d11886bc05b9481d2ad6bfda3 Mon Sep 17 00:00:00 2001 From: Jake Runzer Date: Wed, 14 Aug 2024 13:45:27 -0400 Subject: [PATCH] Connect to public database URLs (#532) * connect to public db url for postgres, redis, and mongo * mysql connect command * grammar --- src/commands/connect.rs | 337 +++++++++++++++++++++++++++++++++++----- src/errors.rs | 6 + 2 files changed, 307 insertions(+), 36 deletions(-) diff --git a/src/commands/connect.rs b/src/commands/connect.rs index 629a1c8..2760b5a 100644 --- a/src/commands/connect.rs +++ b/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, ) -> Result<(String, Vec)> { - 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) -> Result<(String, Vec)> { + 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) -> Result<(String, Vec)> { + 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) -> Result<(String, Vec)> { + 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) -> Result<(String, Vec)> { + 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() + ); + } + } +} diff --git a/src/errors.rs b/src/errors.rs index cccc0b6..193fb09 100644 --- a/src/errors.rs +++ b/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, }