agent_first_psql/
conn.rs

1use crate::types::{RuntimeConfig, SessionConfig};
2use std::error::Error as _;
3use tokio_postgres::Config;
4use tokio_postgres::config::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.";
7
8#[derive(Debug, Clone, PartialEq, Eq)]
9pub struct ConnectionConfigError {
10    message: String,
11    hint: Option<String>,
12}
13
14impl ConnectionConfigError {
15    pub fn new(message: impl Into<String>, hint: Option<String>) -> Self {
16        Self {
17            message: message.into(),
18            hint,
19        }
20    }
21
22    pub fn message(&self) -> &str {
23        &self.message
24    }
25
26    pub fn hint(&self) -> Option<&str> {
27        self.hint.as_deref()
28    }
29}
30
31impl std::fmt::Display for ConnectionConfigError {
32    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
33        f.write_str(&self.message)
34    }
35}
36
37impl std::error::Error for ConnectionConfigError {}
38
39pub fn resolve_session_name(cfg: &RuntimeConfig, requested: Option<&str>) -> String {
40    requested
41        .map(std::string::ToString::to_string)
42        .unwrap_or_else(|| cfg.default_session.clone())
43}
44
45pub fn resolve_pg_config(cfg: &SessionConfig) -> Result<Config, ConnectionConfigError> {
46    if let Some(dsn) = cfg
47        .dsn_secret
48        .clone()
49        .or_else(|| std::env::var("AFPSQL_DSN_SECRET").ok())
50    {
51        validate_dsn_ssl_options(&dsn)?;
52        return dsn.parse().map_err(|e| map_pg_config_parse_error("dsn", e));
53    }
54
55    if let Some(conninfo) = cfg
56        .conninfo_secret
57        .clone()
58        .or_else(|| std::env::var("AFPSQL_CONNINFO_SECRET").ok())
59    {
60        validate_conninfo_ssl_options(&conninfo)?;
61        return conninfo
62            .parse()
63            .map_err(|e| map_pg_config_parse_error("conninfo", e));
64    }
65
66    let host = cfg
67        .host
68        .clone()
69        .or_else(|| std::env::var("AFPSQL_HOST").ok())
70        .or_else(|| std::env::var("PGHOST").ok())
71        .unwrap_or_else(|| "127.0.0.1".to_string());
72    let port = cfg
73        .port
74        .or_else(|| {
75            std::env::var("AFPSQL_PORT")
76                .ok()
77                .and_then(|s| s.parse().ok())
78        })
79        .or_else(|| std::env::var("PGPORT").ok().and_then(|s| s.parse().ok()))
80        .unwrap_or(5432);
81    let user = cfg
82        .user
83        .clone()
84        .or_else(|| std::env::var("AFPSQL_USER").ok())
85        .or_else(|| std::env::var("PGUSER").ok())
86        .unwrap_or_else(|| "postgres".to_string());
87    let dbname = cfg
88        .dbname
89        .clone()
90        .or_else(|| std::env::var("AFPSQL_DBNAME").ok())
91        .or_else(|| std::env::var("PGDATABASE").ok())
92        .unwrap_or_else(|| "postgres".to_string());
93    let password = cfg
94        .password_secret
95        .clone()
96        .or_else(|| std::env::var("AFPSQL_PASSWORD_SECRET").ok())
97        .or_else(|| std::env::var("PGPASSWORD").ok());
98
99    let mut pg_cfg = Config::new();
100    pg_cfg.host(host).port(port).user(user).dbname(dbname);
101    if let Some(pw) = password {
102        pg_cfg.password(pw);
103    }
104    if let Some(sslmode) = env_nonempty("PGSSLMODE") {
105        apply_sslmode(&mut pg_cfg, "PGSSLMODE", &sslmode)?;
106    }
107    Ok(pg_cfg)
108}
109
110fn env_nonempty(name: &str) -> Option<String> {
111    std::env::var(name).ok().filter(|value| !value.is_empty())
112}
113
114pub fn libpq_env_fallbacks_in_use(cfg: &SessionConfig) -> Vec<&'static str> {
115    if cfg.dsn_secret.is_some() || cfg.conninfo_secret.is_some() {
116        return Vec::new();
117    }
118    if std::env::var("AFPSQL_DSN_SECRET").is_ok() || std::env::var("AFPSQL_CONNINFO_SECRET").is_ok()
119    {
120        return Vec::new();
121    }
122    let mut used = Vec::new();
123    if cfg.host.is_none()
124        && std::env::var("AFPSQL_HOST").is_err()
125        && env_nonempty("PGHOST").is_some()
126    {
127        used.push("PGHOST");
128    }
129    if cfg.port.is_none()
130        && std::env::var("AFPSQL_PORT").is_err()
131        && env_nonempty("PGPORT").is_some()
132    {
133        used.push("PGPORT");
134    }
135    if cfg.user.is_none()
136        && std::env::var("AFPSQL_USER").is_err()
137        && env_nonempty("PGUSER").is_some()
138    {
139        used.push("PGUSER");
140    }
141    if cfg.dbname.is_none()
142        && std::env::var("AFPSQL_DBNAME").is_err()
143        && env_nonempty("PGDATABASE").is_some()
144    {
145        used.push("PGDATABASE");
146    }
147    if cfg.password_secret.is_none()
148        && std::env::var("AFPSQL_PASSWORD_SECRET").is_err()
149        && env_nonempty("PGPASSWORD").is_some()
150    {
151        used.push("PGPASSWORD");
152    }
153    if env_nonempty("PGSSLMODE").is_some() {
154        used.push("PGSSLMODE");
155    }
156    used
157}
158
159fn validate_dsn_ssl_options(dsn: &str) -> Result<(), ConnectionConfigError> {
160    let Some(query) = dsn.split_once('?').map(|(_, query)| query) else {
161        return Ok(());
162    };
163    let query = query.split('#').next().unwrap_or(query);
164    for part in query.split('&') {
165        let (key, value) = part.split_once('=').unwrap_or((part, ""));
166        validate_ssl_option(key, value, "dsn")?;
167    }
168    Ok(())
169}
170
171fn validate_conninfo_ssl_options(conninfo: &str) -> Result<(), ConnectionConfigError> {
172    for (key, value) in parse_conninfo_pairs(conninfo) {
173        validate_ssl_option(&key, &value, "conninfo")?;
174    }
175    Ok(())
176}
177
178fn validate_ssl_option(key: &str, value: &str, source: &str) -> Result<(), ConnectionConfigError> {
179    match key {
180        "sslmode" => validate_sslmode(source, value),
181        "sslnegotiation" if value == "postgres" => Ok(()),
182        "sslnegotiation" => Err(ConnectionConfigError::new(
183            format!("unsupported {source} TLS option `sslnegotiation={value}`"),
184            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()),
185        )),
186        "sslrootcert" | "sslcert" | "sslkey" | "sslpassword" | "sslcrl" | "sslcrldir"
187        | "sslcertmode" | "sslsni" | "ssl_min_protocol_version" | "ssl_max_protocol_version"
188        => Err(unsupported_ssl_option(source, key)),
189        _ => Ok(()),
190    }
191}
192
193fn apply_sslmode(
194    pg_cfg: &mut Config,
195    source: &str,
196    value: &str,
197) -> Result<(), ConnectionConfigError> {
198    validate_sslmode(source, value)?;
199    let mode = match value {
200        "disable" => SslMode::Disable,
201        "prefer" => SslMode::Prefer,
202        "require" => SslMode::Require,
203        _ => return Err(unsupported_sslmode(source, value)),
204    };
205    pg_cfg.ssl_mode(mode);
206    Ok(())
207}
208
209fn validate_sslmode(source: &str, value: &str) -> Result<(), ConnectionConfigError> {
210    match value {
211        "disable" | "prefer" | "require" => Ok(()),
212        _ => Err(unsupported_sslmode(source, value)),
213    }
214}
215
216fn unsupported_sslmode(source: &str, value: &str) -> ConnectionConfigError {
217    ConnectionConfigError::new(
218        format!(
219            "unsupported {source} sslmode `{value}`; supported values are disable, prefer, require"
220        ),
221        Some(SUPPORTED_SSLMODE_HINT.to_string()),
222    )
223}
224
225fn unsupported_ssl_option(source: &str, key: &str) -> ConnectionConfigError {
226    ConnectionConfigError::new(
227        format!("unsupported {source} TLS option `{key}`"),
228        Some(SUPPORTED_SSLMODE_HINT.to_string()),
229    )
230}
231
232fn map_pg_config_parse_error(source: &str, err: tokio_postgres::Error) -> ConnectionConfigError {
233    let cause = err.source().map(std::string::ToString::to_string);
234    if let Some(cause) = cause.as_deref() {
235        if cause == "invalid value for option `sslmode`" {
236            return ConnectionConfigError::new(
237                format!("unsupported {source} sslmode"),
238                Some(SUPPORTED_SSLMODE_HINT.to_string()),
239            );
240        }
241        if let Some(key) = cause
242            .strip_prefix("unknown option `")
243            .and_then(|rest| rest.strip_suffix('`'))
244            && is_unsupported_ssl_option(key)
245        {
246            return unsupported_ssl_option(source, key);
247        }
248    }
249
250    let detail = cause
251        .map(|cause| format!("{err}: {cause}"))
252        .unwrap_or_else(|| err.to_string());
253    ConnectionConfigError::new(format!("invalid {source}: {detail}"), None)
254}
255
256fn is_unsupported_ssl_option(key: &str) -> bool {
257    matches!(
258        key,
259        "sslrootcert"
260            | "sslcert"
261            | "sslkey"
262            | "sslpassword"
263            | "sslcrl"
264            | "sslcrldir"
265            | "sslcertmode"
266            | "sslsni"
267            | "ssl_min_protocol_version"
268            | "ssl_max_protocol_version"
269            | "sslnegotiation"
270    )
271}
272
273fn parse_conninfo_pairs(input: &str) -> Vec<(String, String)> {
274    let bytes = input.as_bytes();
275    let mut pairs = Vec::new();
276    let mut i = 0usize;
277
278    while i < bytes.len() {
279        while i < bytes.len() && bytes[i].is_ascii_whitespace() {
280            i += 1;
281        }
282        if i >= bytes.len() {
283            break;
284        }
285
286        let key_start = i;
287        while i < bytes.len() && bytes[i] != b'=' && !bytes[i].is_ascii_whitespace() {
288            i += 1;
289        }
290        let key = &input[key_start..i];
291        while i < bytes.len() && bytes[i].is_ascii_whitespace() {
292            i += 1;
293        }
294        if i >= bytes.len() || bytes[i] != b'=' {
295            break;
296        }
297        i += 1;
298        while i < bytes.len() && bytes[i].is_ascii_whitespace() {
299            i += 1;
300        }
301
302        let mut value = String::new();
303        if i < bytes.len() && bytes[i] == b'\'' {
304            i += 1;
305            while i < bytes.len() {
306                match bytes[i] {
307                    b'\\' if i + 1 < bytes.len() => {
308                        i += 1;
309                        value.push(bytes[i] as char);
310                        i += 1;
311                    }
312                    b'\'' => {
313                        i += 1;
314                        break;
315                    }
316                    b => {
317                        value.push(b as char);
318                        i += 1;
319                    }
320                }
321            }
322        } else {
323            while i < bytes.len() && !bytes[i].is_ascii_whitespace() {
324                if bytes[i] == b'\\' && i + 1 < bytes.len() {
325                    i += 1;
326                }
327                value.push(bytes[i] as char);
328                i += 1;
329            }
330        }
331
332        pairs.push((key.to_string(), value));
333    }
334
335    pairs
336}
337
338#[cfg(test)]
339#[path = "../tests/support/unit_conn.rs"]
340mod tests;
agent_first_psql/conn.rs

agent_first_psql/
conn.rs
