Merge pull request #945 from trungda/tdinh/hostaddr

Add hostaddr support

Author: sfackler

postgres/src/config.rs

@@ -6,6 +6,7 @@ use crate::connection::Connection;
 use crate::Client;
 use log::info;
 use std::fmt;
+use std::net::IpAddr;
 use std::path::Path;
 use std::str::FromStr;
 use std::sync::Arc;
@@ -39,6 +40,19 @@ use tokio_postgres::{Error, Socket};
 ///     path to the directory containing Unix domain sockets. Otherwise, it is treated as a hostname. Multiple hosts
 ///     can be specified, separated by commas. Each host will be tried in turn when connecting. Required if connecting
 ///     with the `connect` method.
+/// * `hostaddr` - Numeric IP address of host to connect to. This should be in the standard IPv4 address format,
+///     e.g., 172.28.40.9. If your machine supports IPv6, you can also use those addresses.
+///     If this parameter is not specified, the value of `host` will be looked up to find the corresponding IP address,
+///     - 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.
+///     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,
+///             in which case it will be used as the host name.
 /// * `port` - The port to connect to. Multiple ports can be specified, separated by commas. The number of ports must be
 ///     either 1, in which case it will be used for all hosts, or the same as the number of hosts. Defaults to 5432 if
 ///     omitted or the empty string.
@@ -70,6 +84,10 @@ use tokio_postgres::{Error, Socket};
 /// ```
 ///
 /// ```not_rust
+/// host=host1,host2,host3 port=1234,,5678 hostaddr=127.0.0.1,127.0.0.2,127.0.0.3 user=postgres target_session_attrs=read-write
+/// ```
+///
+/// ```not_rust
 /// host=host1,host2,host3 port=1234,,5678 user=postgres target_session_attrs=read-write
 /// ```
 ///
