Unsafe Code Review: Casting &mut [u32;32] to &mut [GenericArray<u8, U16>; 8]

Hey all,
I'm writing some cryptographic code where I'm implementing BlockRng by using Aes128 in a CTR-like mode. The implementation is on the playground, but not compilable, because of missing crates.

The only unsafe usage in that code is:

impl BlockRngCore for AesRngCore {
    type Item = u32;
    // This is equivalent to `[Block; 8]`, but we need to use `u32` to be
    // compatible with `RngCore`.
    type Results = [u32; 32];

    #[inline]
    fn generate(&mut self, results: &mut Self::Results) {
        //        here 👇
        let blocks = unsafe { &mut *(results as *mut _ as *mut [GenericArray<u8, U16>; 8]) };
        blocks.into_iter().for_each(|blk| {
            *blk = self.state.to_le_bytes().into();
            self.state += 1;
        });
        self.aes.encrypt_blocks(blocks);
    }
}

a minimal example of the code in question is:

use generic_array::GenericArray; // 0.14.5
use typenum::U16; // 1.15.0

fn main() {
    let results = &mut [0_u32; 32];
    let blocks = unsafe { &mut *(results as *mut _ as *mut [GenericArray<u8, U16>; 8]) }; 
    dbg!(blocks);
}

(playground, Miri at least is not complaining)

The GenericArray is #[repr(transparent)] and contains a single field that has the type ArrayLength<T>::ArrayType. However, I don't quite understand the recursive nature of the ArrayLength implementation.
But I figure, as the generic_array crate is transmuting [T;N] to GenericArray<T, N> in its from implementation, they should have the same representation.
I'd really appreciate it, if someone more knowledgeable in unsafe Rust could have a look at this and hopefully tell me that the code is correct. I'm also open to ideas of how to do the above without resorting to unsafe, but no overhead. The code is in a critical path and my previous version of stack allocating a new block buffer, encrypting it, and then writing the results performed measurably worse.

We don't really have enough details about the memory layout of GenericArray<u8, U16> to decide if your code is sound or not.

The reason is that you're casting a pointer to an array of u32 to a pointer to an array of GenericArray<u8, U16>. By doing so you rely on the types having the same memory layout. Right now, structs (repr(rust) ones, that is) have no guarantees about their mem-layout, so that's gonna be a problem.

You can get around this by manually specifying the memory layout of GenericArray. As long as you set it to have the same alignment as u32 and make sure your struct has the same size as well, this cast is sound (if the arrays are also the same length), since [u32] has a well-defined memory layout.

