mgmt flash: simplify protocol

artiq/coredevice/comm_mgmt.py

@@ -201,42 +201,17 @@ class CommMgmt:
     def debug_allocator(self):
         self._write_header(Request.DebugAllocator)
 
-    def flash(self, **bin_paths):
+    def flash(self, bin_paths):
         self._write_header(Request.Flash)
 
-        addr_table = {}
-        with io.BytesIO() as image_buf, io.BytesIO() as bin_buf:
-            offset = 0
-            # Reserve 4-bytes for CRC
-            image_buf.write(struct.pack(self.endian + "I", 0))
-            # Reserve 4-bytes for header length
-            image_buf.write(struct.pack(self.endian + "I", 0))
-            image_buf.write(struct.pack(self.endian + "I", len(bin_paths)))
-            for bin_name, filename in bin_paths.items():
+        with io.BytesIO() as image_buf:
+            for filename in bin_paths:
                 with open(filename, "rb") as fi:
                     bin_ = fi.read()
-                    length = bin_buf.write(bin_)
+                    image_buf.write(struct.pack(self.endian + "I", len(bin_)))
+                    image_buf.write(bin_)
 
-                    bin_name_str = bin_name.encode("utf-8")
-                    image_buf.write(struct.pack(self.endian + "I", len(bin_name_str)))
-                    image_buf.write(bin_name_str)
-                    image_buf.write(struct.pack(self.endian + "II", offset, length))
-
-                    offset += length
-
-            # header = image_buf.getvalue()
-            # image = image_buf.getvalue()
-
-            assert(image_buf.tell() == len(image_buf.getvalue()))
-            header_len = image_buf.tell() - 8
-            image_buf.seek(4, 0)
-            image_buf.write(struct.pack(self.endian + "I", header_len))
-            image_buf.seek(0, 2)
-            image_buf.write(bin_buf.getvalue())
-
-            image_buf.seek(4, 0)
-            crc = binascii.crc32(image_buf.read())
-            image_buf.seek(0, 0)
+            crc = binascii.crc32(image_buf.getvalue())
             image_buf.write(struct.pack(self.endian + "I", crc))
 
             self._write_bytes(image_buf.getvalue())

artiq/firmware/runtime/mgmt.rs

@@ -14,11 +14,12 @@ impl From<SchedError> for Error<SchedError> {
 
 mod local_coremgmt {
     use alloc::{string::String, vec::Vec};
+    use byteorder::{ByteOrder, NativeEndian};
     use crc::crc32;
     use log::LevelFilter;
 
     use board_misoc::{config, mem, spiflash};
-    use io::{Cursor, Write, ProtoWrite, ProtoRead, Error as IoError};
+    use io::{Write, ProtoWrite, Error as IoError};
     use logger_artiq::BufferLogger;
     use mgmt_proto::{Error, Reply};
     use sched::{Io, TcpStream, Error as SchedError};
@@ -136,44 +137,30 @@ mod local_coremgmt {
         Ok(())
     }
 
-    pub fn flash(_io: &Io, stream: &mut TcpStream, image: &Vec<u8>) -> Result<(), Error<SchedError>> {
-        let mut reader = Cursor::new(&image[..]);
-        let expected_crc = reader.read_u32().unwrap();
+    pub fn flash(_io: &Io, stream: &mut TcpStream, image: &[u8]) -> Result<(), Error<SchedError>> {
+        let (expected_crc, mut image) = {
+            let (image, crc_slice) = image.split_at(image.len() - 4);
+            (NativeEndian::read_u32(crc_slice), image)
+        };
 
-        let image = &image[4..];
         let actual_crc = crc32::checksum_ieee(image);
 
         if actual_crc == expected_crc {
-            info!("Checksum matched");
-            let header_size = reader.read_u32().unwrap() as usize;
-            let header_offset = reader.position();
-            let bin_offset = header_offset + header_size;
+            let bin_origins = [
+                ("gateware"  , 0                      ),
+                ("bootloader", mem::ROM_BASE          ),
+                ("firmware"  , mem::FLASH_BOOT_ADDRESS),
+            ];
 
-            let header = &image[header_offset..bin_offset];
-            let binaries = &image[bin_offset..];
+            for (name, origin) in bin_origins {
+                info!("Flashing {} binary...", name);
+                let size = NativeEndian::read_u32(&image[..4]) as usize;
+                image = &image[4..];
 
-            info!("found header of size {}", header.len());
+                let (bin, remaining) = image.split_at(size);
+                image = remaining;
 
-            let mut reader = Cursor::new(header);
-            let bin_no = reader.read_u32().unwrap() as usize;
-            for _ in 0..bin_no {
-                let bin_name = reader.read_string().unwrap();
-                let offset = reader.read_u32().unwrap() as usize;
-                let len = reader.read_u32().unwrap() as usize;
-
-                let origin = match bin_name.as_str() {
-                    "gateware" => 0,
-                    "bootloader" => mem::ROM_BASE,
-                    "firmware" => mem::FLASH_BOOT_ADDRESS,
-                    _ => {
-                        error!("unexpected binary component {}", bin_name);
-                        return Ok(Reply::Error.write_to(stream)?);
-                    }
-                };
-
-                unsafe {
-                    spiflash::flash_binary(origin, &binaries[offset..offset+len]);
-                }
+                unsafe { spiflash::flash_binary(origin, bin) };
             }
 
             reboot(_io, stream)?;

artiq/frontend/artiq_coremgmt.py

@@ -179,12 +179,8 @@ def main():
                 "Found firmware files: {}".format(" ".join(firmwares)))
         firmware = firmwares[0]
 
-        bins = {
-            "gateware": gateware,
-            "bootloader": bootloader,
-            "firmware": firmware,
-        }
-        mgmt.flash(**bins)
+        bins = [ gateware, bootloader, firmware ]
+        mgmt.flash(bins)
 
     if args.tool == "reboot":
         mgmt.reboot()