Rust’s const fn feature is a game-changer for cryptography. I’ve been exploring how to use it for compile-time key expansion, and the results are impressive. Let’s dig into this powerful technique.

First, what exactly is const fn? It’s a way to define functions that can be evaluated at compile-time. This means we can do complex calculations before our program even runs, baking the results directly into the binary.

For cryptography, this is huge. Many encryption algorithms require an initial key to be “expanded” into a series of round keys. Traditionally, this happens when the program starts up. But with const fn, we can do it at compile-time instead.

Here’s a simple example to illustrate the concept:

const fn expand_key(key: u32) -> [u32; 4] {
    let mut expanded = [0; 4];
    expanded[0] = key;
    expanded[1] = key.rotate_left(8);
    expanded[2] = key.rotate_left(16);
    expanded[3] = key.rotate_left(24);
    expanded
}

const ROUND_KEYS: [u32; 4] = expand_key(0x12345678);

In this code, we’re taking a 32-bit key and expanding it into four round keys. The expand_key function is marked as const, so it runs at compile-time. The ROUND_KEYS constant will be baked into our binary with the pre-computed values.

Now, let’s look at a more realistic example using AES (Advanced Encryption Standard). AES key expansion is more complex, involving substitution boxes (S-boxes) and round constants. Here’s how we might implement it:

const fn sub_word(word: u32) -> u32 {
    // S-box substitution implementation
    // (simplified for brevity)
    word ^ 0x63636363
}

const fn rot_word(word: u32) -> u32 {
    word.rotate_left(8)
}

const fn expand_aes_key(key: &[u8; 16]) -> [u32; 44] {
    let mut expanded = [0u32; 44];
    let mut i = 0;

    // Initial round key
    while i < 4 {
        expanded[i] = u32::from_be_bytes([key[4*i], key[4*i+1], key[4*i+2], key[4*i+3]]);
        i += 1;
    }

    while i < 44 {
        let mut temp = expanded[i-1];
        if i % 4 == 0 {
            temp = sub_word(rot_word(temp)) ^ (1u32 << (i/4 - 1));
        }
        expanded[i] = expanded[i-4] ^ temp;
        i += 1;
    }

    expanded
}

const AES_ROUND_KEYS: [u32; 44] = expand_aes_key(&[0x2b, 0x7e, 0x15, 0x16, 0x28, 0xae, 0xd2, 0xa6,
                                                   0xab, 0xf7, 0x15, 0x88, 0x09, 0xcf, 0x4f, 0x3c]);

This implementation expands a 128-bit AES key into 44 32-bit words, which are used for the 10 rounds of encryption. By making this const, we’re doing all this work at compile-time.

One challenge with const fn is that we’re limited in what operations we can perform. For example, we can’t use standard library functions that aren’t themselves const. This means we often need to reimplement basic operations.

Let’s look at how we might implement a const-compatible S-box lookup:

const fn create_sbox() -> [u8; 256] {
    let mut sbox = [0u8; 256];
    let mut p = 1u8;
    let mut q = 1u8;

    loop {
        // Multiply p by 3 in GF(2^8)
        p = p ^ (p << 1) ^ (if p & 0x80 != 0 { 0x1B } else { 0 });

        // Divide q by 3 in GF(2^8)
        q ^= q << 1;
        q ^= q << 2;
        q ^= q << 4;
        q ^= if q & 0x80 != 0 { 0x09 } else { 0 };

        // Compute the affine transformation
        let affine = q ^ q.rotate_left(1) ^ q.rotate_left(2) ^ q.rotate_left(3) ^ q.rotate_left(4) ^ 0x63;

        sbox[p as usize] = affine;

        if p == 1 { break; }
    }

    sbox[0] = 0x63;
    sbox
}

const SBOX: [u8; 256] = create_sbox();

This creates the AES S-box at compile-time. We can then use this constant SBOX in our encryption functions.

