Educational implementation of SHA256 in Rust

Question

A few weeks ago I decided to have a look at the Rust programming language. After a few standard exercises I chose to try to implement the SHA256 hashing algorithm, because why not.

The algorithm's state is stored in a struct, similar to Python's hashlib. This is the lib.rs file:

use std::convert::TryInto;


pub struct SHA256Hash {
    h0: u32,
    h1: u32,
    h2: u32,
    h3: u32,
    h4: u32,
    h5: u32,
    h6: u32,
    h7: u32,
    finalized: bool,
    total_bits: usize,
    unprocessed_bytes: Vec<u8>,
    input_block_size: usize
}


impl SHA256Hash {

    /// Create a new hash instance
    pub fn new() -> SHA256Hash {
        SHA256Hash {
            h0: 0x6a09e667,
            h1: 0xbb67ae85,
            h2: 0x3c6ef372,
            h3: 0xa54ff53a,
            h4: 0x510e527f,
            h5: 0x9b05688c,
            h6: 0x1f83d9ab,
            h7: 0x5be0cd19,
            finalized: false,
            total_bits: 0,
            unprocessed_bytes: Vec::new(),
            input_block_size: 64
        }
    }

    pub fn block_size(&self) -> usize {
        self.input_block_size
    }

    /// update with a block of 512bit
    fn update_block(&mut self, block: &[u8; 64]) {

        // Initialize array of round constants:
        // (first 32 bits of the fractional parts of the cube roots of the first 64 primes 2..311):
        let round_constants: [u32; 64] = [
            0x428a2f98, 0x71374491, 0xb5c0fbcf, 0xe9b5dba5, 0x3956c25b, 0x59f111f1, 0x923f82a4, 0xab1c5ed5,
            0xd807aa98, 0x12835b01, 0x243185be, 0x550c7dc3, 0x72be5d74, 0x80deb1fe, 0x9bdc06a7, 0xc19bf174,
            0xe49b69c1, 0xefbe4786, 0x0fc19dc6, 0x240ca1cc, 0x2de92c6f, 0x4a7484aa, 0x5cb0a9dc, 0x76f988da,
            0x983e5152, 0xa831c66d, 0xb00327c8, 0xbf597fc7, 0xc6e00bf3, 0xd5a79147, 0x06ca6351, 0x14292967,
            0x27b70a85, 0x2e1b2138, 0x4d2c6dfc, 0x53380d13, 0x650a7354, 0x766a0abb, 0x81c2c92e, 0x92722c85,
            0xa2bfe8a1, 0xa81a664b, 0xc24b8b70, 0xc76c51a3, 0xd192e819, 0xd6990624, 0xf40e3585, 0x106aa070,
            0x19a4c116, 0x1e376c08, 0x2748774c, 0x34b0bcb5, 0x391c0cb3, 0x4ed8aa4a, 0x5b9cca4f, 0x682e6ff3,
            0x748f82ee, 0x78a5636f, 0x84c87814, 0x8cc70208, 0x90befffa, 0xa4506ceb, 0xbef9a3f7, 0xc67178f2
        ];

        let mut w: [u32; 64] = [0; 64];

        for j in 0..16 {
            // let mut bytes: [u8; 4] = Default::default();
            // bytes.copy_from_slice(&block[j*4..(j+1)*4]);
            // w[j] = transform_array_of_u8_big_endian_to_u32(&bytes);
            w[j] = transform_array_of_u8_big_endian_to_u32(block[j*4..(j+1)*4].try_into().expect(""));
            // println!("w[{:2}] {:08X?}", j, w[j]);
        }

        // Extend the first 16 words into the remaining 48 words w[16..63] of the message schedule array:
        for j in 16..64 {
            let s0 = w[j-15].rotate_right(7) ^ w[j-15].rotate_right(18) ^ (w[j-15] >> 3);
            let s1 = w[j-2].rotate_right(17) ^ w[j-2].rotate_right(19) ^ (w[j-2] >> 10);
            w[j] = w[j-16].wrapping_add(s0).wrapping_add(w[j-7]).wrapping_add(s1);
        }

        let mut a = self.h0;
        let mut b = self.h1;
        let mut c = self.h2;
        let mut d = self.h3;
        let mut e = self.h4;
        let mut f = self.h5;
        let mut g = self.h6;
        let mut h = self.h7;

        // Compression function main loop:
        for j in 0..64 {
            let s1 = e.rotate_right(6) ^ e.rotate_right(11) ^ e.rotate_right(25);
            let ch = (e & f) ^ (!e & g);
            let temp1 = h.wrapping_add(s1).wrapping_add(ch).wrapping_add(round_constants[j]).wrapping_add(w[j]);

            let s0 = a.rotate_right(2) ^ a.rotate_right(13) ^ a.rotate_right(22);
            let maj = (a & b) ^ (a & c) ^ (b & c);
            let temp2 = s0.wrapping_add(maj);

            h = g;
            g = f;
            f = e;
            e = d.wrapping_add(temp1);
            d = c;
            c = b;
            b = a;
            a = temp1.wrapping_add(temp2);
        }

        self.h0 = self.h0.wrapping_add(a);
        self.h1 = self.h1.wrapping_add(b);
        self.h2 = self.h2.wrapping_add(c);
        self.h3 = self.h3.wrapping_add(d);
        self.h4 = self.h4.wrapping_add(e);
        self.h5 = self.h5.wrapping_add(f);
        self.h6 = self.h6.wrapping_add(g);
        self.h7 = self.h7.wrapping_add(h);
    }

