Non-affine types, ManuallyDrop and Invariant Lifetimes in Rust - Part One

Part One: Exploring non-affine types and invariant lifetimes in Rust through writing the sosecrets-rs crate.

Exploring non-affine types in Rust and its involvement in writing a crate for secrets management.

Motivation

A couple of years ago, one of my acquaintances commented that anyone can ’log’ secrets (without any sinister intentions) retrieved from the ‘vault’ that kept secrets in their programs. That comment stayed in my mind for years.

I kept wondering if we could write more secure codes that manage our secrets better until Rust offers an opportunity to do so.

It is the only production-ready programming language that has move semantics as its default and one that has a type system expressive enough to encode rank 2 polymorphism.

I believe these are some of the most important reasons behind its mass adoption in the blockchain industry, powering several smart-contract (domain-specific) languages for several layer-one blockchains 1 2.

Since late 2023, I have been working on the sosecrets-rs crate, which aims to take a step forward in enabling engineers to write more secure codes around secrets management.

Design Goals for sosecrets-rs

There is a type named Secret<T>, it should wrap any T and have two associated parameters (either as type parameters or fields) - one representing the current number of times this secret is revealed (because it is a legitimate need to reveal secrets in the execution of a program) and the other representing the maximum allowed number of times this secret can be revealed.

It should be a compile-time error if the secret is exposed more than what it is allowed to.

An example:

use sosecrets_rs::{
  prelude::*,
  traits::ExposeSecret,
};
use typenum::U2;

// Define a secret with a maximum exposure count of 2, specified by `typenum::U2`
let secret = Secret::<_, U2>::new("my_secret_value".to_string());

// Expose the secret and perform some operations with the exposed value; 
// the secret has been exposed once: `EC` = 1, `MEC` = 2;
let (next_secret, exposed_value) = secret.expose_secret(|exposed_secret| {
    assert_eq!(&*exposed_secret.as_str(), "my_secret_value");
    // Perform operations with the exposed value
});

// Expose the secret again and perform some operations with the exposed value;
// the secret has been exposed twice: `EC` = 2, `MEC` = 2;
let (next_secret, exposed_value) = next_secret.expose_secret(|exposed_secret| {
    assert_eq!(&*exposed_secret.as_str(), "my_secret_value");
    // Perform operations with the exposed value
    // ...
});

// **Try** to expose the secret again and perform some operations with the exposed value;
// `EC` = 3, `MEC` = 2;
// The following when uncommented is uncompilable.
// let (next_secret, exposed_value) = next_secret.expose_secret(|exposed_secret| {
//     assert_eq!(&*exposed_secret.as_str(), "my_secret_value");
//     // Perform operations with the exposed value
//     // ...
// });

There should be a lexical scope within which the secret is revealed. The secret when revealed (therein known as exposed secret) should not be leaked.

An example (the following should NOT compile):

fn main() {
  use sosecrets_rs::{prelude::*, traits::ExposeSecret};
  use typenum::consts::U2;

  struct A {
      inner: i32,
  }

  let new_secret: Secret<_, U2> = Secret::new(A { inner: 69 });

  // An attempt to return `exposed_secret` of type `ExposedSecret<'brand, A>`
  let (new_secret, _) = new_secret.expose_secret(|exposed_secret| exposed_secret);
  // `ExposedSecret<'brand, A>` value can only be dereferenced, even so, it cannot be returned out of the
  // scope of the closure (i.e. `|exposed_secret| *exposed_secret`)
  let (_, _) = new_secret.expose_secret(|exposed_secret| *exposed_secret);
}

This goal is not trivial, consider the following:

use std::ops::Deref;

pub struct Secret<T>(T);

pub struct ExposedSecret<'t, T>(&'t T);

impl<'t, T> Deref for ExposedSecret<'t, T> {
    type Target = T;

    fn deref(&self) -> &Self::Target {
        self.0
    }
}

