Safely Deriving Unsafe Traits

The curious case of safe-discriminant

2024-08-22

Table of Contents

Prelude #

Imagine you are tasked with designing a datatype for requests. For this discussion, let us consider the following operations: Login, Register, Logout, and SendData.

 #[repr(u8)]
 pub enum Request {
     Login {
         username: String,
         password: String,
     } = 0x12,
     Logout = 0x13,
     Register {
         username: String,
         password: String,
         name: String,
     } = 0x14,
     SendData {
         data: Vec<u8>,
     } = 0x15,
 }

One way to send and receive this data structure is to use a single byte to represent the discriminant (1) (1)In this context, the term discriminant refers to values such as 0x12, 0x13, 0x14, and 0x15. Some may refer to it as a tag, but I prefer to follow Rust’s terminology and use discriminant. . This discriminant is followed by a variable-sized stream of bytes, which depends on the specific variant in use. This representation raises the question: how can we extract the discriminant?

As I write this post, Rust does not have direct support for discriminant extraction(2) (2)This is currently being tracked in RFC #3607. , which means we’re on our own! The good news is that the Rust documentation provides a way to extract the discriminant using unsafe code(3) (3)Accessing the numeric value of the discriminant . For our Request enum, we can do the following:

 impl Request {
     pub fn discriminant(&self) -> u8 {
         unsafe { *<*const _>::from(self).cast::<u8>() }
     }
 }

Assuming the Rust documentation is accurate and this code is semantically safe, it should be pretty straightforward to generalize this into a trait and make a procedural macro that auto-derives the traits, right? right?

The Trait and the Macro #

The trait(4) (4)If you are interested in using it, feel free to check docs.rs/safe-discriminant. has two objectives. It determines the type of the discriminant using an associated type and retrieves its value through a method. This implementation should be relatively straightforward(5) (5)While retrieving the discriminant should always be safe, implementing this trait is not always guaranteed to be safe, which is why we designate it as an unsafe trait. . We already have a method for calculating the discriminant using an unsafe cast, as previously demonstrated:

 pub unsafe trait Discriminant {
     type Repr: Copy;
     fn discriminant(&self) -> Self::Repr {
         unsafe { *<*const _>::from(self).cast() }
     }
 }

The derive macro will verify certain conditions. If those conditions are met, we deem the trait to be semantically safe for implementation. The conditions are as follows:

The macro is applied to an enum for which we are deriving Discriminant.
The enum is annotated with #[repr(X)], where X is a primitive type.
Each variant of the enum has an explicitly defined discriminant.

In other words, in the followin snippet Discriminant will be successfully derived for Good, while both Bad1 and Bad2 will trigger an error in the derive macro.

 #[derive(Discriminant)]
 #[repr(i64)]
 pub enum Good<T> {
     A = 1,
     B(T) = -1,
     C { fst: T, snd: T } = -2,
 }

 #[derive(Discriminant)]
 #[repr(i64)]
 pub enum Bad1<T> {
     A = 1,
     B(T), // note missing discriminant
     C { fst: T, snd: T } = -2,
 }

 #[derive(Discriminant)]
 // #[repr(i64)]  <--- no repr
 pub enum Bad2<T> {
     A = 1,
     B(T) = -1,
     C { fst: T, snd: T } = -2,
 }

The Unsoundness Problem #

A Look at Procedural Macros #

Before diving into soundness issues, it is beneficial to have a look at the finer points of Rust procedural macros (6) (6)This is not intended as a tutorial on how macros work If you are looking for a guide, I recommend The Little Book of Rust Macros which is an excellent resource. . Rust offers three types of procedural macros: function-like macros, derive macros, and attribute macros. Our primary focus will be on derive macros and attribute macros.

The main difference between attribute macros and derive macros is that an attribute macro takes an Abstract Syntax Tree (AST) as input and can change any part of that tree(7) (7)It does not have to, it can always append to that tree. . This means that the input to an attribute macro does not need to follow proper Rust syntax(8) (8)This is usually rare, and I have never seen it. . That said, the input can still be a valid Rust AST. This AST may then be transformed into a different valid Rust syntax. In contrast, derive macros always receive an input that is a valid Rust AST. While these macros can introduce new code, the original input AST remains unchanged. This behaviour makes derive macros particularly well-suited for deriving traits in Rust.

