FSMs + SansIO refactor

[brainstorm]

This is a relatively big refactor but needed pre-requisite to introduce other I/O (protocols) such as SPI/I2C/CAN to this project.

Current state machine experiments/drafting are happening in: https://github.com/brainstorm/ssh_fsm_tests ... on a std context for testing convenience (faster iteration/experimentation loop).

Moving to compile-time validated FSMs will give some peace of mind w.r.t, for instance, pre-auth attacks that have been unfortunately common in recent SSH server explorations, see:

https://threatprotect.qualys.com/2025/04/21/erlang-otp-ssh-server-remote-code-execution-vulnerability-cve-2025-32433/
https://www.runzero.com/blog/sshamble-unexpected-exposures-in-the-secure-shell/

The trigger of this refactoring idea is https://www.firezone.dev/blog/sans-io, for reference. This approach should be combined with some kind of profiling to ensure we don't regress performance-wise.

Hopefully, by the time of tackling this task, there will be some suitable profiling facility for esp-hal's no-std?: https://docs.espressif.com/projects/esp-idf/en/latest/esp32/api-guides/app_trace.html

[brainstorm]

Early drafts of FSMs for both application states and wifi auth ones:

[brainstorm]

Good ideas to borrow/use in this FSM no_std crate: https://github.com/alexfertel/nanomachine ... biggest drawback is its use of alloc though, unfortunately.

[Autofix]

I've added some additional states and events as I've tried to formalize the FSM. It would be best to constrain all the state changes to purely events but how it is now there are ENV settings and task checks required.

Created using https://mermaid.js.org/

---
config:
  theme: redux
  look: classic
  layout: dagre
---
flowchart TD
    A(["Start"]) --> B["Power On"]
    B --> C["Initialise Peripherals"]
    C -- Peripherals OK --> D{"Device provisioned"}
    D -- Provisioned --> E{"Read stored variables"}
    D -- Not provisioned --> F["Default AP"]
    I["Authorisation Checks - No Bridge"] -- Access Granted --> J["SSH Connect"]
    J -- SSH Connection Established --> K["Read ENV Variables"]
    K -- Variables Read --> L["Store settings"]
    E -- AP mode --> n4["Start AP Mode"]
    E -- STA mode --> n5["Start STA mode"]
    n4 -- AP Up --> n6["SSH Init"]
    n5 -- STA Up --> n6
    n6 -- SSH Ready --> n27{"UART Up?"}
    n11["Idle"] --> n12["Client Connecting"]
    n12 -- Connection Established --> I
    n17["Client Connected - No bridge"] --> n19["Notify no bridge"]
    M["Client Connected"] -- Client Disconnected --> n11
    n12 -- TCP Timeout --> n11
    I -- Access Denied --> n11
    I -- Timeout --> n11
    J -- Connection dropped --> n11
    K -- Input validation error --> n20["Notify ENV variable validation error"]
    n20 -- Notification Sent --> n11
    n21["Update SSH Settings"] -- Notification Sent --> n6
    n22["Update UART Setting"] -- Notification Sent' --> n8["UART Init"]
    n26["Bridge"] -- Bridge OK --> n11
    n8 -- UART Ok --> n26
    n27 -- UART Ok --> n26
    n27 -- UART Error --> n8
    n28{"Bridge Up?"} -- Bridge OK --> M
    n28 -- Bridge Error --> n17
    n23["Update Wifi Settings"] -- Wifi Settings Updated --> E
    n19 -- Client disconnects --> n11
    F -- Default AP up --> n6
    L -- Settings stored --> n29{"Settings Changed"}
    n29 -- No changes --> n28
    n29 -- UART Settings changed --> n22
    n29 -- Wifi settings changed --> n23
    n29 -- SSH Settings Changed --> n21
    n8 -- UART Error --> n11

[Autofix]

Latest FSM diagram which incorporates more error events, a reset state, and all state decisions have been replaced with specific events to work better with the next/run code structure.

Created using https://mermaid.js.org/

---
config:
  theme: redux
  layout: dagre
  look: neo
