1 files changed, 138 insertions, 2 deletions

diff --git a/src/main.rs b/src/main.rs
index 6431c3d..4dde989 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -1,10 +1,11 @@
 use rustkrazy_admind::{Error, Result};
 
 use std::fs::{self, File};
-use std::io::{self, BufReader};
+use std::io::{self, BufReader, Write};
 
 use actix_web::{
-    dev::ServiceRequest, http::header::ContentType, web, App, HttpResponse, HttpServer,
+    dev::ServiceRequest, http::header::ContentType, web, App, FromRequest, HttpRequest,
+    HttpResponse, HttpServer,
 };
 use actix_web_httpauth::extractors::basic::{BasicAuth, Config};
 use actix_web_httpauth::extractors::AuthenticationError;
@@ -36,6 +37,38 @@ async fn handle_shutdown() -> HttpResponse {
     }
 }
 
+async fn handle_update_boot(req: HttpRequest) -> HttpResponse {
+    let boot = match boot_dev() {
+        Ok(v) => v,
+        Err(e) => {
+            return HttpResponse::InternalServerError()
+                .content_type(ContentType::plaintext())
+                .body(format!("can't locate boot partition: {}", e))
+        }
+    };
+
+    match stream_to(boot, req).await {
+        Ok(_) => {}
+        Err(e) => {
+            return HttpResponse::InternalServerError()
+                .content_type(ContentType::plaintext())
+                .body(format!("can't update boot partition: {}", e))
+        }
+    }
+
+    match switch_to_inactive_root() {
+        Ok(_) => HttpResponse::Ok()
+            .content_type(ContentType::plaintext())
+            .body("successfully updated boot partition and switched to inactive root"),
+        Err(e) => HttpResponse::InternalServerError()
+            .content_type(ContentType::plaintext())
+            .body(format!(
+                "can't switch to inactive root (this is probably fatal): {}",
+                e
+            )),
+    }
+}
+
 #[actix_web::main]
 async fn main() -> io::Result<()> {
     match start().await {
@@ -60,6 +93,7 @@ async fn start() -> Result<()> {
             .wrap(auth)
             .service(web::resource("/reboot").to(handle_reboot))
             .service(web::resource("/shutdown").to(handle_shutdown))
+            .service(web::resource("/update/boot").to(handle_update_boot))
     })
     .bind_rustls("[::]:8443", config)?
     .run()
@@ -121,3 +155,105 @@ fn validate_credentials(user_id: &str, user_password: &str) -> io::Result<bool>
         "Invalid credentials",
     ))
 }
+
+fn modify_cmdline(old: &str, new: &str) -> Result<()> {
+    let cmdline = fs::read_to_string("/boot/cmdline.txt")?;
+    fs::write("/boot/cmdline.txt", cmdline.replace(old, new))?;
+
+    Ok(())
+}
+
+fn dev() -> Result<&'static str> {
+    let devs = ["/dev/mmcblk0", "/dev/sda", "/dev/vda"];
+
+    for dev in devs {
+        if fs::metadata(dev).is_ok() {
+            return Ok(dev);
+        }
+    }
+
+    Err(Error::NoDiskDev)
+}
+
+fn boot_dev() -> Result<&'static str> {
+    Ok(match dev()? {
+        "/dev/mmcblk0" => "/dev/mmcblk0p1",
+        "/dev/sda" => "/dev/sda1",
+        "/dev/vda" => "/dev/vda1",
+        _ => unreachable!(),
+    })
+}
+
+fn active_root() -> Result<String> {
+    let cmdline = fs::read_to_string("/proc/cmdline")?;
+
+    for seg in cmdline.split(' ') {
+        if seg.starts_with("root=PARTUUID=00000000-") {
+            let root_id = seg
+                .split("root=PARTUUID=00000000-0")
+                .collect::<Vec<&str>>()
+                .into_iter()
+                .next_back()
+                .ok_or(Error::RootdevUnset)?;
+
+            return Ok(match dev()? {
+                "/dev/mmcblk0" => format!("/dev/mmcblk0p{}", root_id),
+                "/dev/sda" => format!("/dev/sda{}", root_id),
+                "/dev/vda" => format!("/dev/vda{}", root_id),
+                _ => unreachable!(),
+            });
+        }
+    }
+
+    Err(Error::RootdevUnset)
+}
+
+fn inactive_root() -> Result<String> {
+    let cmdline = 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()
+                .ok_or(Error::RootdevUnset)?
+            {
+                "2" => "3",
+                "3" => "2",
+                _ => unreachable!(),
+            };
+
+            return Ok(match dev()? {
+                "/dev/mmcblk0" => format!("/dev/mmcblk0p{}", root_id),
+                "/dev/sda" => format!("/dev/sda{}", root_id),
+                "/dev/vda" => format!("/dev/vda{}", root_id),
+                _ => unreachable!(),
+            });
+        }
+    }
+
+    Err(Error::RootdevUnset)
+}
+
+async fn stream_to(dst: &str, req: HttpRequest) -> Result<()> {
+    let bytes = web::Bytes::extract(&req).await?;
+    let mut file = File::create(dst)?;
+
+    file.write_all(&bytes)?;
+    file.sync_all()?;
+
+    Ok(())
+}
+
+fn switch_to_inactive_root() -> Result<()> {
+    let old = active_root()?;
+    let new = inactive_root()?;
+
+    let old = String::from("root=PARTUUID=00000000-0") + &old.chars().last().unwrap().to_string();
+    let new = String::from("root=PARTUUID=00000000-0") + &new.chars().last().unwrap().to_string();
+
+    modify_cmdline(&old, &new)?;
+    Ok(())
+}