forge
ExploreLog in

Tahinli/forge

public
0 watching
★ 0
10.5 KiB
1//! Self-update for an idle runner. Each claim response may advertise the sha256 of the forge's
2//! binary. When that differs from ours, the runner stops claiming, lets its running jobs finish,
3//! swaps its own binary for the forge's, and exits with [`RESTART_EXIT_CODE`] so systemd starts
4//! the new one.
5//!
6//! [`UpdateState`] makes the decisions and does no I/O. [`install`] does the file work.
7
8use std::io::Read;
9use std::os::unix::fs::PermissionsExt;
10use std::path::{Path, PathBuf};
11use std::process::Stdio;
12use std::time::Duration;
13
14use sha2::{Digest, Sha256};
15
16use crate::client::{ClientError, ForgeClient};
17
18/// The exit code after a successful swap. systemd restarts the runner on any non-zero code.
19pub const RESTART_EXIT_CODE: u8 = 75;
20/// How long the new binary gets to answer `--version`.
21const SMOKE_TIMEOUT: Duration = Duration::from_secs(10);
22/// The temp file, next to the running binary, that the download goes to.
23const NEW_NAME: &str = ".forge-runner.new";
24/// The previous binary, kept for a manual rollback.
25const PREV_SUFFIX: &str = ".prev";
26/// Optional list of bad binaries, one sha256 per line, next to the running binary.
27const BAD_SUFFIX: &str = ".bad";
28
29/// Decides whether an update is due. It remembers the last sha that failed, so the runner does
30/// not retry it on every claim. A different advertised sha clears that memory.
31#[derive(Clone, Debug, PartialEq, Eq)]
32pub struct UpdateState {
33 enabled: bool,
34 /// The sha256 of the running binary. Updates are disabled when it is unknown.
35 current: Option<String>,
36 pending: Option<String>,
37 given_up: Option<String>,
38}
39
40impl UpdateState {
41 /// `current` is the sha256 (lowercase hex) of the running binary.
42 pub fn new(auto_update: bool, current: Option<String>) -> Self {
43 Self {
44 enabled: auto_update && current.is_some(),
45 current,
46 pending: None,
47 given_up: None,
48 }
49 }
50
51 /// Records the sha the forge advertised on a claim response. Returns true when this makes a
52 /// new update due. An absent header changes nothing.
53 pub fn observe(&mut self, advertised: Option<&str>) -> bool {
54 if !self.enabled {
55 return false;
56 }
57 let Some(advertised) = advertised else {
58 return false;
59 };
60 if Some(advertised) == self.current.as_deref()
61 || Some(advertised) == self.given_up.as_deref()
62 {
63 // Already running it, or it failed and the forge has not moved on.
64 self.pending = None;
65 return false;
66 }
67 if self.pending.as_deref() == Some(advertised) {
68 return false;
69 }
70 self.pending = Some(advertised.to_owned());
71 true
72 }
73
74 /// Whether updates are on at all. False when `auto_update` is off or the binary is unknown.
75 pub fn enabled(&self) -> bool {
76 self.enabled
77 }
78
79 /// The sha to update to, if an update is due.
80 pub fn pending(&self) -> Option<&str> {
81 self.pending.as_deref()
82 }
83
84 /// Stops trying `sha` until the forge advertises a different one.
85 pub fn give_up(&mut self, sha: &str) {
86 self.given_up = Some(sha.to_owned());
87 self.pending = None;
88 }
89}
90
91/// Whether a `.bad` list, one sha256 per line, names `sha`.
92pub fn is_listed_bad(list: &str, sha: &str) -> bool {
93 list.lines()
94 .any(|line| line.trim().eq_ignore_ascii_case(sha))
95}
96
97/// The sha256 (lowercase hex) of the file at `path`.
98pub fn file_sha256(path: &Path) -> std::io::Result<String> {
99 let mut file = std::fs::File::open(path)?;
100 let mut hasher = Sha256::new();
101 let mut buf = vec![0u8; 64 * 1024];
102 loop {
103 let n = file.read(&mut buf)?;
104 if n == 0 {
105 break;
106 }
107 hasher.update(&buf[..n]);
108 }
109 Ok(hex::encode(hasher.finalize()))
110}
111
112/// The path of the running binary, and the sha256 of its contents.
113pub fn locate_self() -> std::io::Result<(PathBuf, String)> {
114 let exe = std::env::current_exe()?;
115 let sha = file_sha256(&exe)?;
116 Ok((exe, sha))
117}
118
119/// What [`install`] did when it did not fail.
120#[derive(Clone, Copy, Debug, PartialEq, Eq)]
121pub enum Installed {
122 /// The new binary is in place. The caller should exit with [`RESTART_EXIT_CODE`].
123 Replaced,
124 /// The `.bad` list names the advertised binary, so nothing was replaced.
125 Listed,
126}
127
128/// Why an update failed. The current binary is untouched in every case.
129#[derive(Debug, thiserror::Error)]
130pub enum UpdateError {
131 #[error("local io: {0}")]
132 Io(#[from] std::io::Error),
133 #[error("download: {0}")]
134 Download(#[from] ClientError),
135 #[error("sha256 mismatch: forge advertised {expected}, downloaded {actual}")]
136 Mismatch { expected: String, actual: String },
137 #[error("new binary failed its check: {0}")]
138 Smoke(String),
139}
140
141/// Downloads the forge's binary next to `exe`, checks it, and swaps it in. The old binary is kept
142/// as `<exe>.prev`. On any error the temp file is removed.
143pub async fn install(
144 client: &ForgeClient,
145 runner_secret: &str,
146 exe: &Path,
147 expected: &str,
148) -> Result<Installed, UpdateError> {
149 let new = exe.with_file_name(NEW_NAME);
150 // A temp file left by an earlier crash must not be reused.
151 let _ = std::fs::remove_file(&new);
152 let result = stage_and_swap(client, runner_secret, exe, &new, expected).await;
153 if !matches!(result, Ok(Installed::Replaced)) {
154 let _ = std::fs::remove_file(&new);
155 }
156 result
157}
158
159async fn stage_and_swap(
160 client: &ForgeClient,
161 runner_secret: &str,
162 exe: &Path,
163 new: &Path,
164 expected: &str,
165) -> Result<Installed, UpdateError> {
166 let actual = client.download_binary(runner_secret, new).await?;
167 if actual != expected {
168 return Err(UpdateError::Mismatch {
169 expected: expected.to_owned(),
170 actual,
171 });
172 }
173 std::fs::set_permissions(new, std::fs::Permissions::from_mode(0o755))?;
174 std::fs::File::open(new)?.sync_all()?;
175 smoke_test(new).await?;
176
177 let bad = match std::fs::read_to_string(sibling(exe, BAD_SUFFIX)) {
178 Ok(list) => list,
179 Err(e) if e.kind() == std::io::ErrorKind::NotFound => String::new(),
180 Err(e) => return Err(e.into()),
181 };
182 if is_listed_bad(&bad, expected) {
183 return Ok(Installed::Listed);
184 }
185
186 let prev = sibling(exe, PREV_SUFFIX);
187 match std::fs::remove_file(&prev) {
188 Ok(()) => {}
189 Err(e) if e.kind() == std::io::ErrorKind::NotFound => {}
190 Err(e) => return Err(e.into()),
191 }
192 // A hard link keeps the old inode alive at no cost. Copy only where links are not allowed.
193 if std::fs::hard_link(exe, &prev).is_err() {
194 std::fs::copy(exe, &prev)?;
195 }
196 std::fs::rename(new, exe)?;
197 Ok(Installed::Replaced)
198}
199
200/// Runs `<bin> --version` and requires exit 0 within [`SMOKE_TIMEOUT`].
201async fn smoke_test(bin: &Path) -> Result<(), UpdateError> {
202 let output = tokio::time::timeout(
203 SMOKE_TIMEOUT,
204 tokio::process::Command::new(bin)
205 .arg("--version")
206 .stdin(Stdio::null())
207 .kill_on_drop(true)
208 .output(),
209 )
210 .await
211 .map_err(|_| UpdateError::Smoke("--version did not finish in time".into()))??;
212 if !output.status.success() {
213 return Err(UpdateError::Smoke(format!(
214 "--version exited with {}",
215 output.status
216 )));
217 }
218 Ok(())
219}
220
221/// `exe` with `suffix` appended to its file name, for example `forge.prev`.
222fn sibling(exe: &Path, suffix: &str) -> PathBuf {
223 let mut name = exe.as_os_str().to_owned();
224 name.push(suffix);
225 PathBuf::from(name)
226}
227
228#[cfg(test)]
229mod tests {
230 use super::*;
231
232 const OURS: &str = "aaaa";
233 const THEIRS: &str = "bbbb";
234
235 fn state(auto_update: bool) -> UpdateState {
236 UpdateState::new(auto_update, Some(OURS.to_owned()))
237 }
238
239 #[test]
240 fn a_different_sha_becomes_pending_once() {
241 let mut s = state(true);
242 assert!(s.observe(Some(THEIRS)));
243 assert_eq!(s.pending(), Some(THEIRS));
244 // The same sha again is not news.
245 assert!(!s.observe(Some(THEIRS)));
246 assert_eq!(s.pending(), Some(THEIRS));
247 }
248
249 #[test]
250 fn matching_sha_or_no_header_leaves_nothing_pending() {
251 let mut s = state(true);
252 assert!(!s.observe(Some(OURS)));
253 assert_eq!(s.pending(), None);
254 assert!(!s.observe(None));
255 assert_eq!(s.pending(), None);
256 }
257
258 #[test]
259 fn absent_header_keeps_a_pending_update() {
260 let mut s = state(true);
261 s.observe(Some(THEIRS));
262 s.observe(None);
263 assert_eq!(s.pending(), Some(THEIRS));
264 }
265
266 #[test]
267 fn disabled_or_unknown_current_never_pends() {
268 let mut off = state(false);
269 assert!(!off.observe(Some(THEIRS)));
270 assert_eq!(off.pending(), None);
271
272 let mut unknown = UpdateState::new(true, None);
273 assert!(!unknown.observe(Some(THEIRS)));
274 assert_eq!(unknown.pending(), None);
275 }
276
277 #[test]
278 fn a_failed_sha_is_not_retried_until_the_forge_moves_on() {
279 let mut s = state(true);
280 s.observe(Some(THEIRS));
281 s.give_up(THEIRS);
282 assert_eq!(s.pending(), None);
283
284 // Claims keep advertising the same sha: nothing becomes pending.
285 assert!(!s.observe(Some(THEIRS)));
286 assert_eq!(s.pending(), None);
287
288 // A different sha is due again.
289 assert!(s.observe(Some("cccc")));
290 assert_eq!(s.pending(), Some("cccc"));
291 }
292
293 #[test]
294 fn going_back_to_the_failed_sha_cancels_a_pending_update() {
295 let mut s = state(true);
296 s.observe(Some(THEIRS));
297 s.give_up(THEIRS);
298 s.observe(Some("cccc"));
299 s.observe(Some(THEIRS));
300 assert_eq!(s.pending(), None);
301 }
302
303 #[test]
304 fn bad_list_matches_whole_lines_case_insensitively() {
305 let list = "# comment is not a sha\nBBBB\n cccc \n";
306 assert!(is_listed_bad(list, THEIRS));
307 assert!(is_listed_bad(list, "cccc"));
308 assert!(!is_listed_bad(list, "bbb"));
309 assert!(!is_listed_bad("", THEIRS));
310 }
311
312 #[test]
313 fn sibling_appends_the_suffix() {
314 assert_eq!(
315 sibling(Path::new("/home/t/.local/bin/forge"), PREV_SUFFIX),
316 PathBuf::from("/home/t/.local/bin/forge.prev")
317 );
318 }
319
320 #[test]
321 fn file_sha256_matches_a_known_digest() {
322 let dir = tempfile::tempdir().expect("tempdir");
323 let path = dir.path().join("f");
324 std::fs::write(&path, b"abc").expect("write");
325 assert_eq!(
326 file_sha256(&path).expect("hash"),
327 "ba7816bf8f01cfea414140de5dae2223b00361a396177a9cb410ff61f20015ad"
328 );
329 }
330}