diff --git a/README.md b/README.md index 5efd554..b4f4955 100644 --- a/README.md +++ b/README.md @@ -63,7 +63,9 @@ the lease lifetime (60 seconds by default), and unique `EXECUTOR_WORKER_ID` for a singleton worker deployment to let its next instance immediately recover commands left by the previous instance at startup; otherwise a process-unique identifier is generated and abandoned -work is recovered when its lease expires. +work is recovered when its lease expires. Each poll considers up to +`COMMAND_CLAIM_CANDIDATES` sessions (10 by default) and takes the first whose +per-session claim is free, so one busy session does not delay the others. Commands are parsed with `shlex.split` before execution, so quoting rules follow POSIX shells but features like glob expansion are not performed. @@ -74,8 +76,13 @@ its PostgreSQL credentials, and the host outside `SHELL_ROOT` are trusted and must not be accessible to submitted commands. `run_subprocess` therefore builds a new environment containing only a fixed `PATH`, locale, the command account's identity variables, and values saved in `env_snapshot`. It never copies the -worker environment, and rejects `DATABASE_URL` and `PG_CONN` even if a snapshot -contains them. +worker environment, and rejects reserved variables even if a snapshot contains +them: the worker credentials `DATABASE_URL` and `PG_CONN`, every `LD_*`, +`DYLD_*`, and `BASH_FUNC_*` name, and the glibc path overrides `GCONV_PATH`, +`HOSTALIASES`, `LOCPATH`, `MALLOC_TRACE`, `NLSPATH`, and `RESOLV_HOST_CONF`. +Those loader variables would otherwise run attacker-supplied code inside an +allowlisted binary without executing anything new, defeating +`EXECUTOR_ALLOWED_COMMANDS`. Matching is case-insensitive. Production deployments must create a dedicated, unprivileged OS account and configure an absolute-path executable allowlist. The allowlist is the required @@ -106,7 +113,10 @@ command environment rather than inherited wholesale. You can run `cleanup_agent.py` periodically. Command retention applies only to terminal statuses: `done` and `failed` commands older than `CLEANUP_DAYS` are deleted, while recent terminal commands and all `pending` or `running` commands -are retained. Use `replay_agent.py` for session replays. The optional +are retained. Sessions untouched for `CLEANUP_DAYS` also have their `cwd` and +`env` reset, so give the cleanup agent the same `SHELL_ROOT` as the executor +(both default to `/home/sandbox`); a session reset to a `cwd` outside the +executor's `SHELL_ROOT` cannot run anything. Use `replay_agent.py` for session replays. The optional `monitor_agent.py` emits usage metrics like command counts and average run time to stdout or CSV. diff --git a/sql/html_fragments.sql b/sql/html_fragments.sql index f08d59f..0461c07 100644 --- a/sql/html_fragments.sql +++ b/sql/html_fragments.sql @@ -51,14 +51,14 @@ AS $$ pg_shell_html_escape(result.status), pg_shell_html_escape(result.exit_code::TEXT) ), - '' ORDER BY result.row_number + '' ORDER BY result.id ), '' )::"text/html" - FROM ( - SELECT row_number() OVER () AS row_number, output.* - FROM latest_output(p_user_id, p_session_id, p_since_id) AS output - ) AS result + -- Ordering by id keeps the fragment a chronological transcript. A bare + -- row_number() OVER () would instead depend on the order the planner + -- happened to emit latest_output's rows in. + FROM latest_output(p_user_id, p_session_id, p_since_id) AS result $$; CREATE OR REPLACE FUNCTION submit_command_html( diff --git a/sql/install.sql b/sql/install.sql index 1776fb3..cc11653 100644 --- a/sql/install.sql +++ b/sql/install.sql @@ -4,7 +4,8 @@ \i sql/migrate_replay_provenance_retention.sql \i sql/migrate_command_leases.sql \i sql/submit_command.sql -\i sql/latest_output.sql -- updated latest_output with optional since_id filter +-- latest_output takes an optional since_id filter +\i sql/latest_output.sql \i sql/html_fragments.sql \i sql/fork_session.sql \i sql/replay_session.sql diff --git a/sql/latest_output.sql b/sql/latest_output.sql index 9a6f061..0071c69 100644 --- a/sql/latest_output.sql +++ b/sql/latest_output.sql @@ -18,18 +18,30 @@ RETURNS TABLE( completed_at TIMESTAMP ) LANGUAGE plpgsql AS $$ BEGIN + -- Rows are always returned oldest-first so a caller can render them as a + -- transcript. The p_since_id = 0 backfill still selects the *newest* 20 + -- commands (SPEC.md: "last N commands"), but the DESC + LIMIT is confined + -- to an inner query and the result is re-sorted ascending. Returning that + -- page newest-first would make the initial page load read bottom-to-top + -- and then flip as soon as incremental polling took over. + -- -- Columns are qualified with the table alias so they are not confused -- with the identically named RETURNS TABLE output columns, which would -- otherwise raise "column reference \"id\" is ambiguous". IF p_since_id = 0 THEN RETURN QUERY - SELECT c.id, c.command, c.output, c.exit_code, c.status, - c.submitted_at, c.completed_at - FROM commands c - WHERE c.user_id = p_user_id AND c.session_id = p_session_id - AND c.id > p_since_id - ORDER BY c.id DESC - LIMIT 20; + SELECT recent.id, recent.command, recent.output, recent.exit_code, + recent.status, recent.submitted_at, recent.completed_at + FROM ( + SELECT c.id, c.command, c.output, c.exit_code, c.status, + c.submitted_at, c.completed_at + FROM commands c + WHERE c.user_id = p_user_id AND c.session_id = p_session_id + AND c.id > p_since_id + ORDER BY c.id DESC + LIMIT 20 + ) AS recent + ORDER BY recent.id ASC; ELSE RETURN QUERY SELECT c.id, c.command, c.output, c.exit_code, c.status, diff --git a/tests/test_cleanup_agent.py b/tests/test_cleanup_agent.py index 0f6d3fd..f2e0265 100644 --- a/tests/test_cleanup_agent.py +++ b/tests/test_cleanup_agent.py @@ -107,6 +107,83 @@ def conn(): conn.close() +def test_main_retries_after_a_connection_failure(monkeypatch, caplog): + """A database outage must not silently end the retention daemon.""" + attempts = [] + + def flaky_get_conn(): + attempts.append(len(attempts)) + if len(attempts) < 3: + raise RuntimeError("connection refused") + raise KeyboardInterrupt # stand in for operator shutdown + + sleeps = [] + monkeypatch.setattr(cleanup_agent, "get_conn", flaky_get_conn) + monkeypatch.setattr(cleanup_agent.time, "sleep", sleeps.append) + monkeypatch.setattr(sys, "argv", ["cleanup_agent.py", "--interval", "3600"]) + + with pytest.raises(KeyboardInterrupt): + cleanup_agent.main() + + assert len(attempts) == 3, "daemon gave up instead of retrying" + # Retries back off by RETRY_DELAY_SECONDS, not the full 3600s interval. + assert sleeps == [cleanup_agent.RETRY_DELAY_SECONDS] * 2 + + +def test_main_retries_after_a_failed_cleanup_run(monkeypatch): + """An error mid-run must not end the daemon either.""" + closed = [] + + class FakeConn: + def close(self): + closed.append(True) + + runs = [] + + def flaky_cleanup_once(conn, days): + runs.append(days) + if len(runs) < 2: + raise psycopg2.OperationalError("server closed the connection") + raise KeyboardInterrupt + + monkeypatch.setattr(cleanup_agent, "get_conn", FakeConn) + monkeypatch.setattr(cleanup_agent, "cleanup_once", flaky_cleanup_once) + monkeypatch.setattr(cleanup_agent.time, "sleep", lambda seconds: None) + monkeypatch.setattr(sys, "argv", ["cleanup_agent.py"]) + + with pytest.raises(KeyboardInterrupt): + cleanup_agent.main() + + assert len(runs) == 2 + assert len(closed) == 2, "the connection must be closed on every path" + + +@pytest.mark.parametrize( + "failure", [RuntimeError("no dsn"), psycopg2.OperationalError("boom")] +) +def test_main_once_reports_failure_through_exit_status(monkeypatch, failure): + """--once must exit non-zero so a supervisor or cron notices.""" + if isinstance(failure, RuntimeError): + monkeypatch.setattr( + cleanup_agent, "get_conn", lambda: (_ for _ in ()).throw(failure) + ) + else: + + class FakeConn: + def close(self): + pass + + monkeypatch.setattr(cleanup_agent, "get_conn", FakeConn) + monkeypatch.setattr( + cleanup_agent, + "cleanup_once", + lambda conn, days: (_ for _ in ()).throw(failure), + ) + monkeypatch.setattr(sys, "argv", ["cleanup_agent.py", "--once"]) + + assert cleanup_agent.main() == 1 + + def test_cleanup_once_deletes_only_old_terminal_commands(conn): uid_str = str(uuid.uuid4()) with conn.cursor() as cur: @@ -266,7 +343,37 @@ def test_cleanup_once_resets_env(conn): cur.execute("SELECT cwd, env FROM environments WHERE user_id = %s", (uid_str,)) cwd, env = cur.fetchone() - assert cwd == "/home/sandbox" + assert cwd == cleanup_agent.DEFAULT_SHELL_ROOT + assert env == {} + + +def test_cleanup_once_resets_env_to_configured_shell_root(conn, monkeypatch): + """A reset session must land inside the executor's SHELL_ROOT. + + The executor hands environments.cwd straight to the subprocess, so + resetting to a hardcoded /home/sandbox in a deployment that configured a + different SHELL_ROOT leaves the session unable to run anything. + """ + monkeypatch.setenv("SHELL_ROOT", "/home/pg-shell-command") + uid_str = str(uuid.uuid4()) + with conn.cursor() as cur: + cur.execute( + "INSERT INTO users (id, username) VALUES (%s, %s)", + (uid_str, "shell-root-user"), + ) + cur.execute( + "INSERT INTO environments (user_id, cwd, env, updated_at) " + "VALUES (%s, %s, %s::jsonb, now() - interval '100 days')", + (uid_str, "/home/pg-shell-command/deep", '{"k":1}'), + ) + + cleanup_once(conn, 90) + + with conn.cursor() as cur: + cur.execute("SELECT cwd, env FROM environments WHERE user_id = %s", (uid_str,)) + cwd, env = cur.fetchone() + + assert cwd == "/home/pg-shell-command" assert env == {} diff --git a/tests/test_executor_agent.py b/tests/test_executor_agent.py index fe930f5..4bc2955 100644 --- a/tests/test_executor_agent.py +++ b/tests/test_executor_agent.py @@ -123,7 +123,29 @@ def test_run_subprocess_does_not_expose_worker_secrets(monkeypatch, tmp_path): assert output.splitlines() == ["None", "None"] -@pytest.mark.parametrize("reserved", ["DATABASE_URL", "PG_CONN"]) +@pytest.mark.parametrize( + "reserved", + [ + "DATABASE_URL", + "PG_CONN", + # Dynamic loader controls: these run attacker-supplied code inside an + # allowlisted binary, so they would defeat EXECUTOR_ALLOWED_COMMANDS. + "LD_PRELOAD", + "LD_LIBRARY_PATH", + "LD_AUDIT", + "DYLD_INSERT_LIBRARIES", + "BASH_FUNC_ls%%", + # glibc path overrides with the same effect. + "GCONV_PATH", + "LOCPATH", + "NLSPATH", + "HOSTALIASES", + "RESOLV_HOST_CONF", + "MALLOC_TRACE", + # Names are matched case-insensitively; the loader does the same. + "ld_preload", + ], +) def test_run_subprocess_rejects_reserved_snapshot_variables( monkeypatch, tmp_path, reserved ): @@ -131,6 +153,31 @@ def test_run_subprocess_rejects_reserved_snapshot_variables( run_subprocess("python3 -c 'pass'", str(tmp_path), {reserved: "secret"}) +def test_run_subprocess_rejects_ld_preload_before_launching(monkeypatch, tmp_path): + """LD_PRELOAD must be refused rather than handed to the dynamic loader.""" + + def fail_popen(*args, **kwargs): # pragma: no cover - must not be reached + pytest.fail("subprocess launched with an attacker-controlled loader") + + monkeypatch.setattr(workers.executor_agent.subprocess, "Popen", fail_popen) + + with pytest.raises(ValueError, match="LD_PRELOAD"): + run_subprocess( + "python3 -c 'pass'", + str(tmp_path), + {"LD_PRELOAD": str(tmp_path / "evil.so")}, + ) + + +def test_run_subprocess_allows_unreserved_snapshot_variables(tmp_path): + """The prefix rule must not swallow ordinary variables.""" + cmd = "python3 -c 'import os;print(os.environ[\"LDAP_URL\"])'" + exit_code, output = run_subprocess(cmd, str(tmp_path), {"LDAP_URL": "ldap://x"}) + + assert exit_code == 0 + assert "ldap://x" in output + + def test_handle_command_uses_combined_output(monkeypatch): captured = {} diff --git a/tests/test_executor_integration.py b/tests/test_executor_integration.py index deb51b0..8926ee6 100644 --- a/tests/test_executor_integration.py +++ b/tests/test_executor_integration.py @@ -12,10 +12,16 @@ import uuid import psycopg2 +import psycopg2.extensions import pytest import workers.executor_agent -from workers.executor_agent import fetch_pending, handle_command, recover_worker_commands +from workers.executor_agent import ( + _release_user_claim, + fetch_pending, + handle_command, + recover_worker_commands, +) @pytest.fixture(autouse=True) @@ -111,6 +117,105 @@ def test_executor_cd_updates_environment(db_conn, monkeypatch, tmp_path): assert cur.fetchone()[0] == str(sub) +def test_fetch_pending_skips_a_session_whose_claim_lock_is_held(db_conn): + """A busy session must not hide every session queued behind it. + + SPEC.md section 3 requires per-user serialization "while still allowing + different users to execute concurrently". fetch_pending used to look at a + single candidate row and return None when its advisory lock was taken, so + one session whose lock was held by a peer worker (for instance after its + lease expired but before it finished) starved every other session for as + long as that worker held the lock. + """ + busy_user, busy_session = _create_user_with_env(db_conn, "/tmp") + free_user, free_session = _create_user_with_env(db_conn, "/tmp") + + with db_conn.cursor() as cur: + # The blocked session's command is submitted first, so it sorts ahead. + cur.execute( + "SELECT submit_command(%s, %s, %s)", (busy_user, busy_session, "echo busy") + ) + busy_id = cur.fetchone()[0] + cur.execute( + "SELECT submit_command(%s, %s, %s)", (free_user, free_session, "echo free") + ) + free_id = cur.fetchone()[0] + cur.execute("SELECT hashtextextended(%s, 0)", (busy_session,)) + busy_lock_key = cur.fetchone()[0] + + # A peer worker on its own connection holds the busy session's claim lock. + peer = psycopg2.connect(os.environ["TEST_DATABASE_URL"]) + peer.autocommit = True + try: + with peer.cursor() as cur: + cur.execute("SELECT pg_advisory_lock(%s)", (busy_lock_key,)) + + db_conn.autocommit = False + row = None + try: + row = fetch_pending(db_conn) + assert row is not None, "poll gave up instead of trying the next session" + assert row["id"] == free_id + assert str(row["session_id"]) == free_session + finally: + if row is not None: + _release_user_claim(db_conn, row) + db_conn.autocommit = True + finally: + peer.close() + + with db_conn.cursor() as cur: + cur.execute( + "SELECT id, status FROM commands WHERE id IN (%s, %s) ORDER BY id", + (busy_id, free_id), + ) + statuses = dict(cur.fetchall()) + + # The blocked session is left untouched for the worker that owns it. + assert statuses[busy_id] == "pending" + assert statuses[free_id] == "running" + + +def test_fetch_pending_claims_atomically_without_server_warnings(db_conn): + """The claim is one transaction, and it does not warn on every poll. + + fetch_pending used to open its transaction with an explicit ``BEGIN``, + which psycopg2 had already issued implicitly. PostgreSQL answered every + single poll with ``WARNING: there is already a transaction in progress``, + so a worker polling once a second wrote 86,400 warnings a day into the + server log. + """ + user_id, session_id = _create_user_with_env(db_conn, "/tmp") + with db_conn.cursor() as cur: + cur.execute( + "SELECT submit_command(%s, %s, %s)", (user_id, session_id, "echo warn") + ) + cmd_id = cur.fetchone()[0] + + db_conn.autocommit = False + del db_conn.notices[:] + try: + row = fetch_pending(db_conn) + assert row is not None and row["id"] == cmd_id + # Still a single transaction: the claim is not visible until commit, + # and fetch_pending commits before returning. + assert db_conn.get_transaction_status() == ( + psycopg2.extensions.TRANSACTION_STATUS_IDLE + ) + notices = list(db_conn.notices) + finally: + _release_user_claim(db_conn, row) + db_conn.autocommit = True + + assert not [n for n in notices if "already a transaction in progress" in n], notices + + with db_conn.cursor() as cur: + cur.execute("SELECT status, worker_id FROM commands WHERE id = %s", (cmd_id,)) + status, worker_id = cur.fetchone() + assert status == "running" + assert worker_id is not None + + def test_fetch_pending_returns_none_when_idle(db_conn): # No pending rows -> fetch_pending yields nothing rather than raising. db_conn.autocommit = False diff --git a/tests/test_functions.py b/tests/test_functions.py index 898b0a8..0242c43 100644 --- a/tests/test_functions.py +++ b/tests/test_functions.py @@ -1,5 +1,6 @@ import json import os +import re import select import time import uuid @@ -107,11 +108,14 @@ def test_submit_and_latest_output(conn): cur.execute("SELECT * FROM latest_output(%s, %s)", (user_id, session_id)) rows = cur.fetchall() + # The backfill keeps the newest 20 commands but hands them back + # oldest-first, so a caller can render them as a transcript and then + # append later rows from the p_since_id branch without reversing. assert len(rows) == 20 - expected_ids = list(reversed(cmd_ids))[:20] + expected_ids = cmd_ids[-20:] assert [row[0] for row in rows] == expected_ids - assert rows[0][2] == "msg24" - assert rows[0][6] is not None + assert rows[-1][2] == "msg24" + assert rows[-1][6] is not None def test_submit_command_notifies(conn): @@ -275,6 +279,59 @@ def test_latest_output_since_id(conn): assert [row[0] for row in rows] == expected_ids +def test_latest_output_ordering_is_consistent_across_branches(conn): + """The backfill and the incremental branch must agree on row order. + + The htmx UI replaces #output with the p_since_id = 0 fragment and appends + submissions with hx-swap="beforeend"; the CLI's ``tail`` prints rows in + the order it receives them. Either caller renders a scrambled transcript + if the two branches disagree. + """ + user_id = str(uuid.uuid4()) + with conn.cursor() as cur: + cur.execute("INSERT INTO users(id, username) VALUES (%s, %s)", (user_id, "u4")) + session_id = _create_session(cur, user_id) + ids = [] + for idx in range(5): + cur.execute( + "SELECT submit_command(%s, %s, %s)", + (user_id, session_id, f"echo ord{idx}"), + ) + ids.append(cur.fetchone()[0]) + + cur.execute("SELECT id FROM latest_output(%s, %s, 0)", (user_id, session_id)) + backfill = [row[0] for row in cur.fetchall()] + cur.execute( + "SELECT id FROM latest_output(%s, %s, %s)", (user_id, session_id, ids[1]) + ) + incremental = [row[0] for row in cur.fetchall()] + + assert backfill == sorted(backfill), "backfill must be oldest-first" + assert incremental == sorted(incremental), "incremental must be oldest-first" + assert backfill == ids + assert incremental == ids[2:] + + +def test_latest_output_html_renders_oldest_first(conn): + user_id = str(uuid.uuid4()) + with conn.cursor() as cur: + cur.execute("INSERT INTO users(id, username) VALUES (%s, %s)", (user_id, "u5")) + session_id = _create_session(cur, user_id) + ids = [] + for idx in range(3): + cur.execute( + "SELECT submit_command(%s, %s, %s)", + (user_id, session_id, f"echo html{idx}"), + ) + ids.append(cur.fetchone()[0]) + + cur.execute("SELECT latest_output_html(%s, %s, 0)", (user_id, session_id)) + fragment = cur.fetchone()[0] + + rendered = [int(value) for value in re.findall(r'data-command-id="(\d+)"', fragment)] + assert rendered == ids + + def test_html_rpcs_render_escaped_fragments(conn): user_id = str(uuid.uuid4()) command = '' diff --git a/tests/test_monitor_agent.py b/tests/test_monitor_agent.py index 12525dd..92bfa27 100644 --- a/tests/test_monitor_agent.py +++ b/tests/test_monitor_agent.py @@ -1,7 +1,10 @@ +import os from datetime import datetime import csv +import psycopg2 + import workers.monitor_agent as monitor_agent @@ -173,3 +176,56 @@ class Args: pass else: # pragma: no cover - defensive raise AssertionError("Expected ValueError when both windows are set") + + +def test_ensure_monitor_state_table_commits(monkeypatch): + """The CREATE TABLE must be committed, not left to a later save. + + Runs that never reach save_monitor_state -- every --since-hours or + --since-days run, and any run with nothing new -- otherwise rolled the + table back on close, so the watermark could never be kept. + """ + calls = [] + + class FakeCursor: + def __enter__(self): + return self + + def __exit__(self, *exc): + return False + + def execute(self, query, params=None): + calls.append("execute") + + class FakeConn: + def cursor(self): + return FakeCursor() + + def commit(self): + calls.append("commit") + + monitor_agent.ensure_monitor_state_table(FakeConn()) + + assert calls == ["execute", "commit"] + + +def test_ensure_monitor_state_table_persists_without_a_watermark(db_conn): + """End to end: the table survives a run that saves no watermark.""" + with db_conn.cursor() as cur: + cur.execute("DROP TABLE IF EXISTS monitor_state") + + conn = psycopg2.connect(os.environ["TEST_DATABASE_URL"]) + try: + monitor_agent.ensure_monitor_state_table(conn) + finally: + conn.close() # rolls back anything left uncommitted + + with db_conn.cursor() as cur: + cur.execute("SELECT to_regclass('monitor_state') IS NOT NULL") + assert cur.fetchone()[0] is True + + +def test_module_exposes_its_docstring(): + """A __future__ import above the docstring demotes it to a dead literal.""" + assert monitor_agent.__doc__ is not None + assert "pg_shell monitor agent" in monitor_agent.__doc__ diff --git a/workers/cleanup_agent.py b/workers/cleanup_agent.py index 58f09f5..4ea1955 100755 --- a/workers/cleanup_agent.py +++ b/workers/cleanup_agent.py @@ -6,6 +6,19 @@ from workers.db import get_conn +RETRY_DELAY_SECONDS = 5 +DEFAULT_SHELL_ROOT = "/home/sandbox" + + +def shell_root() -> str: + """Return the sandbox root a reset session should start from. + + Must agree with executor_agent's ``SHELL_ROOT``: a reset session whose cwd + sits outside that root cannot run anything, because the executor passes + the cwd straight to the subprocess. + """ + return os.getenv("SHELL_ROOT", DEFAULT_SHELL_ROOT) + def cleanup_once(conn, days: int) -> None: """Remove commands and reset environments older than ``days`` days. @@ -30,17 +43,17 @@ def cleanup_once(conn, days: int) -> None: cur.execute( """ UPDATE environments - SET cwd = '/home/sandbox', env = '{}'::jsonb, updated_at = now() + SET cwd = %s, env = '{}'::jsonb, updated_at = now() WHERE updated_at < now() - %s * interval '1 day' """, - (days,), + (shell_root(), days), ) reset = cur.rowcount logging.info("Reset %d stale environments", reset) conn.commit() -def main(): +def main() -> int: parser = argparse.ArgumentParser( description="Cleanup old commands and environments" ) @@ -74,19 +87,32 @@ def main(): logging.basicConfig(level=logging.INFO, format="%(asctime)s %(message)s") while True: + # A database outage must not end the daemon: it is the periodic + # retention job, so giving up means retention silently stops until + # somebody notices. --once still reports the failure to its caller. try: conn = get_conn() except RuntimeError as exc: logging.error("Cleanup agent failed to connect to database: %s", exc) - break + if args.once: + return 1 + time.sleep(min(args.interval, RETRY_DELAY_SECONDS)) + continue + try: cleanup_once(conn, args.days) + except Exception: + logging.exception("Cleanup run failed") + failed = True + else: + failed = False finally: conn.close() + if args.once: - break - time.sleep(args.interval) + return 1 if failed else 0 + time.sleep(min(args.interval, RETRY_DELAY_SECONDS) if failed else args.interval) if __name__ == "__main__": - main() + raise SystemExit(main()) diff --git a/workers/executor_agent.py b/workers/executor_agent.py index 039b37f..c7a913d 100644 --- a/workers/executor_agent.py +++ b/workers/executor_agent.py @@ -38,8 +38,30 @@ TRUNCATION_SUFFIX = "...[truncated]" TERMINATION_GRACE_SECONDS = 0.5 PIPE_DRAIN_TIMEOUT_SECONDS = 0.5 -RESERVED_COMMAND_ENV = frozenset({"DATABASE_URL", "PG_CONN"}) +RESERVED_COMMAND_ENV = frozenset( + { + # Worker credentials. + "DATABASE_URL", + "PG_CONN", + # glibc path overrides that make an already-allowlisted binary load or + # read attacker-controlled files without executing anything new. + "GCONV_PATH", + "HOSTALIASES", + "LOCPATH", + "MALLOC_TRACE", + "NLSPATH", + "RESOLV_HOST_CONF", + } +) +# Dynamic loader controls. ``LD_PRELOAD`` and friends run attacker-supplied +# code inside an allowlisted binary, so honouring them from a command's +# ``env_snapshot`` would defeat EXECUTOR_ALLOWED_COMMANDS entirely. +# ``BASH_FUNC_*`` smuggles exported shell functions into any allowlisted shell. +RESERVED_COMMAND_ENV_PREFIXES = ("LD_", "DYLD_", "BASH_FUNC_") DEFAULT_COMMAND_PATH = "/usr/local/bin:/usr/bin:/bin" +# How many sessions a single poll may try before giving up. One session +# whose claim lock is held elsewhere must not hide every other session. +CLAIM_CANDIDATES = int(os.getenv("COMMAND_CLAIM_CANDIDATES", "10")) LEASE_SECONDS = float(os.getenv("COMMAND_LEASE_SECONDS", "60")) LEASE_REFRESH_SECONDS = float( os.getenv("COMMAND_LEASE_REFRESH_SECONDS", str(max(1.0, LEASE_SECONDS / 3))) @@ -144,9 +166,17 @@ def fetch_pending(conn, worker_id: str = WORKER_ID) -> Dict[str, Any] | None: run a command for the same user concurrently. Snapshots are refreshed from the user's effective environment at claim time, after all preceding commands have reached a terminal state. + + Up to ``CLAIM_CANDIDATES`` sessions are considered per poll and the first + one whose claim lock is free is taken. Sessions are independent, so a + session locked by another worker must be stepped over rather than ending + the poll -- otherwise it would hide every session behind it. + + ``conn`` must not be in autocommit mode: the requeue, the claim and the + row lock all have to land in one transaction, which psycopg2 opens + implicitly on the first statement below and this function commits. """ with conn.cursor(cursor_factory=RealDictCursor) as cur: - cur.execute("BEGIN;") # Requeue abandoned work first so it no longer blocks later commands # in the same session. NULL leases are rows created by older workers. cur.execute( @@ -175,14 +205,19 @@ def fetch_pending(conn, worker_id: str = WORKER_ID) -> Dict[str, Any] | None: ) ORDER BY c.submitted_at, c.id FOR UPDATE SKIP LOCKED - LIMIT 1 - """) - row = cur.fetchone() + LIMIT %s + """, (CLAIM_CANDIDATES,)) + # The NOT EXISTS above admits at most one command per session, so + # these candidates are distinct sessions in submission order. + row = None + for candidate in cur.fetchall(): + cur.execute( + "SELECT pg_try_advisory_lock(%s)", (candidate["claim_lock_key"],) + ) + if cur.fetchone()["pg_try_advisory_lock"]: + row = candidate + break if row: - cur.execute("SELECT pg_try_advisory_lock(%s)", (row["claim_lock_key"],)) - if not cur.fetchone()["pg_try_advisory_lock"]: - conn.commit() - return None cur.execute( """ UPDATE commands @@ -237,6 +272,19 @@ def refresh_lease(conn, cmd_id: int, worker_id: str = WORKER_ID) -> bool: return refreshed +def _is_reserved_command_env(name: str) -> bool: + """Report whether ``name`` must never be taken from ``env_snapshot``. + + Snapshots come from untrusted submitters, so they may not carry worker + credentials or any variable that changes what an allowlisted executable + loads or runs. + """ + upper = name.upper() + return upper in RESERVED_COMMAND_ENV or upper.startswith( + RESERVED_COMMAND_ENV_PREFIXES + ) + + def run_subprocess( command: str, cwd: str, @@ -257,7 +305,7 @@ def run_subprocess( snapshot = env_snapshot if not isinstance(snapshot, dict): raise ValueError("env_snapshot must be a JSON object") - reserved = RESERVED_COMMAND_ENV.intersection(snapshot) + reserved = {name for name in snapshot if _is_reserved_command_env(name)} if reserved: names = ", ".join(sorted(reserved)) raise ValueError(f"reserved environment variable(s): {names}") diff --git a/workers/monitor_agent.py b/workers/monitor_agent.py index 682f62b..6b62cf0 100644 --- a/workers/monitor_agent.py +++ b/workers/monitor_agent.py @@ -1,6 +1,4 @@ #!/usr/bin/env python3 -from __future__ import annotations - """pg_shell monitor agent. Collects usage statistics from the ``commands`` table. Metrics include @@ -9,6 +7,8 @@ written to a CSV file. """ +from __future__ import annotations + import argparse import csv import logging @@ -25,6 +25,13 @@ def ensure_monitor_state_table(conn) -> None: + """Create the watermark table, committing so it outlives this run. + + Without the commit the CREATE TABLE is rolled back whenever the run does + not reach save_monitor_state -- always under --since-hours/--since-days, + and on any run with nothing new to report -- so the agent recreated the + table on every poll and never kept a watermark. + """ with conn.cursor() as cur: cur.execute( """ @@ -36,6 +43,7 @@ def ensure_monitor_state_table(conn) -> None: ) """ ) + conn.commit() def load_monitor_state(conn) -> tuple[datetime | None, int]: @@ -216,7 +224,7 @@ def upsert_csv_metrics(path: str, rows: Iterator[Row]) -> None: os.replace(temporary_path, path) -def main() -> None: +def main() -> int: parser = argparse.ArgumentParser(description="Collect usage metrics") parser.add_argument( "--interval",