[opentitanlib] Refactor rescue protocol creation

sw/host/opentitanlib/src/io/uart.rs

@@ -18,7 +18,7 @@ use crate::impl_serializable_error;
 use crate::io::console::ConsoleDevice;
 use crate::transport::TransportError;
 
-#[derive(Clone, Debug, Args, Serialize, Deserialize)]
+#[derive(Clone, Debug, Default, Args, Serialize, Deserialize)]
 pub struct UartParams {
     /// UART instance.
     #[arg(long, default_value = "CONSOLE")]

sw/host/opentitanlib/src/rescue/mod.rs

@@ -2,13 +2,173 @@
 // Licensed under the Apache License, Version 2.0, see LICENSE for details.
 // SPDX-License-Identifier: Apache-2.0
 
+use anyhow::{ensure, Result};
+use clap::{Args, ValueEnum};
 use thiserror::Error;
 
+use crate::app::TransportWrapper;
+use crate::chip::boot_log::BootLog;
+use crate::chip::boot_svc::{BootSlot, BootSvc, OwnershipActivateRequest, OwnershipUnlockRequest};
+use crate::chip::device_id::DeviceId;
+use crate::io::uart::UartParams;
+use crate::with_unknown;
+
 pub mod serial;
 pub mod xmodem;
 
+pub use serial::RescueSerial;
+
 #[derive(Debug, Error)]
 pub enum RescueError {
     #[error("bad mode: {0}")]
     BadMode(String),
+    #[error("configuration error: {0}")]
+    Configuration(String),
+}
+
+#[derive(ValueEnum, Default, Debug, Clone, Copy, PartialEq)]
+pub enum RescueProtocol {
+    #[default]
+    Xmodem,
+    UsbDfu,
+    SpiDfu,
+}
+
+#[derive(ValueEnum, Default, Debug, Clone, Copy, PartialEq)]
+pub enum RescueTrigger {
+    #[default]
+    SerialBreak,
+    Gpio,
+    Strap,
+}
+
+#[derive(Clone, Default, Debug, Args)]
+pub struct RescueParams {
+    /// Rescue Protocol
+    #[arg(short, long, value_enum, default_value_t = RescueProtocol::Xmodem)]
+    pub protocol: RescueProtocol,
+    #[arg(short, long, value_enum, default_value_t = RescueTrigger::SerialBreak)]
+    pub trigger: RescueTrigger,
+    #[arg(short, long, default_value = "")]
+    pub value: String,
+    #[command(flatten)]
+    uart: UartParams,
+}
+
+impl RescueParams {
+    pub fn create(&self, transport: &TransportWrapper) -> Result<Box<dyn Rescue>> {
+        match self.protocol {
+            RescueProtocol::Xmodem => self.create_serial(transport),
+            RescueProtocol::UsbDfu => self.create_usbdfu(transport),
+            RescueProtocol::SpiDfu => self.create_spidfu(transport),
+        }
+    }
+    fn create_serial(&self, transport: &TransportWrapper) -> Result<Box<dyn Rescue>> {
+        ensure!(
+            self.trigger == RescueTrigger::SerialBreak,
+            RescueError::Configuration(format!(
+                "Xmodem does not support trigger {:?}",
+                self.trigger
+            ))
+        );
+        ensure!(
+            self.value.is_empty(),
+            RescueError::Configuration(format!(
+                "Xmodem does not support trigger value {:?}",
+                self.value
+            ))
+        );
+
+        Ok(Box::new(RescueSerial::new(self.uart.create(transport)?)))
+    }
+    fn create_usbdfu(&self, _transport: &TransportWrapper) -> Result<Box<dyn Rescue>> {
+        unimplemented!()
+    }
+    fn create_spidfu(&self, _transport: &TransportWrapper) -> Result<Box<dyn Rescue>> {
+        unimplemented!()
+    }
+}
+
+with_unknown! {
+pub enum RescueMode: u32 {
+    Rescue = u32::from_be_bytes(*b"RESQ"),
+    RescueB = u32::from_be_bytes(*b"RESB"),
+    BootLog = u32::from_be_bytes(*b"BLOG"),
+    BootSvcReq = u32::from_be_bytes(*b"BREQ"),
+    BootSvcRsp = u32::from_be_bytes(*b"BRSP"),
+    OwnerBlock = u32::from_be_bytes(*b"OWNR"),
+    GetOwnerPage0 = u32::from_be_bytes(*b"OPG0"),
+    GetOwnerPage1 = u32::from_be_bytes(*b"OPG1"),
+    DeviceId = u32::from_be_bytes(*b"OTID"),
+    EraseOwner = u32::from_be_bytes(*b"KLBR"),
+}
+}
+
+pub trait Rescue {
+    fn enter(&self, transport: &TransportWrapper, reset_target: bool) -> Result<()>;
+    fn set_mode(&self, mode: RescueMode) -> Result<()>;
+    fn send(&self, data: &[u8]) -> Result<()>;
+    fn recv(&self) -> Result<Vec<u8>>;
+
+    // Not supported by all backends
+    fn set_speed(&self, speed: u32) -> Result<u32>;
+    fn wait(&self) -> Result<()>;
+    fn reboot(&self) -> Result<()>;
+
+    fn get_raw(&self, mode: RescueMode) -> Result<Vec<u8>> {
+        self.set_mode(mode)?;
+        self.recv()
+    }
+
+    fn set_raw(&self, mode: RescueMode, data: &[u8]) -> Result<()> {
+        self.set_mode(mode)?;
+        self.send(data)
+    }
+
+    fn update_firmware(&self, slot: BootSlot, image: &[u8]) -> Result<()> {
+        let mode = if slot == BootSlot::SlotB {
+            RescueMode::RescueB
+        } else {
+            RescueMode::Rescue
+        };
+        self.set_raw(mode, image)
+    }
+
+    fn get_boot_log(&self) -> Result<BootLog> {
+        let blog = self.get_raw(RescueMode::BootLog)?;
+        Ok(BootLog::try_from(blog.as_slice())?)
+    }
+
+    fn get_boot_svc(&self) -> Result<BootSvc> {
+        let bsvc = self.get_raw(RescueMode::BootSvcRsp)?;
+        Ok(BootSvc::try_from(bsvc.as_slice())?)
+    }
+
+    fn get_device_id(&self) -> Result<DeviceId> {
+        let id = self.get_raw(RescueMode::DeviceId)?;
+        DeviceId::read(&mut std::io::Cursor::new(&id))
+    }
+
+    fn set_next_bl0_slot(&self, primary: BootSlot, next: BootSlot) -> Result<()> {
+        let message = BootSvc::next_boot_bl0_slot(primary, next);
+        self.set_raw(RescueMode::BootSvcReq, &message.to_bytes()?)
+    }
+
+    fn ownership_unlock(&self, unlock: OwnershipUnlockRequest) -> Result<()> {
+        let message = BootSvc::ownership_unlock(unlock);
+        self.set_raw(RescueMode::BootSvcReq, &message.to_bytes()?)
+    }
+
+    fn ownership_activate(&self, activate: OwnershipActivateRequest) -> Result<()> {
+        let message = BootSvc::ownership_activate(activate);
+        self.set_raw(RescueMode::BootSvcReq, &message.to_bytes()?)
+    }
+
+    fn set_owner_config(&self, data: &[u8]) -> Result<()> {
+        self.set_raw(RescueMode::OwnerBlock, data)
+    }
+
+    fn erase_owner(&self) -> Result<()> {
+        self.set_raw(RescueMode::EraseOwner, &[])
+    }
 }

sw/host/opentitanlib/src/rescue/serial.rs

@@ -7,12 +7,9 @@ use std::rc::Rc;
 use std::time::Duration;
 
 use crate::app::TransportWrapper;
-use crate::chip::boot_log::BootLog;
-use crate::chip::boot_svc::{BootSlot, BootSvc, OwnershipActivateRequest, OwnershipUnlockRequest};
-use crate::chip::device_id::DeviceId;
 use crate::io::uart::Uart;
 use crate::rescue::xmodem::Xmodem;
-use crate::rescue::RescueError;
+use crate::rescue::{Rescue, RescueError, RescueMode};
 use crate::uart::console::UartConsole;
 
 pub struct RescueSerial {
@@ -23,19 +20,9 @@ pub struct RescueSerial {
 
 impl RescueSerial {
     const ONE_SECOND: Duration = Duration::from_secs(1);
-    pub const RESCUE: [u8; 4] = *b"RESQ";
-    pub const RESCUE_B: [u8; 4] = *b"RESB";
-    pub const REBOOT: [u8; 4] = *b"REBO";
-    pub const BAUD: [u8; 4] = *b"BAUD";
-    pub const BOOT_LOG: [u8; 4] = *b"BLOG";
-    pub const BOOT_SVC_REQ: [u8; 4] = *b"BREQ";
-    pub const BOOT_SVC_RSP: [u8; 4] = *b"BRSP";
-    pub const OWNER_BLOCK: [u8; 4] = *b"OWNR";
-    pub const GET_OWNER_PAGE0: [u8; 4] = *b"OPG0";
-    pub const GET_OWNER_PAGE1: [u8; 4] = *b"OPG1";
-    pub const OT_ID: [u8; 4] = *b"OTID";
-    pub const ERASE_OWNER: [u8; 4] = *b"KLBR";
-    pub const WAIT: [u8; 4] = *b"WAIT";
+    pub const REBOOT: RescueMode = RescueMode(u32::from_be_bytes(*b"REBO"));
+    pub const BAUD: RescueMode = RescueMode(u32::from_be_bytes(*b"BAUD"));
+    pub const WAIT: RescueMode = RescueMode(u32::from_be_bytes(*b"WAIT"));
 
     const BAUD_115K: [u8; 4] = *b"115K";
     const BAUD_230K: [u8; 4] = *b"230K";
@@ -51,8 +38,10 @@ impl RescueSerial {
             enter_delay: Duration::from_secs(5),
         }
     }
+}
 
-    pub fn enter(&self, transport: &TransportWrapper, reset_target: bool) -> Result<()> {
+impl Rescue for RescueSerial {
+    fn enter(&self, transport: &TransportWrapper, reset_target: bool) -> Result<()> {
         log::info!("Setting serial break to trigger rescue mode.");
         self.uart.set_break(true)?;
         if reset_target {
@@ -67,7 +56,7 @@ impl RescueSerial {
         Ok(())
     }
 
-    pub fn set_baud(&self, baud: u32) -> Result<()> {
+    fn set_speed(&self, baud: u32) -> Result<u32> {
         // Make sure the requested rate is a known rate.
         let symbol = match baud {
             115200 => Self::BAUD_115K,
@@ -89,11 +78,13 @@ impl RescueSerial {
             return Err(RescueError::BadMode(result[0].clone()).into());
         }
         // Change our side of the connection to the new rate.
+        let old = self.uart.get_baudrate()?;
         self.uart.set_baudrate(baud)?;
-        Ok(())
+        Ok(old)
     }
 
-    pub fn set_mode(&self, mode: [u8; 4]) -> Result<()> {
+    fn set_mode(&self, mode: RescueMode) -> Result<()> {
+        let mode = mode.0.to_be_bytes();
         self.uart.write(&mode)?;
         let enter = b'\r';
         self.uart.write(std::slice::from_ref(&enter))?;
@@ -110,84 +101,26 @@ impl RescueSerial {
         Ok(())
     }
 
-    pub fn wait(&self) -> Result<()> {
+    fn wait(&self) -> Result<()> {
         self.set_mode(Self::WAIT)?;
         Ok(())
     }
 
-    pub fn reboot(&self) -> Result<()> {
+    fn reboot(&self) -> Result<()> {
         self.set_mode(Self::REBOOT)?;
         Ok(())
     }
 
-    pub fn update_firmware(&self, slot: BootSlot, image: &[u8]) -> Result<()> {
-        self.set_mode(if slot == BootSlot::SlotB {
-            Self::RESCUE_B
-        } else {
-            Self::RESCUE
-        })?;
+    fn send(&self, data: &[u8]) -> Result<()> {
         let xm = Xmodem::new();
-        xm.send(&*self.uart, image)?;
+        xm.send(&*self.uart, data)?;
         Ok(())
     }
 
-    pub fn get_raw(&self, mode: [u8; 4]) -> Result<Vec<u8>> {
-        self.set_mode(mode)?;
+    fn recv(&self) -> Result<Vec<u8>> {
         let mut data = Vec::new();
         let xm = Xmodem::new();
         xm.receive(&*self.uart, &mut data)?;
         Ok(data)
     }
-
-    pub fn get_boot_log(&self) -> Result<BootLog> {
-        let blog = self.get_raw(Self::BOOT_LOG)?;
-        Ok(BootLog::try_from(blog.as_slice())?)
-    }
-
-    pub fn get_boot_svc(&self) -> Result<BootSvc> {
-        let bsvc = self.get_raw(Self::BOOT_SVC_RSP)?;
-        Ok(BootSvc::try_from(bsvc.as_slice())?)
-    }
-
-    pub fn get_device_id(&self) -> Result<DeviceId> {
-        let id = self.get_raw(Self::OT_ID)?;
-        DeviceId::read(&mut std::io::Cursor::new(&id))
-    }
-
-    pub fn set_boot_svc_raw(&self, data: &[u8]) -> Result<()> {
-        self.set_mode(Self::BOOT_SVC_REQ)?;
-        let xm = Xmodem::new();
-        xm.send(&*self.uart, data)?;
-        Ok(())
-    }
-
-    pub fn set_next_bl0_slot(&self, primary: BootSlot, next: BootSlot) -> Result<()> {
-        let message = BootSvc::next_boot_bl0_slot(primary, next);
-        let data = message.to_bytes()?;
-        self.set_boot_svc_raw(&data)
-    }
-
-    pub fn ownership_unlock(&self, unlock: OwnershipUnlockRequest) -> Result<()> {
-        let message = BootSvc::ownership_unlock(unlock);
-        let data = message.to_bytes()?;
-        self.set_boot_svc_raw(&data)
-    }
-
-    pub fn ownership_activate(&self, activate: OwnershipActivateRequest) -> Result<()> {
-        let message = BootSvc::ownership_activate(activate);
-        let data = message.to_bytes()?;
-        self.set_boot_svc_raw(&data)
-    }
-
-    pub fn set_owner_config(&self, data: &[u8]) -> Result<()> {
-        self.set_mode(Self::OWNER_BLOCK)?;
-        let xm = Xmodem::new();
-        xm.send(&*self.uart, data)?;
-        Ok(())
-    }
-
-    pub fn erase_owner(&self) -> Result<()> {
-        self.set_mode(Self::ERASE_OWNER)?;
-        Ok(())
-    }
 }

sw/host/opentitanlib/src/rescue/xmodem.rs

@@ -188,7 +188,8 @@ impl Xmodem {
                 }
                 _ => {
                     return Err(XmodemError::UnsupportedMode(format!(
-                        "bad start of packet: {byte:?}"
+                        "bad start of packet: {byte:02x} ({})",
+                        byte as char
                     ))
                     .into());
                 }