Implement rollback and reboot

This is a copy of the logic from rustkrazy_admind (to prevent a high
amount of dependencies for an HTTPS client). Signals cannot be used
without losing the ability to handle errors.

1 files changed, 84 insertions, 2 deletions

diff --git a/src/main.rs b/src/main.rs
index f08e8a6..ebe921f 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -23,6 +23,9 @@ use std::io::{self, Read, Write};
 use std::net::{Shutdown, SocketAddr, TcpListener, TcpStream};
 use std::time::{Duration, Instant};
 
+use nix::sys::signal::Signal;
+use nix::unistd::Pid;
+
 const UPDATE_FILE: &str = "/data/update";
 const LISTEN_SOCKET: &str = "[::]:12808";
 const MAGIC: [u8; 4] = [0x32, 0x7f, 0xfe, 0x4c];
@@ -283,9 +286,88 @@ fn check_connectivity(
 }
 
 fn rollback() -> io::Result<()> {
-    todo!()
+    switch_to_inactive_root()?;
+    reboot()
 }
 
 fn reboot() -> io::Result<()> {
-    todo!()
+    eprintln!("[info] connection unhealthy, rebooting");
+
+    match nix::sys::signal::kill(Pid::from_raw(1), Signal::SIGUSR1) {
+        Ok(_) => Ok(()),
+        Err(e) => Err(io::Error::from(e)),
+    }
+}
+
+fn modify_cmdline(new: &str) -> io::Result<()> {
+    let boot = boot_dev().ok_or(io::Error::from(io::ErrorKind::NotFound))?;
+    let boot_partition = std::fs::OpenOptions::new()
+        .read(true)
+        .write(true)
+        .open(boot)?;
+    let buf_stream = fscommon::BufStream::new(boot_partition);
+    let bootfs = fatfs::FileSystem::new(buf_stream, fatfs::FsOptions::new())?;
+
+    let mut file = bootfs.root_dir().open_file("cmdline.txt")?;
+
+    file.write_all(new.as_bytes())?;
+
+    Ok(())
+}
+
+fn dev() -> Option<&'static str> {
+    let devs = ["/dev/mmcblk0", "/dev/sda", "/dev/vda"];
+    devs.into_iter().find(|&dev| std::fs::metadata(dev).is_ok())
+}
+
+fn boot_dev() -> Option<&'static str> {
+    Some(match dev()? {
+        "/dev/mmcblk0" => "/dev/mmcblk0p1",
+        "/dev/sda" => "/dev/sda1",
+        "/dev/vda" => "/dev/vda1",
+        _ => unreachable!(),
+    })
+}
+
+fn inactive_root() -> io::Result<String> {
+    let cmdline = std::fs::read_to_string("/proc/cmdline")?;
+
+    for seg in cmdline.split(' ') {
+        if seg.starts_with("root=PARTUUID=00000000-") {
+            let root_id = match seg
+                .split("root=PARTUUID=00000000-0")
+                .collect::<Vec<&str>>()
+                .into_iter()
+                .next_back()
+                .expect("no root device")
+            {
+                "2" => "3",
+                "3" => "2",
+                _ => unreachable!(),
+            };
+
+            return Ok(
+                match dev().ok_or(io::Error::from(io::ErrorKind::NotFound))? {
+                    "/dev/mmcblk0" => format!("/dev/mmcblk0p{}", root_id),
+                    "/dev/sda" => format!("/dev/sda{}", root_id),
+                    "/dev/vda" => format!("/dev/vda{}", root_id),
+                    _ => unreachable!(),
+                },
+            );
+        }
+    }
+
+    Err(io::Error::from(io::ErrorKind::NotFound))
+}
+
+fn switch_to_inactive_root() -> io::Result<()> {
+    let new = inactive_root()?;
+    let new = String::from("root=PARTUUID=00000000-0") + &new.chars().last().unwrap().to_string();
+
+    let cmdline = format!("{} init=/bin/init rootwait console=tty1", new);
+
+    modify_cmdline(&cmdline)?;
+    nix::unistd::sync();
+
+    Ok(())
 }