---
flowchart TD
    n1(["Start"]) --> n2["Power On"]
    n2 -- All Good --> n3["Initialise Peripherals"]
    n3 -- Peripherals OK --> n7["Initialise TCP"]
    n7 -- Start default AP --> n9["Starting Default AP"]
    n10["Authorisation Checks"] -- Access Granted --> n13["SSH Connect"]
    n13 -- SSH Connection Established --> n14["Read ENV Variables"]
    n14 -- Variables Read --> n16["Store settings"]
    n4["Starting AP Mode"] -- AP Up --> n6["SSH Init"]
    n5["Starting STA mode"] -- AP up --> n6
    n11["Idle"] --> n12["Client Connecting"]
    n12 -- Connection Established --> n10
    n17["Client Connected - No bridge"] --> n19["Notify no bridge"]
    n15["Client Connected"] -- Client Disconnected --> n11
    n12 -- TCP Timeout --> n11
    n10 -- Access Denied --> n11
    n10 -- Timeout --> n11
    n13 -- Connection dropped --> n11
    n14 -- Input validation error --> n20["Notify ENV variable validation error"]
    n20 -- Notification Sent --> n11
    n21["Update SSH Settings"] -- Notification Sent --> n6
    n22["Update UART Setting"] -- Notification Sent' --> n30["Reset"]
    n26["Bridge"] -- Bridge OK --> n11
    n8["UART Check/Init"] -- UART Ok --> n26
    n28["Check Bridge"] -- Bridge OK --> n15
    n28 -- Bridge Error --> n17
    n23["Update Wifi Settings"] -- Wifi Settings Updated --> n7
    n19 -- Client disconnects --> n11
    n9 -- AP up --> n6
    n16 -- No changes --> n28
    n16 -- UART Settings changed --> n22
    n16 -- Wifi settings changed --> n23
    n16 -- SSH Settings Changed --> n21
    n8 -- UART Error --> n11
    n7 -- Start AP Mode --> n4
    n7 -- Start STA mode --> n5
    n7 -- Fail --> n30
    n3 -- Fail --> n30
    n30 -- All Good --> n3
    n9 -- Fail --> n30
    n5 -- Fail --> n30
    n4 -- Fail --> n30
    n6 -- SSH Ready --> n8
    n26 -- Bridge Error --> n6

[Autofix]

FSM exploration carried out in a fork of @brainstorm's separate repo: https://github.com/Autofix/ssh_fsm_tests/tree/fsm-library
I have used embassy_sync mutexes for sharing variables between the various state_machine.run functions as a first but incomplete solution and I am looking for a better solution as this does not feel correct or work in all instances.
The mutexes are basically used as a global variable and only used a single time as I am trying to avoid passing everything down from the main function as each variable will go down through an additional 4 function calls adding to the stack size each time.

Create: static RADIO_CLOCK_MUTEX: Mutex<CriticalSectionRawMutex, Option<RADIO_CLK>> = Mutex::new(None);
Initialise: *(RADIO_CLOCK_MUTEX.lock().await) = Some(peripherals.RADIO_CLK);
Retrieve: let radio_clk: RADIO_CLK = RADIO_CLOCK_MUTEX.lock().await.take().unwrap();

Working with @mmalenic we came up with a solution to use Arc to only store the pointer but this requires us to use unsafe code. It also requires heapless 0.9 which was only released in August and embassy is still using heapless 0.8.
Create:

arc_pool!(RadioClkPool: RADIO_CLK<'static>);
static mut RADIO_CLOCK_BLOCK: ArcBlock<RadioClkPool> = ArcBlock::new();
static RADIO_CLOCK_MUTEX: Mutex<CriticalSectionRawMutex, Option<Arc<RadioClkPool>>> = Mutex::new(None);

Initialise:

let radio_clock_block: &'static mut ArcBlock<RADIO_CLK> = unsafe {
        static mut RADIO_CLOCK_BLOCK: ArcBlock<RADIO_CLK> = ArcBlock::new();
        addr_of_mut!(RADIO_CLOCK_BLOCK).as_mut().unwrap()
    };
    RadioClkPool.manage(radio_clock_block);
    *(RADIO_CLOCK_MUTEX.lock().await) = Some(RadioClkPool.alloc(peripherals.RADIO_CLK).unwrap());

Retrieve let radio_clk = unsafe {RADIO_CLOCK_MUTEX.lock().await.take().unwrap().clone_unchecked()};

https://www.firezone.dev/blog/sans-io gives suggestions to avoid this architecture:

To mutate state that isn't local to the Future, you basically have two options: for:

• Use reference-counted pointers and a mutex, i.e. Arc<Mutex>
• Use "actors" and connect them via channels, i.e. spawn multiple tasks with loops that read and write to channels

Both of these options have a runtime overhead: Locks can result in contention and sending messages through channels requires copying. In addition, multiple tasks running inside a runtime operate in a non-deterministic order which can easily lead to race conditions and in the worst case, deadlocks. It appears that with either of these options, we arrive at a design that feels brittle, is prone to deadlocks and no longer employs zero-cost abstractions, yet avoiding all of these is one of the reasons we wanted to use Rust in the first place!

