One-pass encryption/decryption

[tarcieri]

Currently all of the AEAD implementations do two passes over the plaintext/ciphertext when encrypting/decrypting respectively: for encryption, they encrypt the plaintext in the first pass, and authenticate it in the second pass. For decryption, it's vice versa.

A better approach is to pick a number of blocks to operate on in parallel and encrypt/authenticate or authenticate/decrypt in a single pass. This has better cache locality, e.g. when we encrypt data, store the resulting ciphertext, then load it again to do authentication, that is pretty much guaranteed to hit L1 cache when doing it in a single pass (and ideally we could hand off values still stored in e.g. SIMD registers)

This is a tracking issue for converting the implementations of these respective algorithms to be one pass. It also might be good to discuss ways we could have a generic implementation of one pass encryption/decryption in the aead crate (especially one specialized for the non-SIV stream-cipher + universal-hash use case) which can be reused across different algorithm implementations.

• aes-gcm
• aes-gcm-siv†
• aes-siv†
• chacha20poly1305
• xsalsa20poly1305

†NOTE: SIV modes by definition cannot support 1-pass encryption (because the first pass generates the synthetic IV, which must be known in advance before encryption can be performed). However, they can support 1-pass decryption, since the IV is known in advance in that case.

[tarcieri]

Related issue: generalized AEAD implementations based on stream ciphers RustCrypto/traits#45

[nico-abram]

As a complete crypto noob who got here from github explore (But not a complete rust noob), would this be feasible? (Or even just a small part, like trying to make decryption for chacha20poly1305 single pass) How robust are existing tests?

Would the change mostly be changing implementations like https://github.com/RustCrypto/AEADs/blob/master/chacha20poly1305/src/cipher.rs#L66-L91 into something more like https://github.com/RustCrypto/AEADs/blob/master/aes-gcm-siv/src/lib.rs#L317-L347 ?

[tarcieri]

Yes, but it also needs to be done in a way that actually improves performance. I've tried to do this change naively a few times (to aes-gcm, mainly, I might still have the code around) and it decreased performance.

I think doing it properly might require keeping the data flowing through XMM registers... at the very least it needs to all stay in L1 cache.

ChaChaPoly is even trickier because this issue is a micro-optimization and so far we don't have an AVX2 backend for Poly1305 (see RustCrypto/universal-hashes#49)

[nico-abram]

Thanks for the response!

I think I'll give it a shot. Do not let that stop anyone else from trying it since I don't have much hope I'll be able to do much.

How important is performance compiling with avx support vs without? (i.e, would you mostly care about speed when compiling with simd extensions or does the "default" cargo build configuration also matter a lot?)

[newpavlov]

The ccm crate could also use a single pass encryption/decryption.

[str4d]

Let's take chacha20poly1305 in its current form and look at the AVX2 hot path (ignoring all the autodetect code in chacha20 and poly1305).

chacha20poly1305::cipher:

impl<C> Cipher<C> where C: StreamCipher + StreamCipherSeek,
{
    pub(crate) fn encrypt_in_place_detached(
        mut self,
        associated_data: &[u8],
        buffer: &mut [u8],
    ) -> Result<Tag, Error> {
        // ...

        // Not currently implemented, but imagine we did this:
        for chunk in buffer.chunks_mut(BLOCK_SIZE * 4) {
            self.cipher.apply_keystream(chunk);
            self.mac.update_padded(chunk);
        }

        // ...
    }
}

chacha20::backend::autodetect:

pub(crate) const BUFFER_SIZE: usize = BLOCK_SIZE * 4;

chacha20::chacha:

impl<R: Rounds, MC: MaxCounter> StreamCipher for ChaCha<R, MC> {
    fn try_apply_keystream(&mut self, mut data: &mut [u8]) -> Result<(), LoopError> {
        // ...

        let mut chunks = data.chunks_exact_mut(BUFFER_SIZE);
        for chunk in &mut chunks {
            let counter_with_offset = self.counter_offset.checked_add(counter).unwrap();
            self.block.apply_keystream(counter_with_offset, chunk);
            counter = counter.checked_add(COUNTER_INCR).unwrap();
        }

        // ...
    }
}

chacha20::backend::avx2:

const BLOCKS: usize = 4;

impl<R: Rounds> Core<R> {
    pub fn apply_keystream(&self, counter: u64, output: &mut [u8]) {
        debug_assert_eq!(output.len(), BUFFER_SIZE);

        unsafe {
            let state = State {
                a: self.v0,
                b: self.v1,
                c: self.v2,
                d: iv_setup(self.iv, counter),
            };
            let state = self.rounds(state);

            for i in 0..BLOCKS {
                for (chunk, a) in output[i * BLOCK_SIZE..(i + 1) * BLOCK_SIZE]
                    .chunks_mut(0x10)
                    .zip(
                        [state.a, state.b, state.c, state.d]
                            .iter()
                            .map(|s| s.blocks[i]),
                    )
                {
                    let b = _mm_loadu_si128(chunk.as_ptr() as *const __m128i);
                    let out = _mm_xor_si128(a, b);
                    _mm_storeu_si128(chunk.as_mut_ptr() as *mut __m128i, out);
                }
            }
        }
    }
}