@@ -207,6 +225,7 @@ impl Config {
     ///
     /// Multiple hosts can be specified by calling this method multiple times, and each will be tried in order. On Unix
     /// systems, a host starting with a `/` is interpreted as a path to a directory containing Unix domain sockets.
+    /// There must be either no hosts, or the same number of hosts as hostaddrs.
     pub fn host(&mut self, host: &str) -> &mut Config {
         self.config.host(host);
         self
@@ -217,6 +236,11 @@ impl Config {
         self.config.get_hosts()
     }
 
+    /// Gets the hostaddrs that have been added to the configuration with `hostaddr`.
+    pub fn get_hostaddrs(&self) -> &[IpAddr] {
+        self.config.get_hostaddrs()
+    }
+
     /// Adds a Unix socket host to the configuration.
     ///
     /// Unlike `host`, this method allows non-UTF8 paths.
@@ -229,6 +253,15 @@ impl Config {
         self
     }
 
+    /// Adds a hostaddr to the configuration.
+    ///
+    /// 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: IpAddr) -> &mut Config {
+        self.config.hostaddr(hostaddr);
+        self
+    }
+
     /// Adds a port to the configuration.
     ///
     /// Multiple ports can be specified by calling this method multiple times. There must either be no ports, in which

tokio-postgres/src/config.rs

@@ -14,6 +14,8 @@ 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;
 #[cfg(unix)]
@@ -92,6 +94,19 @@ pub enum Host {
 ///     path to the directory containing Unix domain sockets. Otherwise, it is treated as a hostname. Multiple hosts
 ///     can be specified, separated by commas. Each host will be tried in turn when connecting. Required if connecting
 ///     with the `connect` method.
+/// * `hostaddr` - Numeric IP address of host to connect to. This should be in the standard IPv4 address format,
+///     e.g., 172.28.40.9. If your machine supports IPv6, you can also use those addresses.
+///     If this parameter is not specified, the value of `host` will be looked up to find the corresponding IP address,
+///     - 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.
+///     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,
+///             in which case it will be used as the host name.
 /// * `port` - The port to connect to. Multiple ports can be specified, separated by commas. The number of ports must be
 ///     either 1, in which case it will be used for all hosts, or the same as the number of hosts. Defaults to 5432 if
 ///     omitted or the empty string.
@@ -126,6 +141,10 @@ pub enum Host {
 /// ```
 ///
 /// ```not_rust
+/// host=host1,host2,host3 port=1234,,5678 hostaddr=127.0.0.1,127.0.0.2,127.0.0.3 user=postgres target_session_attrs=read-write
+/// ```
+///
+/// ```not_rust
 /// host=host1,host2,host3 port=1234,,5678 user=postgres target_session_attrs=read-write
 /// ```
 ///
@@ -162,6 +181,7 @@ pub struct Config {
     pub(crate) application_name: Option<String>,
     pub(crate) ssl_mode: SslMode,
     pub(crate) host: Vec<Host>,
+    pub(crate) hostaddr: Vec<IpAddr>,
     pub(crate) port: Vec<u16>,
     pub(crate) connect_timeout: Option<Duration>,
     pub(crate) tcp_user_timeout: Option<Duration>,
@@ -189,6 +209,7 @@ impl Config {
             application_name: None,
             ssl_mode: SslMode::Prefer,
             host: vec![],
+            hostaddr: vec![],
             port: vec![],
             connect_timeout: None,
             tcp_user_timeout: None,
@@ -288,6 +309,7 @@ impl Config {
     ///
     /// Multiple hosts can be specified by calling this method multiple times, and each will be tried in order. On Unix
     /// systems, a host starting with a `/` is interpreted as a path to a directory containing Unix domain sockets.
+    /// There must be either no hosts, or the same number of hosts as hostaddrs.
     pub fn host(&mut self, host: &str) -> &mut Config {
         #[cfg(unix)]
         {
@@ -305,6 +327,11 @@ impl Config {
         &self.host
     }
 
+    /// Gets the hostaddrs that have been added to the configuration with `hostaddr`.
+    pub fn get_hostaddrs(&self) -> &[IpAddr] {
+        self.hostaddr.deref()
+    }
+
     /// Adds a Unix socket host to the configuration.
     ///
     /// Unlike `host`, this method allows non-UTF8 paths.
@@ -317,6 +344,15 @@ impl Config {
         self
     }
 
+    /// Adds a hostaddr to the configuration.
+    ///
+    /// 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: IpAddr) -> &mut Config {
+        self.hostaddr.push(hostaddr);
+        self
+    }
+
     /// Adds a port to the configuration.
     ///
     /// Multiple ports can be specified by calling this method multiple times. There must either be no ports, in which
@@ -484,6 +520,14 @@ impl Config {
                     self.host(host);
                 }
             }
+            "hostaddr" => {
+                for hostaddr in value.split(',') {
+                    let addr = hostaddr
+                        .parse()
+                        .map_err(|_| Error::config_parse(Box::new(InvalidValue("hostaddr"))))?;
+                    self.hostaddr(addr);
+                }
+            }
             "port" => {
                 for port in value.split(',') {
                     let port = if port.is_empty() {
@@ -635,6 +679,7 @@ impl fmt::Debug for Config {
             .field("application_name", &self.application_name)
             .field("ssl_mode", &self.ssl_mode)
             .field("host", &self.host)
+            .field("hostaddr", &self.hostaddr)
             .field("port", &self.port)
             .field("connect_timeout", &self.connect_timeout)
             .field("tcp_user_timeout", &self.tcp_user_timeout)
@@ -1025,3 +1070,41 @@ impl<'a> UrlParser<'a> {
             .map_err(|e| Error::config_parse(e.into()))
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use std::net::IpAddr;
+
+    use crate::{config::Host, Config};
+
+    #[test]
+    fn test_simple_parsing() {
+        let s = "user=pass_user dbname=postgres host=host1,host2 hostaddr=127.0.0.1,127.0.0.2 port=26257";
+        let config = s.parse::<Config>().unwrap();
+        assert_eq!(Some("pass_user"), config.get_user());
+        assert_eq!(Some("postgres"), config.get_dbname());
+        assert_eq!(
+            [
+                Host::Tcp("host1".to_string()),
+                Host::Tcp("host2".to_string())
+            ],
+            config.get_hosts(),
+        );
+
+        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_invalid_hostaddr_parsing() {
+        let s = "user=pass_user dbname=postgres host=host1 hostaddr=127.0.0 port=26257";
+        s.parse::<Config>().err().unwrap();
+    }
+}

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,35 +15,70 @@ 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() {
+    if !config.host.is_empty()
+        && !config.hostaddr.is_empty()
+        && config.host.len() != config.hostaddr.len()
+    {
+        let msg = format!(
+            "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)
             .or_else(|| config.port.first())
             .copied()
             .unwrap_or(5432);
 
+        // The value of host is used as the hostname for TLS validation,
+        // if it's not present, use the value of hostaddr.
         let hostname = match host {
-            Host::Tcp(host) => host.as_str(),
+            Some(Host::Tcp(host)) => host.clone(),
             // postgres doesn't support TLS over unix sockets, so the choice here doesn't matter
             #[cfg(unix)]
-            Host::Unix(_) => "",
+            Some(Host::Unix(_)) => "".to_string(),
+            None => hostaddr.map_or("".to_string(), |ipaddr| ipaddr.to_string()),
         };
-
         let tls = tls
-            .make_tls_connect(hostname)
+            .make_tls_connect(&hostname)
             .map_err(|e| Error::tls(e.into()))?;
 
-        match connect_once(host, port, tls, config).await {
+        // 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(&addr, port, tls, config).await {
             Ok((client, connection)) => return Ok((client, connection)),
             Err(e) => error = Some(e),
         }

tokio-postgres/tests/test/runtime.rs

@@ -66,6 +66,58 @@ async fn target_session_attrs_err() {
     .unwrap();
 }
 
+#[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 cancel_query() {
     let client = connect("host=localhost port=5433 user=postgres").await;