But in your specific example you're using a struct with generic types, which is only supported in repr(rust) and therefore has no well-defined memory layout. Also, you're changing the number of elements (so you're casting between arrays with differently sized entries), which may also cause gibberish due to alignment issues.

It will probably work though :person_shrugging:, just don't count on it: the compiler might feel lucky and shuffle around the fields in your struct a lil' s.t. they have different padding and kaboom all your data is now gibberish.

GenericArray is declared as #[repr(transparent)]:

/// Struct representing a generic array - `GenericArray<T, N>` works like [T; N]
#[allow(dead_code)]
#[repr(transparent)]
pub struct GenericArray<T, U: ArrayLength<T>> {
    data: U::ArrayType,
}

And both of the underlying ArrayTypes are declared as #[repr(C)]:

/// Internal type used to generate a struct of appropriate size
#[allow(dead_code)]
#[repr(C)]
#[doc(hidden)]
pub struct GenericArrayImplEven<T, U> {
    parent1: U,
    parent2: U,
    _marker: PhantomData<T>,
}

/* ... */

/// Internal type used to generate a struct of appropriate size
#[allow(dead_code)]
#[repr(C)]
#[doc(hidden)]
pub struct GenericArrayImplOdd<T, U> {
    parent1: U,
    parent2: U,
    data: T,
}

So this pointer cast should always be safe. Note that this must work due to the existence of GenericArray::from_[mut_]slice().

1 Like

Looks good to me. Just note that the endianness of u32 is implementation-defined, so be careful how you use it (you can end up with reversed bytes). For cryptographic code it likely doesn't matter, since it usually just uses u32 or u64 to manipulate several bytes at the time and to use complex operations on bytes (like addition with carrying).

With regards to the implementation of GenericArray, it is intended to be the same as [T; N] with some const generic parameter N. However, at the moment GenericArray was written, const generics didn't exist. Even now they don't support many things, like expressing that concatenation of arrays of length N and M is an array of length N+M. For this reason it is implemented via the trait machinery, which is Turing-complete, and so can express arbitrary compile-time computations (including the lengths of all relevant arrays).

You already know that GenericArray<T, U> has the same layout as <U as ArrayLength<T>>::ArrayType. It is intended that U::ArrayType is the array of type T and length encoded in U. However, we can't directly express such arrays in stable Rust. What we can do is provide an object which has the same layout, i.e. N consecutive copies of T without any padding between them.

The way to do so is via #[repr(C)] structs. Such structs are laid out in memory just like in C language, which means that the fields are laid out in the order of their in-source declaration, with the minimum possible padding required to make all fields aligned. We have two such structs: GenericArrayImplEven and GenericArrayImplOdd, which are defined as

#[allow(dead_code)]
#[repr(C)]
#[doc(hidden)]
pub struct GenericArrayImplOdd<T, U> {
    parent1: U,
    parent2: U,
    data: T,
}

#[allow(dead_code)]
#[repr(C)]
#[doc(hidden)]
pub struct GenericArrayImplEven<T, U> {
    parent1: U,
    parent2: U,
    _marker: PhantomData<T>,
}

This means that GenericArrayImplEven<T, U> has the same alignment as U and its fields are layed out contiguously, while GenericArrayImplOdd<T, U> has the maximum alignment of U and T. If U has the same alignment as T, then GenericArrayImplOdd<T, U> has the same alignment as T, and its fields are layed out consecutively.

Now, the final part of the puzzle is the representation of integers. Natural numbers are represented as singly-linked lists of individual bits, where the bits 0 and 1 are represented as structures typenum::bit::{B0, B1}. The bits of the number are layed out in little-endian order, i.e. the least significant bit is always the first element of the linked list. The inductive definition of ArrayLength<T> consists of the base case and two cases corresponding to individual bits.

unsafe impl<T> ArrayLength<T> for UTerm {
    #[doc(hidden)]
    type ArrayType = [T; 0];
}

unsafe impl<T, N: ArrayLength<T>> ArrayLength<T> for UInt<N, B0> {
    #[doc(hidden)]
    type ArrayType = GenericArrayImplEven<T, N::ArrayType>;
}

unsafe impl<T, N: ArrayLength<T>> ArrayLength<T> for UInt<N, B1> {
    #[doc(hidden)]
    type ArrayType = GenericArrayImplOdd<T, N::ArrayType>;
}

This means that for N == 0 we have N::ArrayLength = [T; 0], for N == 2*K we have N::ArrayLength = GenericArrayImplEven<T, N>, which is layout-equivalent to [K::ArrayLength; 2], and for N == 2*K+1 we have N::ArrayLength = GenericArrayImplOddT, N>, which is the same as contiguously laid out GenericArrayImplEven<T, N> == [K::ArrayLength; 2] and T. You can see that what we have is essentially nested arrays of T with an extra T element attached for odd bits, which means that GenericArray<T, N> has the same layout as [T, N] (where I have identified a compile-time constant N with the linked list of types representing it).

3 Likes

yep that should work, I didn't know about the impl details of GenericArray.

Pretty clever trick to wrap repr(C) structs in a transparent container with a generic type, cool stuff!

Ohhh, that is really neat! I looked at the ArrayLength impls and the Odd and Even structs and noticed their #[repr(C)] but didn't quite understand the inductive definition. What I missed was the encoding of the length as a binary number represented by the recursive types. I mistakenly assumed that typenum's UInt works in a peano kind of way, with one level of nesting for each number. But doing it this way is probably a lot nicer for the compiler :joy:

Regarding the endianness, I briefly thought of it, but apparently not enough. The AesRng is part of a cryptographic protocol executed by two parties which sometimes seed the Rng with the same value, expecting the same output. If I'm simply casting the state counter, the input to the encryption and therefore the output of the rng will differ, breaking the correctness. Thanks for reminding me of that! (For anyone curious, I'm implementing different oblivious transfer protocols).

Thank you very much for your detailed answer :blush: And also thanks to the other people.