universal_hash:

pub trait UniversalHash {
    fn update_padded(&mut self, data: &[u8]) {
        let mut chunks = data.chunks_exact(Self::BlockSize::to_usize());

        for chunk in &mut chunks {
            self.update(GenericArray::from_slice(chunk));
        }

        // ...
    }
}

poly1305::backend::avx2:

impl State {
    pub(crate) unsafe fn compute_block(&mut self, block: &Block, partial: bool) {
        // ...

        self.cached_blocks[self.num_cached_blocks].copy_from_slice(block);
        if self.num_cached_blocks < 3 {
            self.num_cached_blocks += 1;
            return;
        } else {
            self.num_cached_blocks = 0;
        }

        let p = Aligned4x130::from_blocks(&self.cached_blocks);
        // ...
}

poly1305::backend::avx2::helpers:

impl Aligned4x130 {
    pub(super) unsafe fn from_blocks(src: &[Block; 4]) -> Self {
        // 26-bit mask on each 32-bit word.
        let mask_26 = _mm256_set1_epi32(0x3ffffff);
        // Sets bit 24 of each 32-bit word.
        let set_hibit = _mm256_set1_epi32(1 << 24);

        // - Load the four blocks into the following 32-bit word layout:
        //      [b33, b32, b31, b30, b23, b22, b21, b20]
        //      [b13, b12, b11, b10, b03, b02, b01, b00]
        //
        // - Unpack the upper and lower 64 bits:
        //      [b33, b32, b13, b12, b23, b22, b03, b02]
        //      [b31, b30, b11, b10, b21, b20, b01, b00]
        //
        // - Swap the middle two 64-bit words:
        // a0 = [b33, b32, b23, b22, b13, b12, b03, b02]
        // a1 = [b31, b30, b21, b20, b11, b10, b01, b00]
        let (lo, hi) = src.split_at(2);
        let blocks_23 = _mm256_loadu_si256(hi.as_ptr() as *const _);
        let blocks_01 = _mm256_loadu_si256(lo.as_ptr() as *const _);
        // ...
    }
}

So, the hot path above:

• Splits the plaintext into 4-block byte chunks.
• Passes the 4-block byte chunk to chacha20, which:
  • Chunks it into 4-block chunks (no-op).
  • Splits the chunk into individual blocks and:
    • Calls _mm_loadu_si128 on the block to load it into a __m128i.
    • XORs the stream into the __m128i.
    • Stores the __m128i back into the block.
• Passes the (now-encrypted) 4-block byte chunk to poly1305, which:
  • Splits the chunk into individual blocks.
  • Copies each block into a 4-block cache (reconstructing the chunk).
  • Calls _mm256_loadu_si256 on each half of the 4-block cache.
  • Draws the rest of the polynomiOwl.

So the immediate blocker is that the UniversalHash trait doesn't provide any API to process multiple blocks at a time (StreamCipher::try_apply_keystream allows the implementor to choose the chunking, whereas UniversalHash::update is typed on a single block).

Once that is addressed, poly1305 could directly consume 4-block chunks without using its cache, at which point we would be consistently passing around a 4-block chunk size. Then the question becomes the form in which we pass the chunk around. Sketching out two possible directions:

• Add an associated type for the chunk to an aead trait, which is constrained to equal an equivalent associated type in cipher and universal-hash. Then chacha20 and poly1305 would separately set it to the same concrete type.
• Add an an aead::AeadChunk trait, and implement it in chacha20poly1305. Have some way to map it to the chunk inputs of cipher and universal-hash.

[tarcieri]

Yeah, it's definitely a drawback that the UniversalHash trait doesn't provide a multi-block API. It's also problematic that data in AEADs is round-tripping through calls like _mm_storeu_si128/_mm_loadu_si128, especially along crate boundaries where it's pretty much guaranteed not to get optimized away even in cases where it's aligned.

These sorts of optimization problems for passing data between stream ciphers and universal hash functions were the impetus for the simd-buffers work I was attempting here:

RustCrypto/utils#221

I abandoned that, but now I wonder if maybe crypto_bigint::UInt might make a reasonable replacement buffer type for these use cases, especially if we ensured they were properly aligned such that they could be safely converted to similarly-structured SIMD types. Then a multi-block API could operate over slices of those structured SIMD buffers with guaranteed alignment.