When multiple interleaved derive macros and attribute macros are present, they operate in semi-isolation. The macros are evaluated from the outermost to the innermost, allowing the outermost macro to access the inner attribute macros but not the inner derive macros.

It may seem unusual to consider why this might matter for soundness bugs in unsafe code. However, these subtleties can significantly impact code safety.

The Malicious Attribute #

Recall that one of the safety conditions for #[derive(Discriminant)] is that the enum in question must be annotated with #[repr(X)], where X is a primitive type. We can create a macro that invalidates this condition as follows:

 #[proc_macro_attribute]
 pub fn remove_repr(_: TokenStream, item: TokenStream) -> TokenStream {
     let mut tagged_enum = parse_macro_input!(item as ItemEnum);
     tagged_enum
         .attrs
         .retain(|attr| !attr.path().is_ident("repr"));
     quote! {
         #tagged_enum
     }
     .into()
}

To demonstrate how this macro will behave, consider the following snippet(9) (9)Please note that #[remove_repr] does not correctly handle enums with named or unnamed fields. A proper implementation of this macro must also account for the removal of the discriminant. :

 #[remove_repr]
 #[repr(u8)]
 pub enum Request {
     Login  = 0x12,
     Logout = 0x13,
 }

It will expand to:

 pub enum Request {
     Login  = 0x12,
     Logout = 0x13,
 }

To break the safety condition, we can attempt the following:

 #[derive(Discriminant)]
 #[remove_repr]
 #[repr(u8)]
 pub enum Request {
     Login  = 0x12,
     Logout = 0x13,
 }

Due to the order of macro execution, the derive macro will expand first, resulting in:

 #[remove_repr]
 #[repr(u8)]
 pub enum Request {
     Login  = 0x12,
     Logout = 0x13,
 }
 unsafe impl Discriminant for Request {
     type Repr = u8;
 }

Next, the #[remove_repr] macro will expand, violating the soundness invariant, and we end up with code that looks like this:

 pub enum Request {
     Login  = 0x12,
     Logout = 0x13,
 }
 unsafe impl Discriminant for Request {
     type Repr = u8;
 }

To fix this problem, we need to establish one more safety rule for deriving Discriminant as follows (10) (10)Rules 1-3 did not change. :

The macro is applied to an enum for which we are deriving Discriminant.
The enum is annotated with #[repr(X)], where X is a primitive type.
Each variant of the enum has an explicitly defined discriminant.
Aside from #[repr(X)], there are no #[attr] style proc-macros following #[derive(Discriminant)].

Epiloge #

At this point, it is clear that unsafe code can introduce hidden bugs, even when interacting with safe code. This raises the question: how much overhead can we expect from a safe alternative? Referring back to our Request example, a safe solution might look like this:

 impl Request {
     #[inline(never)]
     pub fn discriminant(&self) -> u8 {
         match &self {
             Request::Login {username, password} => 0x12,
             Request::Logout => 0x13,
             Request::Register {
                 username,
                 password,
                 name,
             } => 0x14,
             Request::SendData { data } => 0x15,
         }
     }
 }

The use of #[inline(never)] alludes to the possibility that this function may be optimized out. To analyze the generated assembly thoroughly, we prefer it to be a separate function. When compiled with rustc-1.80.0 using the -C opt-level=3 flag for the x86_64 architecture, we obtain the following assembly:

 discriminant:
    movzx  eax,BYTE PTR [rdi]
    ret

The output is precisely what you would expect from the unsafe trait and macro, yet it is simpler to reason about. However, the unsafe trait is ultimately better(11) (11)Objectively, if we can establish a test suite that verifies the match statement is optimized to a mov instruction, then the match approach is superior because it requires no unsafe code. , especially since I cannot write a post about the match approach.