Skip to content

Add hostaddr support #945

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 17 commits into from
Jul 17, 2023
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
33 changes: 33 additions & 0 deletions postgres/src/config.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -67,6 +81,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
/// ```
///
Expand Down Expand Up @@ -204,6 +222,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
Expand All @@ -214,6 +233,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.
Expand All @@ -226,6 +250,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
Expand Down
83 changes: 83 additions & 0 deletions tokio-postgres/src/config.rs
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,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)]
Expand Down Expand Up @@ -91,6 +93,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.
Expand Down Expand Up @@ -122,6 +137,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
/// ```
///
Expand Down Expand Up @@ -158,6 +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<IpAddr>,
pub(crate) port: Vec<u16>,
pub(crate) connect_timeout: Option<Duration>,
pub(crate) keepalives: bool,
Expand Down Expand Up @@ -188,6 +208,7 @@ impl Config {
application_name: None,
ssl_mode: SslMode::Prefer,
host: vec![],
hostaddr: vec![],
port: vec![],
connect_timeout: None,
keepalives: true,
Expand Down Expand Up @@ -281,6 +302,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)]
{
Expand All @@ -298,6 +320,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.
Expand All @@ -310,6 +337,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
Expand Down Expand Up @@ -455,6 +491,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() {
Expand Down Expand Up @@ -593,6 +637,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("keepalives", &self.keepalives)
Expand Down Expand Up @@ -975,3 +1020,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();
}
}
55 changes: 45 additions & 10 deletions tokio-postgres/src/connect.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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),
}
Expand Down
52 changes: 52 additions & 0 deletions tokio-postgres/tests/test/runtime.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down