impl<T> Secret<T> {
    pub fn new(v: T) -> Self {
        Self(v)
    }

    pub fn expose_secret<'u, 't, ReturnType>(
        &'u self,
        closure: fn(ExposedSecret<'t, T>) -> ReturnType,
    ) -> ReturnType
    where
        'u: 't,
    {
        closure(ExposedSecret(&self.0))
    }
}

fn main() {
    #[derive(Debug)]
    struct BigSecret {
        small_secret: i32,
    }

    let secret = Secret::new(BigSecret { small_secret: 69 });

    let exposed_secret = secret.expose_secret(|exposed_secret| exposed_secret.0);

    println!("{:?}", exposed_secret); // Outputs: BigSecret { small_secret: 69 }
}

The above will output BigSecret { small_secret: 69 }, allowing the ExposedSecret<'t, T> to ’escape’ out of the closure fn(ExposedSecret<'t, T>) -> ReturnType. Playground

Other goals can be:

Secrets are serializable and preserve their MEC and EC type, field information.

Do the checking at runtime instead of compile time.

In this article, let’s explore how non-affine types enabled the first design goal to be met.

Affine Type

Simply put, by default, most of Rust’s types are affine types 3 4; that means, the value of such type can only be used at most once - once you have created a value of an affine type, you can choose to not use it (i.e. use it zero times) or use it once and lose the ‘right’ to use it the second time. In Rust’s parlance, a move-only type, i.e. almost all non-primitive types (specifically, those that do not implement the Copy trait) is an affine type and its value can only be owned by one named variable at any one time. In other words, it is the destructive move semantics. For more in-depth discussion on different substructural type systems (e.g. affine, linear, etc.), you can consult the book, Advanced Topics in Types and Programming Languages.

Secret

The main public type in sosecrets-rs is Secret<T, MEC: typenum::Unsigned, EC: typenum::Unsigned>. It is not an affine type - its value can be used at most MEC (an unsigned positive integer, e.g. 5, 100, etc.) times.

sosecrets-rs makes use of the typenum crate, a crate that provides type-level representations of integers, to raise a compile error only if the secret is exposed (via the EC type parameter) more than what it is maximally allowed to (via the MEC type parameter). This is achieved by imposing additional trait bounds (IsLessOrEqual) on the types that implement the Unsigned trait.

sosecrets-rs crate exposes a trait named ExposeSecret under its traits module. It is not included in prelude, where Secret<T, MEC, EC> resides because users of this crate should have a choice whether they want the secrets to being ’exposed’ or not.

In pseudo-Rust, ExposeSecret’s fn expose_secret<ReturnType, ClosureType>(self, scope: ClosureType) -> (Self::Next, ReturnType) (NOTE: expose_secret(...) function signature is NOT valid Rust) consumes the current secret and gives the caller back a new Secret with its EC type parameter ‘incremented’.

Let’s see a watered-down implementation of ExposeSecret trait for Secret<T, MEC, EC>:

impl<
      'max,
      T,
      MEC: Unsigned,
      EC: Add<U1> + Unsigned + IsLessOrEqual<MEC, Output = True>,
  > ExposeSecret<'max, &'max T, MEC, EC> for Secret<T, MEC, EC>
{
  type Exposed<'brand> = ExposedSecret<'brand, &'brand T>
  where
      'max: 'brand;

  type Next = Secret<T, MEC, Sum<EC, U1>>
  where
      EC: Add<U1> + Unsigned + IsLessOrEqual<MEC, Output = True>,
      Sum<EC, U1>: Unsigned + IsLessOrEqual<MEC, Output = True> + Add<U1>;

  fn expose_secret<ReturnType, ClosureType>(
      mut self,
      scope: ClosureType,
  ) -> (Secret<T, MEC, Sum<EC, U1>>, ReturnType)
  where
      Sum<EC, U1>: Add<U1> + Unsigned + IsLessOrEqual<MEC, Output = True>,
      for<'brand> ClosureType: FnOnce(ExposedSecret<'brand, &'brand T>) -> ReturnType,
  {
      let returned_value = scope(ExposedSecret(&self.0, PhantomData));
      let inner = ManuallyDrop::new(unsafe { ManuallyDrop::take(&mut self.0) });
      forget(self);
      (Secret(inner, PhantomData), returned_value)
  }
}

