encore-tokio-postgres 0.7.13

use std::collections::HashMap;
use std::{cmp, io};

use rand::seq::SliceRandom;
use tokio::net;
use tokio_util::codec::Framed;

use crate::client::{Addr, SocketConfig};
use crate::codec::PostgresCodec;
use crate::config::{Host, LoadBalanceHosts};
use crate::connect_raw::connect_proxy_raw;
use crate::connect_socket::connect_socket;
use crate::maybe_tls_stream::MaybeTlsStream;
use crate::tls::MakeTlsConnect;
use crate::{Config, Error, Socket};

pub(crate) struct ProxyInfo<T>
where
    T: MakeTlsConnect<Socket>,
{
    pub backend: Framed<MaybeTlsStream<Socket, T::Stream>, PostgresCodec>,
    pub socket_config: SocketConfig,
    /// The backend's cancellation key, as process ID and secret key, if it
    /// gave one.
    pub backend_key: Option<(i32, i32)>,
    pub parameters: HashMap<String, String>,
}

/// Connects to the backend `config` describes on behalf of a proxy's client,
/// starting up with the parameters the client passes on.
pub(crate) async fn connect_proxy<T>(
    tls: &mut T,
    config: &Config,
    client_params: &[(String, String)],
) -> Result<ProxyInfo<T>, Error>
where
    T: MakeTlsConnect<Socket>,
{
    if config.host.is_empty() && config.hostaddr.is_empty() {
        return Err(Error::config("both host and hostaddr are missing".into()));
    }

    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 indices = (0..num_hosts).collect::<Vec<_>>();
    if config.load_balance_hosts == LoadBalanceHosts::Random {
        indices.shuffle(&mut rand::rng());
    }

    let mut error = None;
    for i in indices {
        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,
        let hostname = match host {
            Some(Host::Tcp(host)) => Some(host.clone()),
            // postgres doesn't support TLS over unix sockets, so the choice here doesn't matter
            #[cfg(unix)]
            Some(Host::Unix(_)) => None,
            None => None,
        };

        // 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 => host.cloned().unwrap(),
        };

        let addr_str = match addr {
            Host::Tcp(ref host) => host.to_owned(),
            _ => "".to_string(),
        };

        match connect_host(addr, hostname, port, tls, config, client_params).await {
            Ok(info) => return Ok(info),
            Err(e) => {
                log::error!(
                    "failed connecting to database at {}:{}, {e}",
                    addr_str,
                    port
                );
                error = Some(e)
            }
        }
    }

    Err(error.unwrap())
}

async fn connect_host<T>(
    host: Host,
    hostname: Option<String>,
    port: u16,
    tls: &mut T,
    config: &Config,
    client_params: &[(String, String)],
) -> Result<ProxyInfo<T>, Error>
where
    T: MakeTlsConnect<Socket>,
{
    match host {
        Host::Tcp(host) => {
            let mut addrs = net::lookup_host((&*host, port))
                .await
                .map_err(Error::connect)?
                .collect::<Vec<_>>();

            if config.load_balance_hosts == LoadBalanceHosts::Random {
                addrs.shuffle(&mut rand::rng());
            }

            let mut last_err = None;
            for addr in addrs {
                match connect_once(
                    Addr::Tcp(addr.ip()),
                    hostname.as_deref(),
                    port,
                    tls,
                    config,
                    client_params,
                )
                .await
                {
                    Ok(stream) => return Ok(stream),
                    Err(e) => {
                        last_err = Some(e);
                        continue;
                    }
                };
            }

            Err(last_err.unwrap_or_else(|| {
                Error::connect(io::Error::new(
                    io::ErrorKind::InvalidInput,
                    "could not resolve any addresses",
                ))
            }))
        }
        #[cfg(unix)]
        Host::Unix(path) => {
            connect_once(
                Addr::Unix(path),
                hostname.as_deref(),
                port,
                tls,
                config,
                client_params,
            )
            .await
        }
    }
}

async fn connect_once<T>(
    addr: Addr,
    hostname: Option<&str>,
    port: u16,
    tls: &mut T,
    config: &Config,
    client_params: &[(String, String)],
) -> Result<ProxyInfo<T>, Error>
where
    T: MakeTlsConnect<Socket>,
{
    let socket = connect_socket(
        &addr,
        port,
        config.connect_timeout,
        config.tcp_user_timeout,
        if config.keepalives {
            Some(&config.keepalive_config)
        } else {
            None
        },
    )
    .await?;

    let tls = tls
        .make_tls_connect(hostname.unwrap_or(""))
        .map_err(|e| Error::tls(e.into()))?;
    let has_hostname = hostname.is_some();

    let (stream, backend_key, parameters) =
        connect_proxy_raw(socket, tls, has_hostname, config, client_params).await?;

    let socket_config = SocketConfig {
        addr,
        hostname: hostname.map(|s| s.to_string()),
        port,
        connect_timeout: config.connect_timeout,
        tcp_user_timeout: config.tcp_user_timeout,
        keepalive: if config.keepalives {
            Some(config.keepalive_config.clone())
        } else {
            None
        },
    };
    Ok(ProxyInfo {
        backend: stream,
        socket_config,
        backend_key,
        parameters,
    })
}