use crate::config::HookConfig;
use anyhow::{bail, Result};
use serde_json::Value;
use std::process::Stdio;
use std::time::Duration;
use tokio::io::AsyncWriteExt;
const DEFAULT_TIMEOUT_SECS: u64 = 10;
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum HookVerdict {
Allow,
Deny(String),
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum Event {
PreTool,
PostTool,
SessionEnd,
}
#[derive(Debug)]
struct Hook {
event: Event,
command: String,
tools: Vec<String>,
timeout: Duration,
}
impl Hook {
fn matches_tool(&self, tool: &str) -> bool {
self.tools.is_empty() || self.tools.iter().any(|t| t == tool)
}
}
#[derive(Debug, Default)]
pub struct HookSet {
hooks: Vec<Hook>,
}
impl HookSet {
pub fn from_config(configs: &[HookConfig]) -> Result<Self> {
let mut hooks = Vec::new();
for c in configs {
let event = match c.event.as_str() {
"pre_tool" => Event::PreTool,
"post_tool" => Event::PostTool,
"session_end" => Event::SessionEnd,
other => {
bail!("hook event {other:?} is not one of pre_tool, post_tool, session_end")
}
};
if c.command.trim().is_empty() {
bail!("a hook for {:?} has an empty command", c.event);
}
hooks.push(Hook {
event,
command: c.command.clone(),
tools: c.tools.clone(),
timeout: Duration::from_secs(c.timeout_secs.unwrap_or(DEFAULT_TIMEOUT_SECS)),
});
}
Ok(HookSet { hooks })
}
pub fn is_empty(&self) -> bool {
self.hooks.is_empty()
}
pub fn watches_tools(&self) -> bool {
self.hooks
.iter()
.any(|h| matches!(h.event, Event::PreTool | Event::PostTool))
}
async fn run_one(
hook: &Hook,
payload: &Value,
workdir: &std::path::Path,
) -> Result<(i32, String)> {
let mut child = tokio::process::Command::new("sh")
.arg("-c")
.arg(&hook.command)
.current_dir(workdir)
.stdin(Stdio::piped())
.stdout(Stdio::piped())
.stderr(Stdio::piped())
.kill_on_drop(true)
.spawn()?;
let bytes = serde_json::to_vec(payload)?;
let fut = async move {
if let Some(mut stdin) = child.stdin.take() {
let _ = stdin.write_all(&bytes).await;
drop(stdin);
}
child.wait_with_output().await
};
let out = tokio::time::timeout(hook.timeout, fut).await??;
let mut text = String::from_utf8_lossy(&out.stdout).trim().to_string();
if text.is_empty() {
text = String::from_utf8_lossy(&out.stderr).trim().to_string();
}
Ok((out.status.code().unwrap_or(-1), text))
}
pub async fn pre_tool(
&self,
tool: &str,
input: &Value,
workdir: &std::path::Path,
) -> HookVerdict {
for hook in self.hooks.iter().filter(|h| h.event == Event::PreTool) {
if !hook.matches_tool(tool) {
continue;
}
let payload = serde_json::json!({
"event": "pre_tool",
"tool": tool,
"input": input,
});
match Self::run_one(hook, &payload, workdir).await {
Ok((0, _)) => {}
Ok((2, reason)) => {
return HookVerdict::Deny(if reason.is_empty() {
format!("blocked by hook `{}`", hook.command)
} else {
reason
});
}
Ok((code, reason)) => {
return HookVerdict::Deny(format!(
"hook `{}` exited {code} (exit 0 allows, 2 denies){}",
hook.command,
if reason.is_empty() {
String::new()
} else {
format!(": {reason}")
}
));
}
Err(e) => {
return HookVerdict::Deny(format!(
"hook `{}` failed to run: {e}",
hook.command
));
}
}
}
HookVerdict::Allow
}
pub async fn post_tool(
&self,
tool: &str,
input: &Value,
is_error: bool,
content: &str,
workdir: &std::path::Path,
) {
for hook in self.hooks.iter().filter(|h| h.event == Event::PostTool) {
if !hook.matches_tool(tool) {
continue;
}
let payload = serde_json::json!({
"event": "post_tool",
"tool": tool,
"input": input,
"is_error": is_error,
"content": content.chars().take(4000).collect::<String>(),
});
if let Err(e) = Self::run_one(hook, &payload, workdir).await {
tracing::warn!("post_tool hook `{}` failed: {e}", hook.command);
}
}
}
pub async fn session_end(
&self,
session_id: &str,
path: &std::path::Path,
workdir: &std::path::Path,
) {
for hook in self.hooks.iter().filter(|h| h.event == Event::SessionEnd) {
let payload = serde_json::json!({
"event": "session_end",
"session_id": session_id,
"path": path,
});
if let Err(e) = Self::run_one(hook, &payload, workdir).await {
tracing::warn!("session_end hook `{}` failed: {e}", hook.command);
}
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use serde_json::json;
fn cfg(event: &str, command: &str) -> HookConfig {
HookConfig {
event: event.into(),
command: command.into(),
tools: Vec::new(),
timeout_secs: Some(5),
}
}
#[test]
fn an_unknown_event_is_a_startup_error() {
let err = HookSet::from_config(&[cfg("pre-tool", "true")])
.unwrap_err()
.to_string();
assert!(err.contains("pre-tool"), "{err}");
assert!(HookSet::from_config(&[cfg("pre_tool", " ")]).is_err());
}
#[tokio::test]
async fn exit_zero_allows_and_exit_two_denies_with_the_reason() {
let set = HookSet::from_config(&[cfg("pre_tool", "true")]).unwrap();
let v = set
.pre_tool("echo", &json!({}), std::path::Path::new("."))
.await;
assert_eq!(v, HookVerdict::Allow);
let set = HookSet::from_config(&[cfg("pre_tool", "echo not today; exit 2")]).unwrap();
let v = set
.pre_tool("echo", &json!({}), std::path::Path::new("."))
.await;
assert_eq!(v, HookVerdict::Deny("not today".into()));
}
#[tokio::test]
async fn an_undefined_exit_code_fails_closed() {
let set = HookSet::from_config(&[cfg("pre_tool", "exit 1")]).unwrap();
match set
.pre_tool("echo", &json!({}), std::path::Path::new("."))
.await
{
HookVerdict::Deny(reason) => assert!(reason.contains("exited 1"), "{reason}"),
HookVerdict::Allow => panic!("an undefined exit code must not be permission"),
}
}
#[tokio::test]
async fn a_hook_that_hangs_fails_closed_at_its_timeout() {
let mut c = cfg("pre_tool", "sleep 30");
c.timeout_secs = Some(1);
let set = HookSet::from_config(&[c]).unwrap();
match set
.pre_tool("echo", &json!({}), std::path::Path::new("."))
.await
{
HookVerdict::Deny(reason) => assert!(reason.contains("failed to run"), "{reason}"),
HookVerdict::Allow => panic!("a timeout must not be permission"),
}
}
#[tokio::test]
async fn a_hook_that_never_reads_a_large_payload_still_times_out() {
let mut c = cfg("pre_tool", "sleep 30");
c.timeout_secs = Some(1);
let set = HookSet::from_config(&[c]).unwrap();
let big = json!({"content": "x".repeat(256 * 1024)});
let verdict = tokio::time::timeout(
std::time::Duration::from_secs(5),
set.pre_tool("fs_write", &big, std::path::Path::new(".")),
)
.await
.expect("the hook's own timeout must fire; the write must not wedge it");
assert!(matches!(verdict, HookVerdict::Deny(_)));
}
#[tokio::test]
async fn the_tool_filter_scopes_a_hook_and_the_payload_reaches_stdin() {
let marker = std::env::temp_dir().join(format!("mecha-hook-{}", uuid::Uuid::new_v4()));
let mut c = cfg("pre_tool", &format!("cat > {}; exit 2", marker.display()));
c.tools = vec!["shell".into()];
let set = HookSet::from_config(&[c]).unwrap();
let v = set
.pre_tool("echo", &json!({}), std::path::Path::new("."))
.await;
assert_eq!(v, HookVerdict::Allow);
assert!(!marker.exists());
let v = set
.pre_tool(
"shell",
&json!({"command": "rm -rf /"}),
std::path::Path::new("."),
)
.await;
assert!(matches!(v, HookVerdict::Deny(_)));
let written = std::fs::read_to_string(&marker).unwrap();
assert!(written.contains("\"event\":\"pre_tool\""));
assert!(written.contains("rm -rf /"));
std::fs::remove_file(&marker).ok();
}
#[tokio::test]
async fn post_tool_failures_are_swallowed_because_observers_cannot_be_load_bearing() {
let set = HookSet::from_config(&[cfg("post_tool", "exit 7")]).unwrap();
set.post_tool("echo", &json!({}), false, "out", std::path::Path::new("."))
.await;
}
}