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 secret_config;
25pub mod skill_admin;
26pub mod ssh_transport;
27pub mod types;
28pub mod writer;
29
30use agent_first_data::OutputFormat;
31#[derive(Clone, Copy, Debug, PartialEq, Eq)]
32pub enum Capability {
33    ReadWrite,
34    ReadOnly,
35}
36
37impl Capability {
38    pub fn permits(self, permission: types::Permission) -> bool {
39        self == Self::ReadWrite || permission.is_read_only()
40    }
41}
42
43pub async fn run(capability: Capability, bin_name: &str) {
44    let raw_args = std::env::args().collect::<Vec<_>>();
45    if let Err(error) = emit::install_output_to_from_raw(&raw_args) {
46        if emit::emit_cli_error(&error, None, OutputFormat::Json).is_err() {
47            std::process::exit(4);
48        }
49        std::process::exit(2);
50    }
51    let mut locked_profile = None;
52    if capability == Capability::ReadOnly {
53        let profile_name = match raw_args
54            .first()
55            .map(String::as_str)
56            .map(readonly_policy::locked_profile_name)
57            .transpose()
58        {
59            Ok(name) => name.flatten(),
60            Err(error) => reject_readonly(&error, readonly_local_capability_hint()),
61        };
62        if let Err(error) =
63            readonly_policy::validate_raw_args_for_profile(&raw_args, profile_name.is_some())
64        {
65            reject_readonly(&error, readonly_local_capability_hint());
66        }
67        if let Some(name) = profile_name {
68            locked_profile = match readonly_policy::load_locked_profile(&name) {
69                // Pin at load time so every consumer of the profile inherits it:
70                // the administrator's endpoint must not be redirected by the
71                // environment, the way connection flags are already refused.
72                Ok(mut profile) => {
73                    profile.profile_pinned = true;
74                    Some(profile)
75                }
76                Err(error) => reject_readonly(&error, readonly_local_capability_hint()),
77            };
78        }
79    }
80    let _stream_redirect = install_stream_redirect_or_exit();
81    let mode = match cli::parse_args(bin_name) {
82        Ok(mode) => mode,
83        Err(error) => {
84            if emit::emit_cli_error(&error, None, OutputFormat::Json).is_err() {
85                std::process::exit(4);
86            }
87            std::process::exit(2);
88        }
89    };
90
91    match mode {
92        cli::Mode::Cli(request) if capability == Capability::ReadOnly && request.psql_mode => {
93            reject_readonly(
94                "psql mode is unavailable in afpsql-readonly",
95                "use `afpsql` for psql compatibility mode; it intentionally has writable semantics",
96            );
97        }
98        cli::Mode::Cli(mut request) => {
99            let has_locked_profile = locked_profile.is_some();
100            if let Some(profile) = locked_profile.clone() {
101                request.session = profile;
102            }
103            if capability == Capability::ReadOnly
104                && let Err(error) = readonly_policy::validate_session_with_trust(
105                    &request.session,
106                    has_locked_profile,
107                )
108            {
109                reject_readonly(&error, readonly_local_capability_hint());
110            }
111            cli_runner::run(request, capability, has_locked_profile).await
112        }
113        cli::Mode::Pipe(mut init) => {
114            let has_locked_profile = locked_profile.is_some();
115            if let Some(profile) = locked_profile {
116                init.session = profile;
117            }
118            if capability == Capability::ReadOnly
119                && let Err(error) =
120                    readonly_policy::validate_session_with_trust(&init.session, has_locked_profile)
121            {
122                reject_readonly(&error, readonly_local_capability_hint());
123            }
124            pipe::run(init, capability, has_locked_profile).await
125        }
126        cli::Mode::PsqlAdmin(_) if capability == Capability::ReadOnly => {
127            reject_readonly(
128                "the psql wrapper is a writable interface",
129                "use `afpsql psql status`, `afpsql psql install`, or `afpsql psql uninstall`",
130            );
131        }
132        cli::Mode::PsqlAdmin(request) => std::process::exit(psql_admin::run(request)),
133        cli::Mode::SkillAdmin(_)
134            if capability == Capability::ReadOnly && locked_profile.is_some() =>
135        {
136            reject_readonly(
137                "skill management is unavailable through an administrator-locked afpsql-readonly profile",
138                "use the ordinary afpsql-readonly or afpsql entrypoint for skill management",
139            );
140        }
141        cli::Mode::SkillAdmin(request) => std::process::exit(skill_admin::run(request)),
142        cli::Mode::PsqlUnsupported(_) if capability == Capability::ReadOnly => {
143            reject_readonly(
144                "psql mode is unavailable in afpsql-readonly",
145                "use `afpsql` for psql compatibility mode; it intentionally has writable semantics",
146            );
147        }
148        cli::Mode::PsqlUnsupported(request) => {
149            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() {
150                std::process::exit(4);
151            }
152            std::process::exit(2);
153        }
154    }
155}
156
157pub fn readonly_hint() -> &'static str {
158    "write operations require `afpsql`; use afpsql-readonly only for database reads"
159}
160
161pub fn readonly_local_capability_hint() -> &'static str {
162    "afpsql-readonly restricts PostgreSQL writes; an administrator-locked profile may additionally restrict host capabilities"
163}
164
165fn reject_readonly(error: &str, hint: &str) -> ! {
166    if emit::emit_cli_error(error, Some(hint), OutputFormat::Json).is_err() {
167        std::process::exit(4);
168    }
169    std::process::exit(2);
170}
171
172fn install_stream_redirect_or_exit()
173-> Option<agent_first_data::stream_redirect::InstalledStreamRedirect> {
174    match agent_first_data::stream_redirect::install_from_raw_args(std::env::args()) {
175        Ok(redirect) => redirect,
176        Err(error) => {
177            let value = agent_first_data::build_cli_error(&error.to_string(), None);
178            let _ = emit::emit_value(value.into(), OutputFormat::Json);
179            std::process::exit(2);
180        }
181    }
182}
183
184#[cfg(test)]
185#[path = "../tests/support/env.rs"]
186mod test_env;
187
188#[cfg(test)]
189#[path = "../tests/support/unit_main.rs"]
190mod main_tests {
191    use crate::limits::{MAX_PARAMS, MAX_SQL_BYTES};
192    use crate::logutil::build_startup_log;
193    use crate::pipe::{has_session_override, read_limited_line, validate_query_request};
194    use crate::types::{ContainerConfig, Output, SessionConfig};
195    include!("../tests/support/unit_main.rs");
196}
agent_first_psql/lib.rs

agent_first_psql/
lib.rs
