IpAddr + try hostaddr first

Author: trungda

tokio-postgres/src/config.rs

@@ -13,6 +13,7 @@ use crate::{Client, Connection, Error};
 use std::borrow::Cow;
 #[cfg(unix)]
 use std::ffi::OsStr;
+use std::net::IpAddr;
 use std::ops::Deref;
 #[cfg(unix)]
 use std::os::unix::ffi::OsStrExt;
@@ -98,7 +99,9 @@ pub enum Host {
 ///     - or if host specifies an IP address, that value will be used directly.
 ///     Using `hostaddr` allows the application to avoid a host name look-up, which might be important in applications
 ///     with time constraints. However, a host name is required for verify-full SSL certificate verification.
-///     Note that `host` is always required regardless of whether `hostaddr` is present.
+///     Specifically:
+///         * If `hostaddr` is specified without `host`, the value for `hostaddr` gives the server network address.
+///             The connection attempt will fail if the authentication method requires a host name;
 ///         * If `host` is specified without `hostaddr`, a host name lookup occurs;
 ///         * If both `host` and `hostaddr` are specified, the value for `hostaddr` gives the server network address.
 ///             The value for `host` is ignored unless the authentication method requires it,
@@ -174,7 +177,7 @@ pub struct Config {
     pub(crate) application_name: Option<String>,
     pub(crate) ssl_mode: SslMode,
     pub(crate) host: Vec<Host>,
-    pub(crate) hostaddr: Vec<String>,
+    pub(crate) hostaddr: Vec<IpAddr>,
     pub(crate) port: Vec<u16>,
     pub(crate) connect_timeout: Option<Duration>,
     pub(crate) keepalives: bool,
@@ -317,7 +320,7 @@ impl Config {
     }
 
     /// Gets the hostaddrs that have been added to the configuration with `hostaddr`.
-    pub fn get_hostaddrs(&self) -> &[String] {
+    pub fn get_hostaddrs(&self) -> &[IpAddr] {
         self.hostaddr.deref()
     }
 
@@ -337,8 +340,8 @@ impl Config {
     ///
     /// Multiple hostaddrs can be specified by calling this method multiple times, and each will be tried in order.
     /// There must be either no hostaddrs, or the same number of hostaddrs as hosts.
-    pub fn hostaddr(&mut self, hostaddr: &str) -> &mut Config {
-        self.hostaddr.push(hostaddr.to_string());
+    pub fn hostaddr(&mut self, hostaddr: IpAddr) -> &mut Config {
+        self.hostaddr.push(hostaddr);
         self
     }
 
@@ -489,7 +492,10 @@ impl Config {
             }
             "hostaddr" => {
                 for hostaddr in value.split(',') {
-                    self.hostaddr(hostaddr);
+                    let addr = hostaddr
+                        .parse()
+                        .map_err(|_| Error::config_parse(Box::new(InvalidValue("hostaddr"))))?;
+                    self.hostaddr(addr);
                 }
             }
             "port" => {
@@ -1016,6 +1022,8 @@ impl<'a> UrlParser<'a> {
 
 #[cfg(test)]
 mod tests {
+    use std::net::IpAddr;
+
     use crate::{config::Host, Config};
 
     #[test]
@@ -1032,16 +1040,14 @@ mod tests {
             config.get_hosts(),
         );
 
-        assert_eq!(["127.0.0.1", "127.0.0.2"], config.get_hostaddrs(),);
+        assert_eq!(
+            [
+                "127.0.0.1".parse::<IpAddr>().unwrap(),
+                "127.0.0.2".parse::<IpAddr>().unwrap()
+            ],
+            config.get_hostaddrs(),
+        );
 
         assert_eq!(1, 1);
     }
-
-    #[test]
-    fn test_empty_hostaddrs() {
-        let s =
-            "user=pass_user dbname=postgres host=host1,host2,host3 hostaddr=127.0.0.1,,127.0.0.2";
-        let config = s.parse::<Config>().unwrap();
-        assert_eq!(["127.0.0.1", "", "127.0.0.2"], config.get_hostaddrs(),);
-    }
 }

tokio-postgres/src/connect.rs

@@ -5,8 +5,8 @@ use crate::connect_socket::connect_socket;
 use crate::tls::{MakeTlsConnect, TlsConnect};
 use crate::{Client, Config, Connection, Error, SimpleQueryMessage, Socket};
 use futures_util::{future, pin_mut, Future, FutureExt, Stream};
-use std::io;
 use std::task::Poll;
+use std::{cmp, io};
 
 pub async fn connect<T>(
     mut tls: T,
@@ -15,25 +15,35 @@ pub async fn connect<T>(
 where
     T: MakeTlsConnect<Socket>,
 {
-    if config.host.is_empty() {
-        return Err(Error::config("host missing".into()));
+    if config.host.is_empty() && config.hostaddr.is_empty() {
+        return Err(Error::config("both host and hostaddr are missing".into()));
     }
 
-    if config.port.len() > 1 && config.port.len() != config.host.len() {
-        return Err(Error::config("invalid number of ports".into()));
-    }
-
-    if !config.hostaddr.is_empty() && config.hostaddr.len() != config.host.len() {
+    if !config.host.is_empty()
+        && !config.hostaddr.is_empty()
+        && config.host.len() != config.hostaddr.len()
+    {
         let msg = format!(
-            "invalid number of hostaddrs ({}). Possible values: 0 or number of hosts ({})",
-            config.hostaddr.len(),
+            "number of hosts ({}) is different from number of hostaddrs ({})",
             config.host.len(),
+            config.hostaddr.len(),
         );
         return Err(Error::config(msg.into()));
     }
 
+    // At this point, either one of the following two scenarios could happen:
+    // (1) either config.host or config.hostaddr must be empty;
+    // (2) if both config.host and config.hostaddr are NOT empty; their lengths must be equal.
+    let num_hosts = cmp::max(config.host.len(), config.hostaddr.len());
+
+    if config.port.len() > 1 && config.port.len() != num_hosts {
+        return Err(Error::config("invalid number of ports".into()));
+    }
+
     let mut error = None;
-    for (i, host) in config.host.iter().enumerate() {
+    for i in 0..num_hosts {
+        let host = config.host.get(i);
+        let hostaddr = config.hostaddr.get(i);
         let port = config
             .port
             .get(i)
@@ -42,27 +52,30 @@ where
             .unwrap_or(5432);
 
         // The value of host is always used as the hostname for TLS validation.
+        // postgres doesn't support TLS over unix sockets, so the choice for Host::Unix variant here doesn't matter
         let hostname = match host {
-            Host::Tcp(host) => host.as_str(),
-            // postgres doesn't support TLS over unix sockets, so the choice here doesn't matter
-            #[cfg(unix)]
-            Host::Unix(_) => "",
+            Some(Host::Tcp(host)) => host.as_str(),
+            _ => "",
         };
         let tls = tls
             .make_tls_connect(hostname)
             .map_err(|e| Error::tls(e.into()))?;
 
-        // If both host and hostaddr are specified, the value of hostaddr is used to to establish the TCP connection.
-        let hostaddr = match host {
-            Host::Tcp(_hostname) => match config.hostaddr.get(i) {
-                Some(hostaddr) if hostaddr.is_empty() => Host::Tcp(hostaddr.clone()),
-                _ => host.clone(),
-            },
-            #[cfg(unix)]
-            Host::Unix(_v) => host.clone(),
+        // Try to use the value of hostaddr to establish the TCP connection,
+        // fallback to host if hostaddr is not present.
+        let addr = match hostaddr {
+            Some(ipaddr) => Host::Tcp(ipaddr.to_string()),
+            None => {
+                if let Some(host) = host {
+                    host.clone()
+                } else {
+                    // This is unreachable.
+                    return Err(Error::config("both host and hostaddr are empty".into()));
+                }
+            }
         };
 
-        match connect_once(&hostaddr, port, tls, config).await {
+        match connect_once(&addr, port, tls, config).await {
             Ok((client, connection)) => return Ok((client, connection)),
             Err(e) => error = Some(e),
         }

tokio-postgres/tests/test/main.rs

@@ -147,6 +147,58 @@ async fn scram_password_ok() {
     connect("user=scram_user password=password dbname=postgres").await;
 }
 
+#[tokio::test]
+async fn host_only_ok() {
+    let _ = tokio_postgres::connect(
+        "host=localhost port=5433 user=pass_user dbname=postgres password=password",
+        NoTls,
+    )
+    .await
+    .unwrap();
+}
+
+#[tokio::test]
+async fn hostaddr_only_ok() {
+    let _ = tokio_postgres::connect(
+        "hostaddr=127.0.0.1 port=5433 user=pass_user dbname=postgres password=password",
+        NoTls,
+    )
+    .await
+    .unwrap();
+}
+
+#[tokio::test]
+async fn hostaddr_and_host_ok() {
+    let _ = tokio_postgres::connect(
+        "hostaddr=127.0.0.1 host=localhost port=5433 user=pass_user dbname=postgres password=password",
+        NoTls,
+    )
+    .await
+    .unwrap();
+}
+
+#[tokio::test]
+async fn hostaddr_host_mismatch() {
+    let _ = tokio_postgres::connect(
+        "hostaddr=127.0.0.1,127.0.0.2 host=localhost port=5433 user=pass_user dbname=postgres password=password",
+        NoTls,
+    )
+    .await
+    .err()
+    .unwrap();
+}
+
+#[tokio::test]
+async fn hostaddr_host_both_missing() {
+    let _ = tokio_postgres::connect(
+        "port=5433 user=pass_user dbname=postgres password=password",
+        NoTls,
+    )
+    .await
+    .err()
+    .unwrap();
+}
+
 #[tokio::test]
 async fn pipelined_prepare() {
     let client = connect("user=postgres").await;