From 1548dc15e415a2298c25af9b04f33bba3d28b998 Mon Sep 17 00:00:00 2001
From: Laura K <6276139+alula@users.noreply.github.com>
Date: Mon, 24 Jul 2023 14:17:02 +0200
Subject: [PATCH] Fix wrong handling of IPv6 addresses in URLs for PostgreSQL
 and MySQL (#4051)

---
 quaint/src/connector/mysql.rs    | 21 +++++++++++++++++++--
 quaint/src/connector/postgres.rs | 26 ++++++++++++++++++++------
 2 files changed, 39 insertions(+), 8 deletions(-)

diff --git a/quaint/src/connector/mysql.rs b/quaint/src/connector/mysql.rs
index 68f4ac28fc35..ce116dd2df6f 100644
--- a/quaint/src/connector/mysql.rs
+++ b/quaint/src/connector/mysql.rs
@@ -22,7 +22,7 @@ use std::{
     time::Duration,
 };
 use tokio::sync::Mutex;
-use url::Url;
+use url::{Host, Url};
 
 /// The underlying MySQL driver. Only available with the `expose-drivers`
 /// Cargo feature.
@@ -98,7 +98,18 @@ impl MysqlUrl {
 
     /// The database host. If `socket` and `host` are not set, defaults to `localhost`.
     pub fn host(&self) -> &str {
-        self.url.host_str().unwrap_or("localhost")
+        match (self.url.host(), self.url.host_str()) {
+            (Some(Host::Ipv6(_)), Some(host)) => {
+                // The `url` crate may return an IPv6 address in brackets, which must be stripped.
+                if host.starts_with('[') && host.ends_with(']') {
+                    &host[1..host.len() - 1]
+                } else {
+                    host
+                }
+            }
+            (_, Some(host)) => host,
+            _ => "localhost",
+        }
     }
 
     /// If set, connected to the database through a Unix socket.
@@ -604,6 +615,12 @@ mod tests {
         assert!(!url.query_params.ssl_opts.accept_invalid_certs());
     }
 
+    #[test]
+    fn should_parse_ipv6_host() {
+        let url = MysqlUrl::new(Url::parse("mysql://[2001:db8:1234::ffff]:5432/testdb").unwrap()).unwrap();
+        assert_eq!("2001:db8:1234::ffff", url.host());
+    }
+
     #[test]
     fn should_allow_changing_of_cache_size() {
         let url = MysqlUrl::new(Url::parse("mysql:///root:root@localhost:3307/foo?statement_cache_size=420").unwrap())
diff --git a/quaint/src/connector/postgres.rs b/quaint/src/connector/postgres.rs
index 50ecdf09eb16..efa414dc2461 100644
--- a/quaint/src/connector/postgres.rs
+++ b/quaint/src/connector/postgres.rs
@@ -25,7 +25,7 @@ use tokio_postgres::{
     config::{ChannelBinding, SslMode},
     Client, Config, Statement,
 };
-use url::Url;
+use url::{Host, Url};
 
 pub(crate) const DEFAULT_SCHEMA: &str = "public";
 
@@ -223,11 +223,19 @@ impl PostgresUrl {
     ///
     /// If none of them are set, defaults to `localhost`.
     pub fn host(&self) -> &str {
-        match (self.query_params.host.as_ref(), self.url.host_str()) {
-            (Some(host), _) => host.as_str(),
-            (None, Some("")) => "localhost",
-            (None, None) => "localhost",
-            (None, Some(host)) => host,
+        match (self.query_params.host.as_ref(), self.url.host_str(), self.url.host()) {
+            (Some(host), _, _) => host.as_str(),
+            (None, Some(""), _) => "localhost",
+            (None, None, _) => "localhost",
+            (None, Some(host), Some(Host::Ipv6(_))) => {
+                // The `url` crate may return an IPv6 address in brackets, which must be stripped.
+                if host.starts_with('[') && host.ends_with(']') {
+                    &host[1..host.len() - 1]
+                } else {
+                    host
+                }
+            }
+            (None, Some(host), _) => host,
         }
     }
 
@@ -1142,6 +1150,12 @@ mod tests {
         assert_eq!("localhost", url.host());
     }
 
+    #[test]
+    fn should_parse_ipv6_host() {
+        let url = PostgresUrl::new(Url::parse("postgresql://[2001:db8:1234::ffff]:5432/dbname").unwrap()).unwrap();
+        assert_eq!("2001:db8:1234::ffff", url.host());
+    }
+
     #[test]
     fn should_handle_options_field() {
         let url = PostgresUrl::new(Url::parse("postgresql:///localhost:5432?options=--cluster%3Dmy_cluster").unwrap())