In the sans-IO world, these problems don't exist. Our protocol code doesn't spawn any tasks and thus, &mut self is all we need to mutate state. Without tasks or threads, we also don't need synchronisation primitives like Mutex. Without channels, there is no need to copy data: The state machine can simply directly reference the buffer we passed to the socket.

I believe the non-recommended way is along the lines of the current path I have gone down. I have not yet worked out the best way the recommended way could be implemented with my current FSM structure.

[brainstorm]

@jamesmunns We are a bit on an embedded Rust architecture pickle here w.r.t: how to use of Mutex/Arc (or, preferably, not use them)... after many discussions and refactoring attempts, we thought about summarising our attempts (as well written by @Autofix above ☝🏻) and then asking somebody with more Rust embedded experience.

What would you do in our current situation? What is the "best" way forward or at least, what's the better pattern here?

[mmalenic]

An approach with generics and traits (doesn't necessarily have to be this way if it's too cluttered). This would let the "ssh-stamp" side of the code depend on the generic FSM code:

mod ssh_stamp {
    use crate::fsm::StateMachineTasks;

    pub struct SSHStamp {
        power_up: u64,
    }

    impl SSHStamp {
        pub fn new() -> SSHStamp {
            Self { power_up: 1 }
        }
    }

    impl StateMachineTasks for SSHStamp {
        fn power_up(&mut self) {
            self.power_up = 2;
            println!("{:#?}", self.power_up);
        }
    }
}

mod fsm {
    pub trait StateMachineTasks {
        fn power_up(&mut self);
    }

    pub enum State {
        PowerUp,
        Done
    }

    pub struct StateMachine<C> {
        state: State,
        context: C
    }

    impl<C: StateMachineTasks> StateMachine<C> {
        pub fn new(context: C) -> StateMachine<C> {
            Self { state: State::PowerUp, context }
        }

        pub fn run(&mut self) {
            match self.state {
                State::PowerUp => {
                    self.context.power_up();
                    self.state = State::Done;
                },
                State::Done => {}
            }
        }
    }
}

use crate::fsm::StateMachine;
use crate::ssh_stamp::SSHStamp;

fn main() {
    let mut state_machine = StateMachine::new(SSHStamp::new());
    state_machine.run();
}

[jamesmunns]

@brainstorm if you want to find a time to chat (text or a call), I'm happy to find a 1h block or something this week to discuss. Feel free to DM me.

I don't want to "parachute in" and give a bunch of advice that is either confusing or misleading because I don't totally understand things. I also have my own opinions and bias, and probably just dumping all my (uninformed) thoughts probably is going to add noise rather than take it away :).

That being said, some vague suggestions are below, if they feel like they are not useful, feel free to ignore. It's mostly "stock" advice I give regularly, if that makes sense.

Architecturally, it might be worth identifying:

• What logic/behavior needs to run concurrently, e.g. what definitely needs to be in separate tasks.
• What resources are required in MULTIPLE of those concurrent tasks (e.g.: your shared resources), and which resources are only required in one of those tasks (e.g.: your exclusive resources)
• Of those resources, which can be statically initialized, and which require runtime init.
  • Static init is much easier: you can use a const fn constructor
  • Runtime init is a bit harder, but for this I very commonly use a construct like StaticCell<Mutex<Resource, _>>, and share the &'static Mutex<Resource, _> around the program. This might be preferrable to some of the unsafe init code listed above.
    • Note that embassy-sync's Mutex has sub-optimal behavior when multiple tasks are waiting on it at the same time: it only has a single Waker slot, which means two tasks can "fight over" the mutex, instead of acting in a fairly queued behavior. maitake-sync::Mutex is an alternative that uses intrusive queues to allow multiple wakers.
    • I greatly prefer this kind of pattern to a static Mutex<Option<Resource>, _>.
• How, if at all, your tasks need to interact with each other, either through shared resources, or some other medium.

I personally find a LOT of value of focusing on this aspect of the architecture, ESPECIALLY:

• If two "components" are actually totally overlapping in terms of resources, they probably don't need to be two components, this can help eliminate sharing. This runs against the desire to say "every aspect of the system stands alone", but is more a more data(-structure) driven design process.
• If two "resources" are ever (even only sometimes) taken together: they probably should be one resource and live together. This can help avoid deadlocks of trying to acquire two (or more) resources at the same time.
• And in general: try to limit sharing AS MUCH AS possible. It's definitely not always possible, usually because you find that two things need to be done concurrently, but it's worth being aware that anywhere you have concurrent sharing, it will "cost" more in terms of verifying correctness.

