Fix wrong handling of IPv6 addresses in URLs for PostgreSQL and MySQL (…

…#4051)
prisma · Jul 24, 2023 · 1548dc1 · 1548dc1
1 parent f5c61f9
commit 1548dc1
Show file tree

Hide file tree

Showing 2 changed files with 39 additions and 8 deletions.
diff --git 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
@@ -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())