From fc1fd96f715a8a8bd30154aca666436ff6ae8892 Mon Sep 17 00:00:00 2001
From: occheung <dc@m-labs.hk>
Date: Wed, 4 Sep 2024 13:21:25 +0800
Subject: [PATCH] satman: implement boot file rewrite sequence

---
 src/Cargo.lock            |  2 ++
 src/satman/Cargo.toml.tpl |  2 ++
 src/satman/src/main.rs    | 51 ++++++++++++++++++++++++++++++++++++++-
 src/satman/src/mgmt.rs    | 31 ++++++++++++++++++++++++
 4 files changed, 85 insertions(+), 1 deletion(-)

diff --git a/src/Cargo.lock b/src/Cargo.lock
index c14d523..e00a922 100644
--- a/src/Cargo.lock
+++ b/src/Cargo.lock
@@ -501,7 +501,9 @@ name = "satman"
 version = "0.0.0"
 dependencies = [
  "build_zynq",
+ "byteorder",
  "core_io",
+ "crc",
  "cslice",
  "embedded-hal",
  "io",
diff --git a/src/satman/Cargo.toml.tpl b/src/satman/Cargo.toml.tpl
index d4fa63c..38f64e6 100644
--- a/src/satman/Cargo.toml.tpl
+++ b/src/satman/Cargo.toml.tpl
@@ -15,7 +15,9 @@ build_zynq = { path = "../libbuild_zynq" }
 
 [dependencies]
 log = { version = "0.4", default-features = false }
+byteorder = { version = "1.3", default-features = false }
 core_io = { version = "0.1", features = ["collections"] }
+crc = { version = "1.7", default-features = false }
 cslice = "0.3"
 embedded-hal = "0.2"
 
diff --git a/src/satman/src/main.rs b/src/satman/src/main.rs
index 7ec3aa0..bf8065b 100644
--- a/src/satman/src/main.rs
+++ b/src/satman/src/main.rs
@@ -4,7 +4,9 @@
 
 #[macro_use]
 extern crate log;
+extern crate byteorder;
 extern crate core_io;
+extern crate crc;
 extern crate cslice;
 extern crate embedded_hal;
 
@@ -1143,7 +1145,7 @@ fn process_aux_packet(
                         },
                     )
                 } else {
-                    drtioaux::send(0, &drtioaux::Packet::CoreMgmtReply { succeeded: false  })
+                    drtioaux::send(0, &drtioaux::Packet::CoreMgmtReply { succeeded: false })
                 }
             }
         }
@@ -1279,6 +1281,53 @@ fn process_aux_packet(
             error!("debug allocator not supported on zynq device");
             drtioaux::send(0, &drtioaux::Packet::CoreMgmtReply { succeeded: false })
         }
+        drtioaux::Packet::CoreMgmtFlashRequest {
+            destination: _destination,
+            last,
+            length,
+            data,
+        } => {
+            forward!(
+                router,
+                _routing_table,
+                _destination,
+                *rank,
+                *self_destination,
+                _repeaters,
+                &packet,
+                timer
+            );
+
+            core_manager.add_image_data(&data, length as usize);
+
+            if last {
+                drtioaux::send(0, &drtioaux::Packet::CoreMgmtDropLink)
+            } else {
+                drtioaux::send(0, &drtioaux::Packet::CoreMgmtReply { succeeded: true })
+            }
+        }
+        drtioaux::Packet::CoreMgmtDropLinkAck {
+            destination: _destination,
+        } => {
+            forward!(
+                router,
+                _routing_table,
+                _destination,
+                *rank,
+                *self_destination,
+                _repeaters,
+                &packet,
+                timer
+            );
+
+            unsafe {
+                csr::gt_drtio::txenable_write(0);
+            }
+            core_manager.write_image();
+            info!("reboot imminent");
+            slcr::reboot();
+            Ok(())
+        }
 
         p => {
             warn!("received unexpected aux packet: {:?}", p);
diff --git a/src/satman/src/mgmt.rs b/src/satman/src/mgmt.rs
index d667483..8b7e569 100644
--- a/src/satman/src/mgmt.rs
+++ b/src/satman/src/mgmt.rs
@@ -1,5 +1,7 @@
 use alloc::vec::Vec;
 
+use byteorder::{ByteOrder, NativeEndian};
+use crc::crc32;
 use io::{Cursor, ProtoRead, ProtoWrite};
 use libboard_artiq::{drtioaux_proto::SAT_PAYLOAD_MAX_SIZE,
                      logger::{BufferLogger, LogBufferRef}};
@@ -48,6 +50,7 @@ pub struct Manager<'a> {
     last_log: Sliceable,
     config_payload: Cursor<Vec<u8>>,
     last_value: Sliceable,
+    image_payload: Vec<u8>,
 }
 
 impl<'a> Manager<'_> {
@@ -57,6 +60,7 @@ impl<'a> Manager<'_> {
             last_log: Sliceable::new(0, Vec::new()),
             config_payload: Cursor::new(Vec::new()),
             last_value: Sliceable::new(0, Vec::new()),
+            image_payload: Vec::new(),
         }
     }
 
@@ -113,4 +117,31 @@ impl<'a> Manager<'_> {
             .map(|()| debug!("erase success"))
             .map_err(|err| warn!("failed to erase: {:?}", err))
     }
+
+    pub fn add_image_data(&mut self, data: &[u8], data_len: usize) {
+        self.image_payload.extend(&data[..data_len]);
+    }
+
+    pub fn write_image(&self) {
+        let mut image = self.image_payload.clone();
+        let image_ref = &image[..];
+        let bin_len = image.len() - 4;
+
+        let (image_ref, expected_crc) = {
+            let (image_ref, crc_slice) = image_ref.split_at(bin_len);
+            (image_ref, NativeEndian::read_u32(crc_slice))
+        };
+
+        let actual_crc = crc32::checksum_ieee(image_ref);
+
+        if actual_crc == expected_crc {
+            image.truncate(bin_len);
+            self.cfg.write("boot", image).expect("failed to write boot image");
+        } else {
+            panic!(
+                "CRC failed in SDRAM (actual {:08x}, expected {:08x})",
+                actual_crc, expected_crc
+            );
+        }
+    }
 }