    /// consume blocks of unprocessed bytes
    fn consume(&mut self) {
        assert_eq!(self.unprocessed_bytes.len() % 64, 0);
        let n_blocks = self.unprocessed_bytes.len() / 64;
        for i in 0..n_blocks {
            let mut block: [u8; 64] = [0; 64];
            // the copy seems to be necessary because of a multiple
            // mutable/immutable borowing situation I've set up for myself
            block.copy_from_slice(&self.unprocessed_bytes[i*64..(i+1)*64]);
            self.update_block(&block);
            // it's a pitty that try_into does not work here
        }
        self.unprocessed_bytes.clear();
    }

    /// query the hash digest
    pub fn digest(&mut self) -> [u32; 8] {
        self.finalize();
        [self.h0, self.h1, self.h2, self.h3, self.h4, self.h5, self.h6, self.h7]
    }

    /// query the hex digest, formatted as hexadecimal string
    pub fn hex_digest(&mut self) -> String {
        let digest = self.digest();
        format!(
            "{:08X?}{:08X?}{:08X?}{:08X?}{:08X?}{:08X?}{:08X?}{:08X?}",
            digest[0], digest[1], digest[2], digest[3],
            digest[4], digest[5], digest[6], digest[7]
        )
    }

    /// update the hash state
    pub fn update(&mut self, input: &[u8]) -> bool {
        if self.finalized || input.len() == 0 {
            return false;
        }

        let input_len_bytes = input.len();
        self.total_bits += input_len_bytes * 8;
        let remaining_bytes = (input_len_bytes + self.unprocessed_bytes.len()) % 64;
        self.unprocessed_bytes.extend(input[..(input_len_bytes-remaining_bytes)].iter().clone());

        self.consume();

        if remaining_bytes > 0 {
            self.unprocessed_bytes.extend(input[input_len_bytes-remaining_bytes..].iter().clone());
        }

        true
    }

    fn finalize(&mut self) {
        self.unprocessed_bytes.push(0x80);

        let n_padding_bits = 512 - (self.total_bits + 8 + 64) % 512;
        let n_padding_bytes = n_padding_bits / 8;
        self.unprocessed_bytes.extend(vec![0; n_padding_bytes]);
        let length: u64 = self.total_bits.try_into().unwrap();
        self.unprocessed_bytes.extend(&transform_u64_to_array_of_u8_big_endian(length));

        self.consume();

        self.finalized = true;
    }
}

///////////////////////////////////////////////////////////////////////////////

