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