hatchery_node/
node_control.rs

1//! Local control channel of this node (`POST /control/{drain,reload,handoff}`
2//! on the ops HTTP listener). Owner-only: a random bearer token written 0600
3//! (Unix) into the node's own state directory, a per-user location on Windows.
4//! `/health` publishes the token file path and the capabilities, never the token.
5//!
6//! - drain: mark the node draining (reported in `/health`); new work is the
7//!   caller's to stop sending. `?draining=false` clears it.
8//! - reload: no hot-reloadable config yet; answers what needs a handoff.
9//! - handoff: drain, answer 202, then after `release_after_ms` stop the ops
10//!   listener, which runs the node's normal graceful shutdown (its own PTY
11//!   sessions stop; nothing it did not spawn is touched) and frees every port
12//!   for the successor.
13
14use std::path::{Path, PathBuf};
15use std::sync::atomic::{AtomicBool, Ordering};
16use std::sync::OnceLock;
17use std::time::Duration;
18
19pub const CAPABILITIES: [&str; 3] = ["control.drain", "control.reload", "control.handoff"];
20const MAX_RELEASE_AFTER_MS: u64 = 10_000;
21
22struct Installed {
23    token: String,
24    file: PathBuf,
25}
26
27static INSTALLED: OnceLock<Installed> = OnceLock::new();
28static DRAINING: AtomicBool = AtomicBool::new(false);
29static HANDOFF: OnceLock<tokio::sync::Notify> = OnceLock::new();
30
31fn handoff_notify() -> &'static tokio::sync::Notify {
32    HANDOFF.get_or_init(tokio::sync::Notify::new)
33}
34
35/// Token file beside the node state.
36pub fn token_file_for(state_path: &Path) -> PathBuf {
37    state_path.parent().map_or_else(|| PathBuf::from("control.token"), |d| d.join("control.token"))
38}
39
40/// Create the token and write it owner-only. Call once at start.
41pub fn install(file: PathBuf) -> std::io::Result<()> {
42    let mut bytes = [0u8; 32];
43    getrandom::fill(&mut bytes).map_err(|e| std::io::Error::other(e.to_string()))?;
44    let token: String = bytes.iter().map(|b| format!("{b:02x}")).collect();
45    write_owner_only(&file, token.as_bytes())?;
46    let _ = INSTALLED.set(Installed { token, file });
47    Ok(())
48}
49
50pub fn write_owner_only(path: &Path, contents: &[u8]) -> std::io::Result<()> {
51    use std::io::Write;
52    if let Some(dir) = path.parent() {
53        std::fs::create_dir_all(dir)?;
54    }
55    let tmp = path.with_extension("token.tmp");
56    let mut options = std::fs::OpenOptions::new();
57    options.create(true).write(true).truncate(true);
58    #[cfg(unix)]
59    {
60        use std::os::unix::fs::OpenOptionsExt;
61        options.mode(0o600);
62    }
63    let mut file = options.open(&tmp)?;
64    file.write_all(contents)?;
65    file.sync_all()?;
66    drop(file);
67    std::fs::rename(&tmp, path)
68}
69
70pub fn installed() -> bool {
71    INSTALLED.get().is_some()
72}
73
74pub fn token_file() -> Option<String> {
75    INSTALLED.get().map(|i| i.file.display().to_string())
76}
77
78pub fn is_draining() -> bool {
79    DRAINING.load(Ordering::SeqCst)
80}
81
82pub fn authorized(header: Option<&str>) -> bool {
83    let (Some(installed), Some(value)) = (INSTALLED.get(), header) else { return false };
84    let Some(given) = value.strip_prefix("Bearer ") else { return false };
85    given.len() == installed.token.len()
86        && given.bytes().zip(installed.token.bytes()).fold(0u8, |a, (x, y)| a | (x ^ y)) == 0
87}
88
89/// Resolves when a handoff asked the ops listener to stop.
90pub async fn handoff_requested() {
91    handoff_notify().notified().await;
92}
93
94/// Execute a verb. Returns (status, JSON body).
95pub fn execute(verb: &str, query: &str) -> (u16, serde_json::Value) {
96    let param = |name: &str| {
97        query.split('&').find_map(|kv| kv.split_once('=').filter(|(k, _)| *k == name).map(|(_, v)| v.to_owned()))
98    };
99    match verb {
100        "drain" => {
101            let on = param("draining").map_or(true, |v| v != "false");
102            DRAINING.store(on, Ordering::SeqCst);
103            tracing::info!(draining = on, "control: drain");
104            (200, serde_json::json!({"ok": true, "draining": on}))
105        }
106        "reload" => (200, serde_json::json!({"ok": true, "applied": [], "needs_handoff_for": ["flags", "env"]})),
107        "handoff" => {
108            let after = param("release_after_ms").and_then(|v| v.parse().ok()).unwrap_or(300u64).min(MAX_RELEASE_AFTER_MS);
109            DRAINING.store(true, Ordering::SeqCst);
110            tracing::info!(release_after_ms = after, "control: handoff accepted; node will shut down gracefully");
111            tokio::spawn(async move {
112                tokio::time::sleep(Duration::from_millis(after)).await;
113                handoff_notify().notify_waiters();
114                handoff_notify().notify_one();
115            });
116            (202, serde_json::json!({"ok": true, "release_after_ms": after}))
117        }
118        _ => (404, serde_json::json!({"error": "unknown control verb"})),
119    }
120}
121
122/// Remove our token file at exit if it is still ours.
123pub fn cleanup() {
124    if let Some(i) = INSTALLED.get() {
125        if std::fs::read_to_string(&i.file).is_ok_and(|t| t.trim() == i.token) {
126            let _ = std::fs::remove_file(&i.file);
127        }
128    }
129}
hatchery_node/node_control.rs

hatchery_node/
node_control.rs