And finally, my bias/opinions:

I don't personally find a ton of value in using sans-io or explicit state machine patterns. This is definitely a "hot take", and there are lots of people that disagree.

I personally prefer shaping my (shared) resources into specific data structures that enforce the invariants I care about, often using an async mutex inside of them, but enforcing "transactionality" of state transitions in the public API.

I prefer to just write my state machines as async code, that operates on the exclusive/shared resources that they contain, with state transitions and atomicity of those transactions being enforced by the public APIs of the data structures. I think writing FSMs as docs is very helpful for exploring the space/defining the behavior, but I prefer encoding them as regular async code, rather than explicit enum + match or similar macro constructs.

Take all of these opinions for what you paid for them, but that's my take on designing systems like these. I'm not saying it's the only "right" way to do this, but it's how I tend to approach architectural design.

[Autofix]

Following the discussion with @jamesmunns, I have made an initial implementation of a Hierarchical State Machine (HSM) which is harder to test but greatly reduces the complexity and reasoning of the program flow.
https://github.com/Autofix/ssh-stamp/tree/ssh-stamp-hsm

I have tried to reduce concurrency by reducing the use of embassy tasks & mutexes, and constraining each variable to the HSM level that they are required. It's probably not the most ideal layout based on the number of structs I ended up using, so hopefully additional eyes may be able to simplify and point me towards a better rust way of achieving this.

The next step is to implement the cleanup functions that will run whenever a certain level is required to return, either from an error or for example when changing config.

config:
  layout: dagre
---
stateDiagram
  direction TB
  state Power_Good {
    direction TB
    Initialise_Peripherals
    state Peripherals_Initialised {
      direction TB
      Initialise_TCP
      state TCP_Initialised {
        direction TB
        Start_Wifi --> Start_Default_AP
        Start_Wifi --> Start_AP_Mode
        Start_Wifi --> Start_STA_Mode
        state Wifi_Initialised {
          direction TB
          SSH_Initialise
          state SSH_Initialised {
            direction TB
            Config_UART
            state UART_Pins_Configured {
              direction TB
              UART_Init
              state UART_Initialised {
                direction TB
                Idle --> ClientConnecting:Entry - Client Connects
                ClientConnecting --> Idle:Exit - Drop ssh connection

                state ClientConnected {
                direction TB
                    ReadConfig --> UpdateConfig:Config Changed

                    ReadConfig --> Check_Bridge:No config change, Start bridge
                    Check_Bridge --> Connected_No_Bridge:Bridge error
                    Connected_No_Bridge --> Notify_Client:Bridge Error
                    state BridgeConnected {
                    direction TB
                    ClientConnectedWithBridge
                    }
                }
              }
            }
          }
        }
      }
    }
  }
  Power_On --> Initialise_Peripherals:All Good
  ClientConnectedWithBridge --> Idle:Exit - Client Disconnects, Drop Bridge
  Config_UART --> UART_Init:Entry - Pins Ok
  UART_Init --> Config_UART:Exit - Drop Pins
  SSH_Initialise --> Config_UART:Entry - SSH Ok
  Config_UART --> SSH_Initialise:Exit - Drop SSH
  Start_Wifi --> Initialise_TCP:Exit - Drop TCP
  UART_Init --> Idle:Entry - Uart Ok
  Idle --> UART_Init:Exit - Drop Uart
  Initialise_Peripherals --> Initialise_TCP:Entry - Peripherals Ok
  Initialise_TCP --> Initialise_Peripherals:Exit - Drop peripherals
  UpdateConfig --> Config_UART:Exit - UART Pins Changed
  UpdateConfig --> SSH_Initialise:Exit - SSH Settings Changed
  UpdateConfig --> Start_Wifi:Exit - Wifi Settings Changed
  Initialise_TCP --> Start_Wifi:Entry - TCP Ok
  Start_Default_AP --> SSH_Initialise:Entry - Wifi Ok
  Start_AP_Mode --> SSH_Initialise:Entry - Wifi Ok
  Start_STA_Mode --> SSH_Initialise:Entry - Wifi Ok
  SSH_Initialise --> Start_Wifi:Exit - Drop Wifi
  Check_Bridge --> ClientConnectedWithBridge:Entry - Bridge Ok
    ClientConnecting --> ReadConfig
    Notify_Client --> Idle: Client Disconnects