From 2e03c63d1b2ed7873eb4846525684f05ea9108a7 Mon Sep 17 00:00:00 2001
From: moss <moss@banay.me>
Date: Sat, 14 Jan 2023 17:07:50 +1100
Subject: [PATCH] Add support for using TLS with PostgreSQL (#260)

---
 refinery_cli/Cargo.toml             |  2 +-
 refinery_core/Cargo.toml            |  2 ++
 refinery_core/src/config.rs         | 27 +++++++++++++++++++++++++++
 refinery_core/src/drivers/config.rs | 11 ++++++++++-
 4 files changed, 40 insertions(+), 2 deletions(-)

diff --git a/refinery_cli/Cargo.toml b/refinery_cli/Cargo.toml
index ae58a2ee..8712cd79 100644
--- a/refinery_cli/Cargo.toml
+++ b/refinery_cli/Cargo.toml
@@ -16,7 +16,7 @@ path = "src/main.rs"
 
 [features]
 default = ["mysql", "postgresql", "sqlite-bundled", "mssql"]
-postgresql = ["refinery-core/postgres"]
+postgresql = ["refinery-core/postgres", "refinery-core/postgres-openssl", "refinery-core/openssl"]
 mysql = ["refinery-core/mysql", "refinery-core/flate2"]
 sqlite = ["refinery-core/rusqlite"]
 sqlite-bundled = ["sqlite", "refinery-core/rusqlite-bundled"]
diff --git a/refinery_core/Cargo.toml b/refinery_core/Cargo.toml
index 11b46917..b3fe9cc1 100644
--- a/refinery_core/Cargo.toml
+++ b/refinery_core/Cargo.toml
@@ -32,6 +32,8 @@ walkdir = "2.3.1"
 # allow multiple versions of the same dependency if API is similar
 rusqlite = { version = ">= 0.23, <= 0.28", optional = true }
 postgres = { version = "0.19", optional = true }
+postgres-openssl = { version = "0.5", optional = true }
+openssl = { version = "0.10", optional = true }
 tokio-postgres = { version = "0.7", optional = true }
 mysql = { version = ">= 21.0.0, <= 23", optional = true, default-features = false}
 mysql_async = { version = ">= 0.28, <= 0.30", optional = true }
diff --git a/refinery_core/src/config.rs b/refinery_core/src/config.rs
index c1fac221..0195dd48 100644
--- a/refinery_core/src/config.rs
+++ b/refinery_core/src/config.rs
@@ -34,6 +34,7 @@ impl Config {
                 db_user: None,
                 db_pass: None,
                 db_name: None,
+                use_tls: false,
                 #[cfg(feature = "tiberius-config")]
                 trust_cert: false,
             },
@@ -138,6 +139,10 @@ impl Config {
         self.main.db_port.as_deref()
     }
 
+    pub fn use_tls(&self) -> bool {
+        self.main.use_tls
+    }
+
     pub fn set_db_user(self, db_user: &str) -> Config {
         Config {
             main: Main {
@@ -182,6 +187,15 @@ impl Config {
             },
         }
     }
+
+    pub fn set_use_tls(self, use_tls: bool) -> Config {
+        Config {
+            main: Main {
+                use_tls: use_tls,
+                ..self.main
+            },
+        }
+    }
 }
 
 impl TryFrom<Url> for Config {
@@ -219,6 +233,17 @@ impl TryFrom<Url> for Config {
                             None,
                         )
                     })?;
+
+                let use_tls = query_params.
+                    get("sslmode")
+                    .unwrap_or(&Cow::Borrowed("false"))
+                    .parse::<bool>()
+                    .map_err(|_| {
+                        Error::new(
+                            Kind::ConfigError("Invalid sslmode value, please use true/false".into()),
+                            None,
+                        )
+                    })?;
             }
         }
 
@@ -237,6 +262,7 @@ impl TryFrom<Url> for Config {
                 db_user: Some(url.username().to_string()),
                 db_pass: url.password().map(|r| r.to_string()),
                 db_name: Some(url.path().trim_start_matches('/').to_string()),
+                use_tls,
                 #[cfg(feature = "tiberius-config")]
                 trust_cert,
             },
@@ -268,6 +294,7 @@ struct Main {
     db_user: Option<String>,
     db_pass: Option<String>,
     db_name: Option<String>,
+    use_tls: bool,
     #[cfg(feature = "tiberius-config")]
     #[serde(default)]
     trust_cert: bool,
diff --git a/refinery_core/src/drivers/config.rs b/refinery_core/src/drivers/config.rs
index 92a00582..6d619108 100644
--- a/refinery_core/src/drivers/config.rs
+++ b/refinery_core/src/drivers/config.rs
@@ -80,7 +80,16 @@ macro_rules! with_connection {
                 cfg_if::cfg_if! {
                     if #[cfg(feature = "postgres")] {
                         let path = build_db_url("postgresql", &$config);
-                        let conn = postgres::Client::connect(path.as_str(), postgres::NoTls).migration_err("could not connect to database", None)?;
+
+                        let conn;
+                        if $config.use_tls() {
+                            let builder = openssl::ssl::SslConnector::builder(openssl::ssl::SslMethod::tls()).unwrap();
+                            let connector = postgres_openssl::MakeTlsConnector::new(builder.build());
+                            conn = postgres::Client::connect(path.as_str(), connector).migration_err("could not connect to database", None)?;
+                        } else {
+                            conn = postgres::Client::connect(path.as_str(), postgres::NoTls).migration_err("could not connect to database", None)?;
+                        }
+
                         $op(conn)
                     } else {
                         panic!("tried to migrate from config for a postgresql database, but feature postgres not enabled!");