Typestate Pattern

// #Rust #Craft

The handshake at the beginning of a BIP-324 connection is a non-trivial dance between the two peers. This isn’t a simple syn/ack. There is secret material to exchange, two forms of traffic shape hiding (garbage and decoys), and protocol version negotiation (although in BIP-324’s current form this is pretty much a no-op). And half way through the handshake, things go encrypted and the ciphers need to fire up and stay in-sync with each other so the rest of the session works.

This is a fair bit of code to implement and I ended up breaking it out into its own module. But more concerning than the amount of code, is the ominous error variant of the crate, HandshakeOutOfOrder. The Handshake type has methods which should be called in a specific order as the handshake progresses. They can’t be called all at the same time because the caller has to exchange some data with the remote during the process. But this implementation opens the door for the caller to do the handshake out of order (e.g. look for garbage when they should be sending a public key). This is detected based on the state of the handshake and a runtime error is raised.

Ideally, this is detected at compile time, not runtime. I attempted to accomplish this by introducing what I have learned is called the typestate pattern. Essentially, I give the type system enough information that it can enforce the handshake methods are called in the correct order.

/// **Initial state** of the handshake state machine which holds local secret materials.
pub struct Initialized {
    point: EcdhPoint,
}

/// **Second state** after sending the local public key.
pub struct SentKey {
    point: EcdhPoint,
    bytes_written: usize,
}

/// **Third state** after receiving the remote's public key and
/// generating the shared secret materials for the session.
pub struct ReceivedKey {
    session_keys: SessionKeyMaterial,
}

/// **Fourth state** after sending the version packet.
pub struct SentVersion {
    cipher: CipherSession,
    remote_garbage_terminator: [u8; NUM_GARBAGE_TERMINTOR_BYTES],
    bytes_written: usize,
    ciphertext_index: usize,
    remote_garbage_authenticated: bool,
}

pub struct Handshake<'a, State> {
    /// Bitcoin network both peers are operating on.
    network: Network,
    /// Local role in the handshake, initiator or responder.
    role: Role,
    /// Optional garbage bytes to send along in handshake.
    /// Optional local garbage bytes to send along in handshake.
    garbage: Option<&'a [u8]>,
    /// State-specific data.
    state: State,
}

// Initialized state implementation
impl<'a> Handshake<'a, Initialized> {
    pub fn send_key(
        mut self,
        garbage: Option<&'a [u8]>,
        output_buffer: &mut [u8],
    ) -> Result<Handshake<'a, SentKey>, Error> { ... }
}

Four handshake states are defined and the Handshake type is made generic across them. State specific impls are left off, but you get the idea.

The only way for a caller to get a Handshake<Initialized> into the next state, Handshake<SentKey>, is to call Handshake<Initialized>::send_key(mut self). send_key is only defined on the Initialized type and it consumes itself, so you are left with just one state. The path is now enforced by the types.

One thing I am still curious about is if the State generic should be bound by a trait which the state types implement. Maybe even just a marker trait. This might make it more difficult for a caller to do something weird like try to create a Handshake<String>, where String is not one of the four valid states. But since the fields of Handshake are private and the only public constructor returns a valid state, I don’t think there is a way for a caller to pull this off. However, if a constructor accepted a State (or if the fields were public), it might be best to introduce a sealed trait (cannot be implemented outside of its own crate) for the states.

// A public marker trait whose name is not publicly exported.
mod sealed {
    pub trait Sealed {}
}

// The Sealed trait is public, which is must be since it's in an exposed API of
// the module. But its *name* is not exposed. It can't be named by a caller,
// so they can't define their own HandshakeState.
pub trait HandshakeState: sealed::Sealed {}

// Boilerplate.
impl sealed::Sealed for Initialized {}
impl sealed::Sealed for SentKey {}
impl sealed::Sealed for ReceivedKey {}
impl sealed::Sealed for SentVersion {}

impl HandshakeState for Initialized {}
impl HandshakeState for SentKey {}
impl HandshakeState for ReceivedKey {}
impl HandshakeState for SentVersion {}

// Prevents invalid states at the type level.
pub struct Handshake<'a, State: HandshakeState> {}

The traditional two-step sealed trait pattern. Sealed traits impls have subtle trade-offs, so I think I’ll avoid if possible.

Each of the four handshake states has state-specific data, which is kinda nice because then I avoid the weirdness of PhantomData. What if each of the handshake state types was essentially just a marker type (aka phantom type or zero sized type) with no data itself? Technically, this type will only mean something at compile time, at runtime it is just nothing. So can we just stick it in the handshake usage like so Handshake<Initialized> and have it be handled by the compiler? Turns out, no, if the compiler sees a generic parameter it needs to actually be used somewhere in the type. There is probably a good reason for this. I believe just keeping the phantom type as a field in the handshake type, state: State, would work. Although the compiler will get pissed that the field is never accessed, so have to slap a _ to the front. But to take it one step further, you can also use _state: PhantomData<State> so the semantics are clear, this thing is just a marker and not used on purpose.

One last thing I found interesting is that this is kinda the inverse of the enum static dispatch pattern I explored in the Static Dispatch log. In that scenario, I am attempting to hide from the caller the generic implementation type. They just care that something is a Connection, not how it is implemented under the hood. But with the handshake its the opposite, I want the caller to be exposed to which generic state type they are working with.