self upgrade

Cargo.lock

@@ -400,6 +400,7 @@ dependencies = [
  "serde",
  "serde_json",
  "serde_yml",
+ "sha2",
  "tar",
  "tokio",
  "tokio-stream",

Cargo.toml

@@ -23,6 +23,7 @@ home = "0.5.11"
 ed25519-dalek = { version = "2.1.1", features = ["rand_core"] }
 bs58 = "0.5.1"
 chrono = "0.4.39"
+sha2 = "0.10.8"
 
 detee-shared = { git = "ssh://git@gitea.detee.cloud/noormohammedb/detee-shared", branch = "stable_01" }
 # detee-shared = { path = "../detee-shared" }

src/global.rs

@@ -1,8 +1,9 @@
 use anyhow::Result;
 use ed25519_dalek::SigningKey;
 use log::{info, warn};
+use sha2::{Digest, Sha256};
 use std::fs::File;
-use std::io::Write;
+use std::io::{Read, Write};
 use std::sync::LazyLock;
 
 pub const PACKAGE_ARCHIVE_POSTFIX: &str = "-enclave_package.tar.gz";
@@ -102,3 +103,18 @@ pub fn get_public_key() -> String {
     log::info!("Loaded the following public key: {pubkey}");
     pubkey
 }
+
+pub fn compute_sha256<P: AsRef<std::path::Path>>(path: P) -> Result<String> {
+    let mut file = File::open(path)?;
+    let mut hasher = Sha256::new();
+    let mut buffer = [0u8; 8192];
+    loop {
+        let bytes_read = file.read(&mut buffer).unwrap();
+        if bytes_read == 0 {
+            break;
+        }
+        hasher.update(&buffer[..bytes_read]);
+    }
+    let result = hasher.finalize();
+    Ok(format!("{:x}", result))
+}

src/main.rs

@@ -5,6 +5,7 @@ pub mod global;
 pub mod grpc;
 pub mod utils;
 
+use anyhow::anyhow;
 use anyhow::Result;
 use data::App;
 use detee_shared::sgx::pb::brain::brain_message_app;
@@ -19,6 +20,8 @@ use global::PUBLIC_KEY;
 use log::info;
 use log::warn;
 use std::collections::HashSet;
+use std::fs::File;
+use std::path::Path;
 use std::time::Duration;
 use tokio::sync::mpsc::Receiver;
 use tokio::sync::mpsc::Sender;
@@ -188,6 +191,13 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
     log::info!("Detee daemon running");
 
     loop {
+        if std::env::var("DAEMON_AUTO_UPGRADE") != Ok("OFF".to_string()) {
+            // This upgrade procedure will get replaced in prod. We need this for the testnet.
+            if let Err(e) = download_and_replace_binary() {
+                log::error!("Failed to upgrade detee-snp-daemon to newer version: {e}");
+            }
+        }
+
         let (brain_msg_tx, brain_msg_rx) = tokio::sync::mpsc::channel(6);
         let (daemon_msg_tx, daemon_msg_rx) = tokio::sync::mpsc::channel(6);
 
@@ -238,3 +248,26 @@ fn set_logging() {
         .format_timestamp(None)
         .init();
 }
+
+fn download_and_replace_binary() -> Result<()> {
+    use reqwest::blocking::get;
+    use std::os::unix::fs::PermissionsExt;
+    const TMP_DAEMON: &str = "/usr/local/bin/detee/new-daemon";
+    const BINARY: &str = "/usr/local/bin/detee-sgx-daemon";
+    let response = get("https://registry.detee.ltd/sgx/daemon/detee-sgx-daemon")?;
+    if !response.status().is_success() {
+        return Err(anyhow!("Failed to download file: {}", response.status()));
+    }
+    let mut tmp_file = File::create(Path::new(&TMP_DAEMON))?;
+    std::io::copy(&mut response.bytes()?.as_ref(), &mut tmp_file)?;
+    let new_hash = crate::global::compute_sha256(TMP_DAEMON)?;
+    let old_hash = crate::global::compute_sha256(BINARY)?;
+    log::debug!("Old binary hash: {old_hash}. New binary hash: {new_hash}");
+    if new_hash != old_hash {
+        std::fs::rename(BINARY, BINARY.to_string() + "_BACKUP")?;
+        std::fs::rename(TMP_DAEMON, BINARY)?;
+        std::fs::set_permissions(BINARY, std::fs::Permissions::from_mode(0o775))?;
+        std::process::exit(0);
+    }
+    Ok(())
+}