fn transform_u64_to_array_of_u8_big_endian(x: u64) -> [u8; 8] {
    let b1 = ((x >> 56) & 0xff) as u8;
    let b2 = ((x >> 48) & 0xff) as u8;
    let b3 = ((x >> 40) & 0xff) as u8;
    let b4 = ((x >> 32) & 0xff) as u8;
    let b5 = ((x >> 24) & 0xff) as u8;
    let b6 = ((x >> 16) & 0xff) as u8;
    let b7 = ((x >> 8) & 0xff) as u8;
    let b8 = (x & 0xff) as u8;
    [b1, b2, b3, b4, b5, b6, b7, b8]
}


fn transform_array_of_u8_big_endian_to_u32(arr_of_u8: &[u8; 4]) -> u32 {
    let mut x: u32 = 0;
    x |= (arr_of_u8[0] as u32) << 24;
    x |= (arr_of_u8[1] as u32) << 16;
    x |= (arr_of_u8[2] as u32) << 8;
    x |= arr_of_u8[3] as u32;
    x
}

///////////////////////////////////////////////////////////////////////////////

#[cfg(test)]
mod tests {
    use super::SHA256Hash;

    /// Run empty test input from FIPS 180-2
    #[test]
    fn sha256_nist_empty() {
        let mut hasher = SHA256Hash::new();
        hasher.update(&[]);
        let digest = hasher.digest();
        assert_eq!(
            digest,
            [0xe3b0c442, 0x98fc1c14, 0x9afbf4c8, 0x996fb924,
             0x27ae41e4, 0x649b934c, 0xa495991b, 0x7852b855]
        );
    }

    /// Run abc test from FIPS 180-2
    #[test]
    fn sha256_nist_abc() {
        let mut hasher = SHA256Hash::new();
        hasher.update(b"abc");
        let digest = hasher.digest();
        assert_eq!(
            digest,
            [0xba7816bf, 0x8f01cfea, 0x414140de, 0x5dae2223,
             0xb00361a3, 0x96177a9c, 0xb410ff61, 0xf20015ad]
        )
    }

    /// Run two-block test from FIPS 180-2
    #[test]
    fn sha256_nist_two_blocks() {
        let mut hasher = SHA256Hash::new();
        hasher.update(b"abcdbcdecdefdefgefghfghighijhijkijkljklmklmnlmnomnopnopq");
        let digest = hasher.digest();
        assert_eq!(
            digest,
            [0x248d6a61, 0xd20638b8, 0xe5c02693, 0x0c3e6039,
             0xa33ce459, 0x64ff2167, 0xf6ecedd4, 0x19db06c1]
        );
    }

    /// Run large input test (1,000,000 x a) from FIPS 180-2
    #[test]
    fn sha256_nist_large_input() {
        let input_str = std::iter::repeat("a").take(1_000_000).collect::<String>();
        let input = input_str.as_bytes();
        let mut hasher = SHA256Hash::new();
        hasher.update(&input);
        let digest = hasher.digest();
        assert_eq!(
            digest,
            [0xcdc76e5c, 0x9914fb92, 0x81a1c7e2, 0x84d73e67,
             0xf1809a48, 0xa497200e, 0x046d39cc, 0xc7112cd0]
        );
    }
}

lib.rs comes with a small suite of unit tests as specified in FIPS 180-2 Appendix B. The test suite can be run using cargo test [--release]

Then there is also main.rs which generates an executable to be used a command-line tool:

use std::env;
use std::fs::File;
use std::io::{BufRead, BufReader};

use sha256::SHA256Hash;


fn main() -> std::io::Result<()> {
    let mut hasher = SHA256Hash::new();

    let args: Vec<String> = env::args().collect();
    if args.len() < 2 {
        eprintln!("Please provide a filename as command line argument!");
        return Ok(());
    }

    let filename = &args[1];
    let file = File::open(filename)?;
    let chunk_size: usize = hasher.block_size() * 1024;
    let mut reader = BufReader::with_capacity(chunk_size, file);

    loop {
        let length = {
            let buffer = reader.fill_buf()?;
            hasher.update(buffer);
            buffer.len()
        };
        if length == 0 {
            break;
        }
        reader.consume(length);
    }

    println!("{} {}", hasher.hex_digest().to_ascii_lowercase(), filename);

    Ok(())
}