The following code snippet on the associated type Exposed in the trait implementation will be explained in the next post of this series when I talk about invariant lifetimes.

type Exposed<'brand> = ExposedSecret<'brand, &'brand T>
  where
      'max: 'brand;

The Next associated type in the trait definition of ExposeSecret is set to the following in the implementation of the trait for Secret<T, MEC, EC>:

type Next = Secret<T, MEC, Sum<EC, U1>>
  where
    EC: Add<U1> + Unsigned + IsLessOrEqual<MEC, Output = True>,
    Sum<EC, U1>: Unsigned + IsLessOrEqual<MEC, Output = True> + Add<U1>;

The first bound EC: Add<U1> + Unsigned + IsLessOrEqual<MEC, Output = True> means the following:

EC must be ‘addable’ to type U1 (typenum’s type to represent a type level unsigned integer of value 1) AND
EC must be Unsigned and most importantly,
EC must be less than or equal to MEC.

The IsLessOrEqual is a trait (or, in this context, it is a meta-function, in the C++-sense of the word, or just a type-level function) that takes in a generic type parameter. The associated type Output of the trait (IsLessOrEqual) implementation for EC will either be the typenum’s True type if it is indeed ’lesser or equal’ than MEC or False otherwise.

Now, the resulting type of ‘adding’ EC to U1 must satisfy the same bounds again so you can keep adding U1 to EC and check if it is less than MEC.

Now let’s look at the definition of the only required method expose_secret of ExposeSecret trait.

fn expose_secret<ReturnType, ClosureType>(
    mut self,
    scope: ClosureType,
) -> (Secret<T, MEC, Sum<EC, U1>>, ReturnType)
where
    Sum<EC, U1>: Add<U1> + Unsigned + IsLessOrEqual<MEC, Output = True>,
    for<'brand> ClosureType: FnOnce(ExposedSecret<'brand, &'brand T>) -> ReturnType,
{
    let returned_value = scope(ExposedSecret(&self.0, PhantomData));
    let inner = ManuallyDrop::new(unsafe { ManuallyDrop::take(&mut self.0) });
    forget(self);
    (Secret(inner, PhantomData), returned_value)
}

The return type of expose_secret is a tuple with its first element as the type Secret<T, MEC, Sum<EC, U1>>, the bound Sum<EC, U1>: Add<U1> + Unsigned + IsLessOrEqual<MEC, Output = True> is there to specify the bounds specified by the associated type Next in the definition of the trait ExposeSecret.

Whenever the bounds are not satisfied, e.g. calling the method expose_secret(|exposed_secret| {}); of a value of type Secret<i32, U5, U5> will raise a compile error because e.g. Sum<U5, U1> did not satisfy the bound IsLessOrEqual<U5, Output = True>. This is how the check is done and how the compile error is raised whenever the number of exposures exceeds the maximum allowed number.

Summary

Both the counting of exposure times at the type level and the representation of unsigned integers at the type level are enabled by the typenum crate.
The checks that raise compile error when they fail are caused by the resultant type parameters (i.e. Sum<EC, U1>) not satisfying specific trait bounds.

In my next post, I will share how the second design goal is implemented and what each of these lines in the implementation of expose_secret(...) means. These will involve concepts like invariant lifetimes, rank 2 polymorphism, and how to ‘drop’ a ManuallyDrop<T>.

{
  let returned_value = scope(ExposedSecret(&self.0, PhantomData));
  let inner = ManuallyDrop::new(unsafe { ManuallyDrop::take(&mut self.0) });
  forget(self);
  (Secret(inner, PhantomData), returned_value)
}