agent_first_psql/
conn.rs

1use crate::types::{RuntimeConfig, SessionConfig};
2use std::error::Error as _;
3use tokio_postgres::Config;
4use tokio_postgres::config::{Host, SslMode};
5
6const SUPPORTED_SSLMODE_HINT: &str = "afpsql supports sslmode=disable, prefer, and require. It does not implement libpq verify-ca/verify-full or client certificate options yet; use psql/libpq when certificate verification or client certificates are required.";
7const DEFAULT_POSTGRES_HOST: &str = "127.0.0.1";
8const DEFAULT_POSTGRES_PORT: u16 = 5432;
9
10#[derive(Debug, Clone, PartialEq, Eq)]
11pub(crate) enum PostgresEndpoint {
12    Tcp { host: String, port: u16 },
13    UnixSocket { directory: String, port: u16 },
14}
15
16#[derive(Debug, Clone, PartialEq, Eq)]
17pub struct ConnectionConfigError {
18    message: String,
19    hint: Option<String>,
20}
21
22impl ConnectionConfigError {
23    pub fn new(message: impl Into<String>, hint: Option<String>) -> Self {
24        Self {
25            message: message.into(),
26            hint,
27        }
28    }
29
30    pub fn message(&self) -> &str {
31        &self.message
32    }
33
34    pub fn hint(&self) -> Option<&str> {
35        self.hint.as_deref()
36    }
37}
38
39impl std::fmt::Display for ConnectionConfigError {
40    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
41        f.write_str(&self.message)
42    }
43}
44
45impl std::error::Error for ConnectionConfigError {}
46
47pub fn resolve_session_name(cfg: &RuntimeConfig, requested: Option<&str>) -> String {
48    requested
49        .map(std::string::ToString::to_string)
50        .unwrap_or_else(|| cfg.default_session.clone())
51}
52
53pub fn resolve_pg_config(cfg: &SessionConfig) -> Result<Config, ConnectionConfigError> {
54    // An administrator-locked readonly profile pins the endpoint, and the CLI
55    // already rejects every connection flag. Environment variables were the one
56    // remaining way to move such an executable off its target, so a pinned
57    // session reads none of them — including PGSSLMODE, which would otherwise
58    // let a caller downgrade TLS on a connection it does not control.
59    let env = |name: &str| {
60        if cfg.profile_pinned {
61            None
62        } else {
63            env_nonempty(name)
64        }
65    };
66
67    if let Some(dsn) = cfg.dsn_secret.clone().or_else(|| env("AFPSQL_DSN_SECRET")) {
68        validate_dsn_ssl_options(&dsn)?;
69        return dsn.parse().map_err(|e| map_pg_config_parse_error("dsn", e));
70    }
71
72    if let Some(conninfo) = cfg
73        .conninfo_secret
74        .clone()
75        .or_else(|| env("AFPSQL_CONNINFO_SECRET"))
76    {
77        validate_conninfo_ssl_options(&conninfo)?;
78        return conninfo
79            .parse()
80            .map_err(|e| map_pg_config_parse_error("conninfo", e));
81    }
82
83    let host = cfg
84        .host
85        .clone()
86        .or_else(|| env("AFPSQL_HOST"))
87        .or_else(|| env("PGHOST"))
88        .unwrap_or_else(|| "127.0.0.1".to_string());
89    let port = cfg
90        .port
91        .or_else(|| env("AFPSQL_PORT").and_then(|s| s.parse().ok()))
92        .or_else(|| env("PGPORT").and_then(|s| s.parse().ok()))
93        .unwrap_or(5432);
94    let user = cfg
95        .user
96        .clone()
97        .or_else(|| env("AFPSQL_USER"))
98        .or_else(|| env("PGUSER"))
99        .unwrap_or_else(|| "postgres".to_string());
100    let dbname = cfg
101        .dbname
102        .clone()
103        .or_else(|| env("AFPSQL_DBNAME"))
104        .or_else(|| env("PGDATABASE"))
105        .unwrap_or_else(|| "postgres".to_string());
106    let password = cfg
107        .password_secret
108        .clone()
109        .or_else(|| env("AFPSQL_PASSWORD_SECRET"))
110        .or_else(|| env("PGPASSWORD"));
111
112    let mut pg_cfg = Config::new();
113    pg_cfg.host(host).port(port).user(user).dbname(dbname);
114    if let Some(pw) = password {
115        pg_cfg.password(pw);
116    }
117    if let Some(sslmode) = env("PGSSLMODE") {
118        apply_sslmode(&mut pg_cfg, "PGSSLMODE", &sslmode)?;
119    }
120    Ok(pg_cfg)
121}
122
123pub(crate) fn resolve_single_postgres_endpoint(
124    pg_cfg: &Config,
125    transport_name: &str,
126) -> Result<PostgresEndpoint, String> {
127    let hosts = pg_cfg.get_hosts();
128    let hostaddrs = pg_cfg.get_hostaddrs();
129    let ports = pg_cfg.get_ports();
130    if hosts.len() > 1 || hostaddrs.len() > 1 || ports.len() > 1 {
131        return Err(format!(
132            "{transport_name} supports a single PostgreSQL host and port; the connection source resolved to multiple targets"
133        ));
134    }
135
136    let port = ports.first().copied().unwrap_or(DEFAULT_POSTGRES_PORT);
137    if let Some(hostaddr) = hostaddrs.first() {
138        #[cfg(unix)]
139        if matches!(hosts.first(), Some(Host::Unix(_))) {
140            return Err(format!(
141                "{transport_name} cannot combine a PostgreSQL Unix socket with hostaddr"
142            ));
143        }
144        return Ok(PostgresEndpoint::Tcp {
145            host: hostaddr.to_string(),
146            port,
147        });
148    }
149
150    match hosts.first() {
151        Some(Host::Tcp(host)) if host.starts_with('/') => Ok(PostgresEndpoint::UnixSocket {
152            directory: host.clone(),
153            port,
154        }),
155        Some(Host::Tcp(host)) => Ok(PostgresEndpoint::Tcp {
156            host: host.clone(),
157            port,
158        }),
159        #[cfg(unix)]
160        Some(Host::Unix(path)) => Ok(PostgresEndpoint::UnixSocket {
161            directory: path.to_string_lossy().into_owned(),
162            port,
163        }),
164        None => Ok(PostgresEndpoint::Tcp {
165            host: DEFAULT_POSTGRES_HOST.to_string(),
166            port,
167        }),
168    }
169}
170
171/// TLS server name for a stdio-bridged connection.
172///
173/// Derived from the endpoint the bridge actually targets, so it is never empty:
174/// OpenSSL rejects a zero-length SNI name outright, which would turn a hostless
175/// DSN into a TLS handshake failure. An IP literal or a Unix-socket target
176/// yields an address rather than a name, and the TLS backend then skips SNI —
177/// which is the correct behavior for both.
178pub(crate) fn postgres_tls_server_name(endpoint: &PostgresEndpoint) -> String {
179    match endpoint {
180        PostgresEndpoint::Tcp { host, .. } => host.clone(),
181        PostgresEndpoint::UnixSocket { .. } => DEFAULT_POSTGRES_HOST.to_string(),
182    }
183}
184
185pub(crate) fn make_supported_tls()
186-> Result<postgres_native_tls::MakeTlsConnector, native_tls::Error> {
187    let tls = native_tls::TlsConnector::builder()
188        // Supported sslmode=prefer/require encrypts without certificate
189        // verification. verify-ca/verify-full are rejected during config parse.
190        .danger_accept_invalid_certs(true)
191        .danger_accept_invalid_hostnames(true)
192        .build()?;
193    Ok(postgres_native_tls::MakeTlsConnector::new(tls))
194}
195
196fn env_nonempty(name: &str) -> Option<String> {
197    std::env::var(name).ok().filter(|value| !value.is_empty())
198}
199
200pub fn libpq_env_fallbacks_in_use(cfg: &SessionConfig) -> Vec<&'static str> {
201    // `resolve_pg_config` reads no environment at all for a pinned profile, so
202    // reporting a fallback here would name variables that had no effect on the
203    // connection.
204    if cfg.profile_pinned {
205        return Vec::new();
206    }
207    if cfg.dsn_secret.is_some() || cfg.conninfo_secret.is_some() {
208        return Vec::new();
209    }
210    if std::env::var("AFPSQL_DSN_SECRET").is_ok() || std::env::var("AFPSQL_CONNINFO_SECRET").is_ok()
211    {
212        return Vec::new();
213    }
214    let mut used = Vec::new();
215    if cfg.host.is_none()
216        && std::env::var("AFPSQL_HOST").is_err()
217        && env_nonempty("PGHOST").is_some()
218    {
219        used.push("PGHOST");
220    }
221    if cfg.port.is_none()
222        && std::env::var("AFPSQL_PORT").is_err()
223        && env_nonempty("PGPORT").is_some()
224    {
225        used.push("PGPORT");
226    }
227    if cfg.user.is_none()
228        && std::env::var("AFPSQL_USER").is_err()
229        && env_nonempty("PGUSER").is_some()
230    {
231        used.push("PGUSER");
232    }
233    if cfg.dbname.is_none()
234        && std::env::var("AFPSQL_DBNAME").is_err()
235        && env_nonempty("PGDATABASE").is_some()
236    {
237        used.push("PGDATABASE");
238    }
239    if cfg.password_secret.is_none()
240        && std::env::var("AFPSQL_PASSWORD_SECRET").is_err()
241        && env_nonempty("PGPASSWORD").is_some()
242    {
243        used.push("PGPASSWORD");
244    }
245    if env_nonempty("PGSSLMODE").is_some() {
246        used.push("PGSSLMODE");
247    }
248    used
249}
250
251fn validate_dsn_ssl_options(dsn: &str) -> Result<(), ConnectionConfigError> {
252    let Some(query) = dsn.split_once('?').map(|(_, query)| query) else {
253        return Ok(());
254    };
255    let query = query.split('#').next().unwrap_or(query);
256    for part in query.split('&') {
257        let (key, value) = part.split_once('=').unwrap_or((part, ""));
258        validate_ssl_option(key, value, "dsn")?;
259    }
260    Ok(())
261}
262
263fn validate_conninfo_ssl_options(conninfo: &str) -> Result<(), ConnectionConfigError> {
264    for (key, value) in parse_conninfo_pairs(conninfo) {
265        validate_ssl_option(&key, &value, "conninfo")?;
266    }
267    Ok(())
268}
269
270fn validate_ssl_option(key: &str, value: &str, source: &str) -> Result<(), ConnectionConfigError> {
271    match key {
272        "sslmode" => validate_sslmode(source, value),
273        "sslnegotiation" if value == "postgres" => Ok(()),
274        "sslnegotiation" => Err(ConnectionConfigError::new(
275            format!("unsupported {source} TLS option `sslnegotiation={value}`"),
276            Some("afpsql supports PostgreSQL's standard TLS negotiation path only; remove sslnegotiation=direct or use psql/libpq for PostgreSQL 17 direct TLS negotiation.".to_string()),
277        )),
278        "sslrootcert" | "sslcert" | "sslkey" | "sslpassword" | "sslcrl" | "sslcrldir"
279        | "sslcertmode" | "sslsni" | "ssl_min_protocol_version" | "ssl_max_protocol_version"
280        => Err(unsupported_ssl_option(source, key)),
281        _ => Ok(()),
282    }
283}
284
285fn apply_sslmode(
286    pg_cfg: &mut Config,
287    source: &str,
288    value: &str,
289) -> Result<(), ConnectionConfigError> {
290    validate_sslmode(source, value)?;
291    let mode = match value {
292        "disable" => SslMode::Disable,
293        "prefer" => SslMode::Prefer,
294        "require" => SslMode::Require,
295        _ => return Err(unsupported_sslmode(source, value)),
296    };
297    pg_cfg.ssl_mode(mode);
298    Ok(())
299}
300
301fn validate_sslmode(source: &str, value: &str) -> Result<(), ConnectionConfigError> {
302    match value {
303        "disable" | "prefer" | "require" => Ok(()),
304        _ => Err(unsupported_sslmode(source, value)),
305    }
306}
307
308fn unsupported_sslmode(source: &str, value: &str) -> ConnectionConfigError {
309    ConnectionConfigError::new(
310        format!(
311            "unsupported {source} sslmode `{value}`; supported values are disable, prefer, require"
312        ),
313        Some(SUPPORTED_SSLMODE_HINT.to_string()),
314    )
315}
316
317fn unsupported_ssl_option(source: &str, key: &str) -> ConnectionConfigError {
318    ConnectionConfigError::new(
319        format!("unsupported {source} TLS option `{key}`"),
320        Some(SUPPORTED_SSLMODE_HINT.to_string()),
321    )
322}
323
324fn map_pg_config_parse_error(source: &str, err: tokio_postgres::Error) -> ConnectionConfigError {
325    let cause = err.source().map(std::string::ToString::to_string);
326    if let Some(cause) = cause.as_deref() {
327        if cause == "invalid value for option `sslmode`" {
328            return ConnectionConfigError::new(
329                format!("unsupported {source} sslmode"),
330                Some(SUPPORTED_SSLMODE_HINT.to_string()),
331            );
332        }
333        if let Some(key) = cause
334            .strip_prefix("unknown option `")
335            .and_then(|rest| rest.strip_suffix('`'))
336            && is_unsupported_ssl_option(key)
337        {
338            return unsupported_ssl_option(source, key);
339        }
340    }
341
342    let detail = cause
343        .map(|cause| format!("{err}: {cause}"))
344        .unwrap_or_else(|| err.to_string());
345    ConnectionConfigError::new(format!("invalid {source}: {detail}"), None)
346}
347
348fn is_unsupported_ssl_option(key: &str) -> bool {
349    matches!(
350        key,
351        "sslrootcert"
352            | "sslcert"
353            | "sslkey"
354            | "sslpassword"
355            | "sslcrl"
356            | "sslcrldir"
357            | "sslcertmode"
358            | "sslsni"
359            | "ssl_min_protocol_version"
360            | "ssl_max_protocol_version"
361            | "sslnegotiation"
362    )
363}
364
365fn parse_conninfo_pairs(input: &str) -> Vec<(String, String)> {
366    let bytes = input.as_bytes();
367    let mut pairs = Vec::new();
368    let mut i = 0usize;
369
370    while i < bytes.len() {
371        while i < bytes.len() && bytes[i].is_ascii_whitespace() {
372            i += 1;
373        }
374        if i >= bytes.len() {
375            break;
376        }
377
378        let key_start = i;
379        while i < bytes.len() && bytes[i] != b'=' && !bytes[i].is_ascii_whitespace() {
380            i += 1;
381        }
382        let key = &input[key_start..i];
383        while i < bytes.len() && bytes[i].is_ascii_whitespace() {
384            i += 1;
385        }
386        if i >= bytes.len() || bytes[i] != b'=' {
387            break;
388        }
389        i += 1;
390        while i < bytes.len() && bytes[i].is_ascii_whitespace() {
391            i += 1;
392        }
393
394        let mut value = String::new();
395        if i < bytes.len() && bytes[i] == b'\'' {
396            i += 1;
397            while i < bytes.len() {
398                match bytes[i] {
399                    b'\\' if i + 1 < bytes.len() => {
400                        i += 1;
401                        value.push(bytes[i] as char);
402                        i += 1;
403                    }
404                    b'\'' => {
405                        i += 1;
406                        break;
407                    }
408                    b => {
409                        value.push(b as char);
410                        i += 1;
411                    }
412                }
413            }
414        } else {
415            while i < bytes.len() && !bytes[i].is_ascii_whitespace() {
416                if bytes[i] == b'\\' && i + 1 < bytes.len() {
417                    i += 1;
418                }
419                value.push(bytes[i] as char);
420                i += 1;
421            }
422        }
423
424        pairs.push((key.to_string(), value));
425    }
426
427    pairs
428}
429
430#[cfg(test)]
431#[path = "../tests/support/unit_conn.rs"]
432mod tests;
agent_first_psql/conn.rs

agent_first_psql/
conn.rs
