Skip to content

Lecture 8.2: State Machines

Last lecture, we saw a simple example of a state machine modeling a vending machine. We specified it via three types:

  • a type of inputs, which is how we interact with the state machine;
  • a type of state, which holds the internal state of the machine;
  • and a type of outputs, which specifies how the machine returns information to the user.

Given these three states, we then defined a type of traces: a list of four-tuples, where each tuple holds the current state, an input, an output, and a new state. Then, we defined a way to randomly generate traces and check trace properties, which specifies whether a trace that the machine emits is a "good" one or not.

State Machines, more generally

To enable more abstract reasoning, we can generalize our definitions from last lecture to a general notion of a state machine. First, we can specify the behavior of the machine via the following type:

type ('s, 'i, 'o) machine = {
    init : 's; 
    step : 's -> 'i -> ('s * 'o)
}

This type of machine has three type parameters: 's, for the state, and 'i and 'o for inputs and outputs, respectively. The machine then has two components: an initial state, of type 's; and a step function (similar to our vending_machine function from last time) which takes in a state and an input, and returns a state and an output.

Now, we can talk about traces in general: 

type ('s, 'i, 'o) trace = ('s * 'i * 'o * 's) list

Just like our type of machines, our type of trace is parameterzed by the same three type variables. A trace is then this four-tuple of state, input, output, and new state. Now, let's run the machine for many steps to generate a trace:

let make_trace (m : ('s, 'i, 'o) machine) (is : 'i list) : ('s, 'i, 'o) trace = 
    let rec go curr is = 
        match is with 
        | [] -> []
        | i :: is' -> 
            let (s', o) = m.step curr i in 
            let rest = go s' is' in 
            (curr, i, o, s') :: rest
    in
    go m.init is

This make_trace function relies on a recursive function go, which builds up the trace one-by-one recursively on the list of inputs, is. The state curr acts as an accumulator for the new current state. We begin the execution by feeding m.init to go.

Finally, we can generate random traces, given a generator for inputs:

let gen_trace (m : ('s, 'i, 'o) machine) (gen_input : unit -> 'i) (n : int) : ('s, 'i, 'o) trace = 
    make_trace m (List.init (Random.int n) (fun _ -> gen_input ()))

Instantiating the generic framework

First, let's just see that the more general framework can be fit to our vending machine example from last time: 

let vend_machine : (vend_state, vend_input, vend_output) machine = {
    init = init_vend_state;
    step = vending_machine
}

let gen_vending_trace () : (vend_state, vend_input, vend_output) trace = 
    gen_trace vend_machine gen_input 10

Hardware Security Modules

We can use state machines for lots of different problems. Today, we'll now look at using a state machine to specify the behavior of a hardware security module, like a YubiKey. We will model a simplified version, that behaves as follows: on the device is a PIN pad, which lets you enter numbers 0-9; additionally, there is a button labeled "OK". The device holds a secret string, which is only output when the correct password is given.

In particular, the device works like this: when it is powered on, it lets you input a password on the PIN pad. If you input the correct password and press "OK", it will display "UNLOCKED" along with the password, and not respond to any more input after that. If you press "OK" without the correct password, it will display "ERROR", and not accept any more input. (If you type in the wrong password, you have to power cycle the device and start over).

We can model this desired behavior via a state machine and a corresponding trace property, which serves as the specification. First, let's define the inputs and outputs: 

type hsm_input = | PressNum of int | PressUnlock

type hsm_output = | ERROR | UNLOCKED of string | Ok

For outputs, Ok is the default output given when the device isn't unlocked or in an error state.

Now, let's mathematically specify the correctness of our machine. For security, we want to say that the device is only unlocked if the correct password is passed in. We can specify that as a trace property. Suppose the secret is S, the password is pw, and we are given a trace T. Then, the property we want to formalize is that if T contains the output UNLOCKED S, the trace must begin with exactly the inputs to compose the password pw.

Let's implement this specification: 

let hsm_trace_correct (secret : string) (pw : int list) (tr : ('s, hsm_input, hsm_output) trace) = 
    if List.exists (fun (_, _, out, _) -> out = UNLOCKED secret) tr then 
        (* Get the N-length prefix of inputs, where N is length of password *)
        let tr_prefix = List.take (List.length pw) tr in 
        let inputs_in_prefix = List.map (fun (_, inp, _, _) -> inp) tr_prefix in 

        (* The sequence we expect, of the password being input *)
        let expected_input = List.map (fun x -> PressNum x) pw in 
        inputs_in_prefix = expected_input
    else true

Note that the trace leaves the state type of the HSM underspecified (as just a type variable), since we don't know what that is yet. It doesn't matter, because our trace property only inspects the inputs and outputs; not the state.

Now, we can implement the HSM. Below is our type of state: 

type hsm_state = {
  secret : string;
  password : int list;
  inputs_so_far : int list;
  curr_display : hsm_output;
  finished : bool
}
It contains: the secret; the password; the inputs pressed so far; the current display on the screen; and whether we are accepting more inputs.

Now, we are ready for the state machine implementation. 

let step_hsm (s : hsm_state) (i : hsm_input) : hsm_state * hsm_output = 
          if s.finished then (s, s.curr_display)
          else 
            match i with 
            | PressNum x ->
                let s' = { s with inputs_so_far = s.inputs_so_far @ [x] } in 
                (s', s'.curr_display)
            | PressUnlock -> 
              if s.inputs_so_far = s.password then 
                let s' = {s with curr_display = UNLOCKED s.secret; finished = true } in 
                (s', s'.curr_display)
              else 
                (* BUG: don't set finished = true *)
                (* This will allow the user to keep querying after they try. *)
                let s' = {s with curr_display = ERROR; finished = true } in 
                (s', s'.curr_display)

If s.finished is true, we don't change the state, and output whatever the state says it should. Otherwise, we match on i. If i is PressNum x, we add x to s.inputs_so_far. If i is PressUnlock, we see if we have input the password in s.inputs_so_far. If we have, then we set the display to UNLOCKED s.secret, and set s.finished = true. Otherwise, the input is wrong, so we set s.curr_display to ERROR, and again set s.finished to true.

Finally, we now can define our machine. Our machine's state expects to be filled in with a specific secret and password, so we define a function of type string -> int list -> (hsm_state, hsm_input, hsm_output) machine: 

let hsm my_secret my_pw : (hsm_state, hsm_input, hsm_output) machine = {
      init = { secret = my_secret; password = my_pw; inputs_so_far = []; curr_display = Ok; finished = false };
      step = step_hsm
}

Testing the HSM

Let's now test the HSM by setting up some tests that use hsm_trace_correct.

(* Given a secret, a password, and some inputs, see if the trace generated by the HSM is valid. *)
let test_hsm_trace_correct secret pw inputs = 
  hsm_trace_correct secret pw (make_trace (hsm secret pw) inputs)

(* Randomly generate inputs. *)
let gen_input () = 
    if Random.bool () then 
        PressUnlock
  else PressNum (Random.int 5)

let gen_inputs () =  
    List.init 5 (fun _ -> gen_input ())

(* Use `forall` to test if the HSM is correct, when given a secret "secret" and a password [0; 1; 2]. *)
let test_property () = 
    forall gen_inputs (fun ins -> test_prop "secret" [0; 1; 2] ins) 5000