The benefits of this approach are significant. First, there’s zero runtime overhead for key expansion or S-box creation. This is especially valuable in resource-constrained environments like embedded systems.

Second, it can improve security. Since the expanded keys are baked into the binary, they’re not sitting in mutable memory where they might be vulnerable to certain types of attacks.

However, there are also some drawbacks to consider. Compile times can increase significantly, especially for more complex algorithms. Also, if you need to support multiple keys, you’ll need to expand each one at compile-time, which could bloat your binary.

Let’s look at how we might use our compile-time expanded keys in an actual encryption function:

fn aes_encrypt_block(input: &[u8; 16], round_keys: &[u32; 44]) -> [u8; 16] {
    let mut state = [0u32; 4];
    for i in 0..4 {
        state[i] = u32::from_be_bytes([input[4*i], input[4*i+1], input[4*i+2], input[4*i+3]]);
    }

    // Initial round
    for i in 0..4 {
        state[i] ^= round_keys[i];
    }

    // Main rounds
    for round in 1..10 {
        // Sub Bytes
        for i in 0..4 {
            let mut s = state[i].to_be_bytes();
            for j in 0..4 {
                s[j] = SBOX[s[j] as usize];
            }
            state[i] = u32::from_be_bytes(s);
        }

        // Shift Rows
        state = [
            state[0],
            state[1].rotate_right(8),
            state[2].rotate_right(16),
            state[3].rotate_right(24),
        ];

        // Mix Columns
        // (implementation omitted for brevity)

        // Add Round Key
        for i in 0..4 {
            state[i] ^= round_keys[4*round + i];
        }
    }

    // Final round
    // (similar to main round, but without Mix Columns)

    let mut output = [0u8; 16];
    for i in 0..4 {
        let bytes = state[i].to_be_bytes();
        output[4*i..4*i+4].copy_from_slice(&bytes);
    }

    output
}

// Usage
let input = [0x32, 0x43, 0xf6, 0xa8, 0x88, 0x5a, 0x30, 0x8d,
             0x31, 0x31, 0x98, 0xa2, 0xe0, 0x37, 0x07, 0x34];
let encrypted = aes_encrypt_block(&input, &AES_ROUND_KEYS);

This function uses our pre-computed round keys and S-box. The encryption itself still happens at runtime, but all the key expansion and S-box computation was done at compile-time.

The const fn feature isn’t just useful for symmetric encryption. We can also use it for other cryptographic operations. For example, we could compute lookup tables for elliptic curve operations, or pre-compute values for certain hash functions.

Here’s a simple example of using const fn for a hash function’s initial values:

const fn rotate_left(x: u32, n: u32) -> u32 {
    (x << n) | (x >> (32 - n))
}

const fn initial_hash_values() -> [u32; 8] {
    let mut h = [0u32; 8];
    let primes = [2, 3, 5, 7, 11, 13, 17, 19];

    let mut i = 0;
    while i < 8 {
        let mut h_i = (primes[i] as f32).sqrt().fract();
        h_i = (h_i * 2.0f32.powi(32)) as u32;
        h[i] = rotate_left(h_i, i as u32);
        i += 1;
    }

    h
}

const HASH_INIT: [u32; 8] = initial_hash_values();

This computes initial hash values similar to those used in SHA-256, but at compile-time.

As powerful as const fn is, it’s important to use it judiciously. Not everything needs to be or should be computed at compile-time. It’s best suited for values that are truly constant and computationally expensive to calculate at runtime.

Also, remember that const fn is still evolving. Each new version of Rust tends to expand what’s possible in const contexts. Keep an eye on the release notes to see what new capabilities become available.

In conclusion, Rust’s const fn feature opens up exciting possibilities for cryptography. By moving complex computations to compile-time, we can create more efficient and potentially more secure implementations. Whether you’re working on embedded systems, high-performance servers, or anything in between, it’s a technique worth mastering. As always in cryptography, be sure to thoroughly test and audit any implementations before using them in production. Happy coding, and stay secure!