agent_first_psql/
lib.rs

1#![deny(
2    clippy::unwrap_used,
3    clippy::expect_used,
4    clippy::panic,
5    clippy::print_stdout,
6    clippy::print_stderr
7)]
8
9pub mod cli;
10pub mod cli_runner;
11pub mod config;
12pub mod conn;
13pub mod container_transport;
14pub mod db;
15pub mod emit;
16pub mod handler;
17pub mod limits;
18pub mod logutil;
19pub mod output_fmt;
20pub mod pipe;
21pub mod protocol;
22pub mod psql_admin;
23pub mod readonly_policy;
24pub mod runtime_env;
25pub mod secret_config;
26pub mod skill_admin;
27pub mod ssh_transport;
28pub mod types;
29pub mod writer;
30
31use agent_first_data::OutputFormat;
32#[derive(Clone, Copy, Debug, PartialEq, Eq)]
33pub enum Capability {
34    ReadWrite,
35    ReadOnly,
36}
37
38pub async fn run(capability: Capability, bin_name: &str) {
39    let raw_args = std::env::args().collect::<Vec<_>>();
40    let mut locked_profile = None;
41    if capability == Capability::ReadOnly {
42        let profile_name = match raw_args
43            .first()
44            .map(String::as_str)
45            .map(readonly_policy::locked_profile_name)
46            .transpose()
47        {
48            Ok(name) => name.flatten(),
49            Err(error) => reject_readonly(&error, readonly_local_capability_hint()),
50        };
51        if let Err(error) =
52            readonly_policy::validate_raw_args_for_profile(&raw_args, profile_name.is_some())
53        {
54            reject_readonly(&error, readonly_local_capability_hint());
55        }
56        if let Some(name) = profile_name {
57            locked_profile = match readonly_policy::load_locked_profile(&name) {
58                // Pin at load time so every consumer of the profile inherits it:
59                // the administrator's endpoint must not be redirected by the
60                // environment, the way connection flags are already refused.
61                Ok(mut profile) => {
62                    profile.profile_pinned = true;
63                    Some(profile)
64                }
65                Err(error) => reject_readonly(&error, readonly_local_capability_hint()),
66            };
67        }
68    }
69    // The registry decides the destination and the file sinks, so redirection
70    // is installed from the resolved plan rather than from a second scan of
71    // argv — and only after the readonly capability checks above have had
72    // their say about creating local files.
73    let cli::Parsed {
74        mode,
75        redirect: _redirect,
76    } = match cli::parse_args(bin_name) {
77        Ok(parsed) => parsed,
78        Err(error) => {
79            if emit::emit_coded_error(
80                &error.code,
81                &error.message,
82                error.hint.as_deref(),
83                OutputFormat::Json,
84            )
85            .is_err()
86            {
87                std::process::exit(4);
88            }
89            std::process::exit(2);
90        }
91    };
92
93    match mode {
94        cli::Mode::Cli(request) if capability == Capability::ReadOnly && request.psql_mode => {
95            reject_readonly(
96                "psql mode is unavailable in afpsql-readonly",
97                "use `afpsql` for psql compatibility mode; it intentionally has writable semantics",
98            );
99        }
100        cli::Mode::Cli(mut request) => {
101            let has_locked_profile = locked_profile.is_some();
102            if let Some(profile) = locked_profile.clone() {
103                request.session = profile;
104            }
105            if capability == Capability::ReadOnly
106                && let Err(error) = readonly_policy::validate_session_with_trust(
107                    &request.session,
108                    has_locked_profile,
109                )
110            {
111                reject_readonly(&error, readonly_local_capability_hint());
112            }
113            cli_runner::run(request, capability, has_locked_profile).await
114        }
115        cli::Mode::Pipe(mut init) => {
116            let has_locked_profile = locked_profile.is_some();
117            if let Some(profile) = locked_profile {
118                init.session = profile;
119            }
120            if capability == Capability::ReadOnly
121                && let Err(error) =
122                    readonly_policy::validate_session_with_trust(&init.session, has_locked_profile)
123            {
124                reject_readonly(&error, readonly_local_capability_hint());
125            }
126            pipe::run(init, capability, has_locked_profile).await
127        }
128        cli::Mode::PsqlAdmin(_) if capability == Capability::ReadOnly => {
129            reject_readonly(
130                "the psql wrapper is a writable interface",
131                "use `afpsql psql status`, `afpsql psql install`, or `afpsql psql uninstall`",
132            );
133        }
134        cli::Mode::PsqlAdmin(request) => std::process::exit(psql_admin::run(request)),
135        cli::Mode::SkillAdmin(_)
136            if capability == Capability::ReadOnly && locked_profile.is_some() =>
137        {
138            reject_readonly(
139                "skill management is unavailable through an administrator-locked afpsql-readonly profile",
140                "use the ordinary afpsql-readonly or afpsql entrypoint for skill management",
141            );
142        }
143        cli::Mode::SkillAdmin(request) => std::process::exit(skill_admin::run(request)),
144        cli::Mode::PsqlUnsupported(_) if capability == Capability::ReadOnly => {
145            reject_readonly(
146                "psql mode is unavailable in afpsql-readonly",
147                "use `afpsql` for psql compatibility mode; it intentionally has writable semantics",
148            );
149        }
150        cli::Mode::PsqlUnsupported(request) => {
151            if emit::emit_cli_error(&format!("unsupported psql mode: {}", request.reason), Some("run the original psql binary directly, for example /path/to/postgresql/bin/psql, or put that PostgreSQL bin directory before the afpsql wrapper in PATH"), OutputFormat::Json).is_err() {
152                std::process::exit(4);
153            }
154            std::process::exit(2);
155        }
156    }
157}
158
159pub fn readonly_hint() -> &'static str {
160    "write operations require `afpsql`; use afpsql-readonly only for database reads"
161}
162
163pub fn readonly_local_capability_hint() -> &'static str {
164    "afpsql-readonly restricts PostgreSQL writes; an administrator-locked profile may additionally restrict host capabilities"
165}
166
167fn reject_readonly(error: &str, hint: &str) -> ! {
168    if emit::emit_cli_error(error, Some(hint), OutputFormat::Json).is_err() {
169        std::process::exit(4);
170    }
171    std::process::exit(2);
172}
173
174#[cfg(test)]
175#[path = "../tests/support/env.rs"]
176mod test_env;
177
178#[cfg(test)]
179#[path = "../tests/support/unit_main.rs"]
180mod main_tests {
181    use crate::limits::{MAX_PARAMS, MAX_SQL_BYTES};
182    use crate::logutil::build_startup_log;
183    use crate::pipe::{has_session_override, read_limited_line, validate_query_request};
184    use crate::types::{ContainerConfig, Output, SessionConfig};
185    include!("../tests/support/unit_main.rs");
186}
agent_first_psql/lib.rs

agent_first_psql/
lib.rs