As an example, let's compute the hash of lib.rs ;-)

cargo run --release src\lib.rs
86dccf89c7cb4de0837a2c4a3c709ae7845151338a1a21a8321fcbf9e4d8dcf7 src\lib.rs

Computing the hash for an ISO image of 3.64 GiB takes roughly 39s here on my laptop (35s with 7zip's built-in SHA256).

I'm open to all kinds of feedback, be it style, idiomatic rust, performance, usability, whatever comes to your mind.

---

Cargo.toml for completeness:

[package]
name = "hash_sha256"
version = "0.1.0"
authors = ["AlexV"]
edition = "2018"

[lib]
name = "sha256"
path = "src/lib.rs"

---

I'm aware of a similar question here on Code Review, but unfortunately it has no answer. Also from what I can judge, the approaches are reasonably different to exist in their own right.

Answer 1

First off, let's start with a simple assumption: SIMD is off the table. As such, your kernel computation (update_block) is mostly optimized, and instead, I ended up focusing on the rest.

Some bad news

After profiling, and as expected, the bulk of the CPU time of your SHA256 implementation is in update_block(). No real surprise there. As a result, we know ahead of time that we can do very little to ease the pains of the 64 rounds of operations without resorting to some devious trickery.

We can, however, make one simple change. We're going to pull in the byteorder crate and use it to convert our u8 arrays into u32s. This will both ease our pain and provide a better version of what you already had. Since we're always feeding it 4 bytes (verifiable by code logic), we can safely drop the slice size prefix and benefit from the blanket &[u8] Read implementation.

It also looks neat and inlinable:

fn transform_array_of_u8_big_endian_to_u32(mut arr_of_u8: &[u8]) -> u32 {
    arr_of_u8.read_u32::<BigEndian>().unwrap()
}

Some good news

There's a ton of allocations we can remove!

We're going to change the method signature of consume from fn consume(&mut self); to fn consume(&mut self, bytes: &[u8]);. The reason for this is the pretty neat performance gain we can score by doing the following (playground):

fn consume(&mut self, mut bytes: &[u8]) {
    let input_len_bytes = bytes.len();
    let unprocessed_len = self.unprocessed_bytes.len();
    self.total_bits += input_len_bytes * 8;
    // Do we have bytes in the unprocessed buffer?
    if unprocessed_len > 0 {
        if (unprocessed_len + input_len_bytes) < 64 {
            // Nothing to do, we just append
            // Copy up to 63 bytes to our Vec
            self.unprocessed_bytes.extend_from_slice(bytes);
            return;
        }
        let (additional, new_bytes) = bytes.split_at(64 - unprocessed_len);
        // Reassign
        bytes = new_bytes;
        // Copy up to 64 bytes from what we just took
        self.unprocessed_bytes.extend_from_slice(additional);
        // We can afford a 64-byte clone
        self.update_block(self.unprocessed_bytes.clone().as_slice());
        self.unprocessed_bytes.clear();
        // Call ourselves
        //return self.inner_consume(new_bytes);
    }
    let iter = bytes.chunks_exact(64);
    let remainder_i = iter.clone();
    for block in iter {
        self.update_block(&block)
    }
    let bytes = remainder_i.remainder();
    self.unprocessed_bytes.extend_from_slice(bytes); // max 64bytes allocated
}

This also opens up repeated calls to update(), something which the previous version fell over on (index out of bounds errors). The chunks_exact() iterator takes advantage that the slice is not mutable and quite literally just moves a pointer around. This both frees us from copies of the data and issues regarding ownership.

Honestly, after having a proper look at profiling data, that's as far as you'll go without venturing into funky territory on the update function, namely simd (via packed_simd); as a learning experiment, I'd strongly recommend it. It is currently nightly-only (it was stable until Rust 1.33, then it broke, and it's on the way to being stable again), but it's worth the hassle - we're talking 10-15x speedups. Due to the nature of the repetitive operation in SHA256, it is ideal for a 256bit (u32x8) vector.