p3_binary_field/aes/

mod.rs

//! `GF(2^8)` in the polynomial basis of `x^8 + x^4 + x^3 + x + 1`, the block cipher's field.
//!
//! The tower already carries a `GF(2^8)`, and the two are the same field up to isomorphism.
//!
//! What this representation buys is hardware.
//!
//! The byte-wise product, linear map and inverse are all defined for this modulus alone.
//!
//! Sixty-four bytes therefore multiply, map or invert in one instruction.
//!
//! The tower's byte field takes a table lookup per element instead.

mod engine;
mod frobenius;
mod packed;

use core::fmt::{self, Debug, Display, Formatter};
use core::iter::{Product, Sum};
use core::ops::{Add, AddAssign, Div, DivAssign, Mul, MulAssign, Neg, Sub, SubAssign};

use num_bigint::BigUint;
use p3_field::op_assign_macros::{
    impl_add_assign, impl_add_base_field, impl_div_methods, impl_mul_base_field, impl_mul_methods,
    impl_sub_assign, impl_sub_base_field, ring_sum,
};
use p3_field::{Algebra, Field, Packable, PrimeCharacteristicRing, RawDataSerializable};
use rand::Rng;
use rand::distr::{Distribution, StandardUniform};
use serde::{Deserialize, Serialize};

pub use crate::aes::engine::ByteMatrix;
pub use crate::aes::frobenius::LinearizedPoly8b;
pub use crate::aes::packed::PackedRijndael8b;
use crate::cantor::CANTOR_BASIS_128;
use crate::gf2::characteristic_two_methods;
use crate::tower::TowerLevel;
use crate::{BinaryField8, Gf2};

/// The carryless product of a value with the modulus below its leading term.
///
/// That tail is `x^4 + x^3 + x + 1`, so its four set bits are the four shifts.
const fn fold_tail(v: u16) -> u16 {
    (v << 4) ^ (v << 3) ^ (v << 1) ^ v
}

/// The carryless product of two bytes, as a polynomial of degree at most fourteen.
const fn carryless(a: u8, b: u8) -> u16 {
    let a = a as u16;
    let mut acc = 0u16;
    let mut i = 0;
    while i < 8 {
        // All ones when bit `i` of the second operand is set, all zeros otherwise.
        let selected = 0u16.wrapping_sub(((b >> i) & 1) as u16);
        acc ^= (a << i) & selected;
        i += 1;
    }
    acc
}

/// The product of two elements, as bytes.
///
/// The modulus rewrites `x^8` as a tail of degree 4, so the high half folds down by that tail.
///
/// That fold reaches degree `6 + 4`, spilling three bits back over the top.
///
/// A second fold of those lands at degree `2 + 4`, below the top, so two rounds are exact.
#[inline]
pub(crate) const fn mul_bytes(a: u8, b: u8) -> u8 {
    let product = carryless(a, b);

    // The half above the modulus, rewritten by the tail.
    let folded = fold_tail(product >> 8);

    // What that fold pushed back over the top, rewritten once more.
    (product ^ folded ^ fold_tail(folded >> 8)) as u8
}

/// The inverse of an element, with zero left alone.
///
/// The multiplicative group has order 255, so inversion is the map raising to the power 254.
///
/// Zero is outside that group, and the same power sends it to itself.
///
/// That is also the convention the byte-wise hardware inverse follows.
///
/// # Algorithm
///
/// An addition chain on the exponent, eleven products of which seven are squarings:
///
/// ```text
///     1, 2, 3, 6, 12, 15, 30, 60, 120, 240, 252, 254
/// ```
#[inline]
pub(crate) const fn invert_byte(a: u8) -> u8 {
    let x2 = mul_bytes(a, a);
    let x3 = mul_bytes(x2, a);
    let x6 = mul_bytes(x3, x3);
    let x12 = mul_bytes(x6, x6);
    let x15 = mul_bytes(x12, x3);
    let x30 = mul_bytes(x15, x15);
    let x60 = mul_bytes(x30, x30);
    let x120 = mul_bytes(x60, x60);
    let x240 = mul_bytes(x120, x120);
    let x252 = mul_bytes(x240, x12);
    mul_bytes(x252, x2)
}

/// The images of the tower generators here, as the smallest solution at each level.
///
/// The tower is generated by elements obeying `X_k^2 + X_{k-1} X_k + 1 = 0`, with `X_-1 = 1`.
///
/// Any images obeying those same relations extend to a homomorphism of `F_2`-algebras.
///
/// Both sides are fields of the same size, so that homomorphism is an isomorphism.
const XI: [u8; 3] = {
    let mut xi = [0u8; 3];
    let mut k = 0;
    while k < 3 {
        let previous = if k == 0 { 1 } else { xi[k - 1] };
        let mut candidate = 1u16;
        while candidate < 256 {
            let c = candidate as u8;
            if mul_bytes(c, c) ^ mul_bytes(previous, c) ^ 1 == 0 {
                xi[k] = c;
                break;
            }
            candidate += 1;
        }
        assert!(xi[k] != 0, "the tower relations have no solution here");
        k += 1;
    }
    xi
};

/// The bit pattern of the multiplicative generator of the tower representation.
///
/// Its image here is the generator, so the change of basis carries one onto the other.
const TOWER_GENERATOR: u8 = 0x13;

/// The map out of the tower basis, column by column.
///
/// The tower basis vector indexed by a bit pattern is the product of the generators it names.
///
/// Stripping the lowest set bit reuses the image already computed for the rest.
const fn tower_images() -> [u8; 8] {
    let mut images = [0u8; 8];
    images[0] = 1;
    let mut j = 1usize;
    while j < 8 {
        let lowest = j.trailing_zeros() as usize;
        images[j] = mul_bytes(XI[lowest], images[j & (j - 1)]);
        j += 1;
    }
    images
}

/// Where each basis vector of the tower representation lands here.
const FROM_TOWER: ByteMatrix = ByteMatrix::from_images(tower_images());

/// Where each basis vector here lands in the tower representation.
///
/// An injective map on a byte is a permutation of the 256 values.
///
/// Searching for the preimage of each basis vector therefore recovers the inverse.
const INTO_TOWER: ByteMatrix = {
    let mut images = [0u8; 8];
    let mut b = 0;
    while b < 8 {
        let target = 1u8 << b;
        let mut x = 1u16;
        while x < 256 {
            if FROM_TOWER.apply(x as u8) == target {
                images[b] = x as u8;
                break;
            }
            x += 1;
        }
        b += 1;
    }
    ByteMatrix::from_images(images)
};

// Composing the two directions must leave every byte alone.
//
// That rules out both a wrong preimage search and a change of basis that is not invertible.
const _: () = assert!(
    INTO_TOWER.compose(FROM_TOWER).to_quadword() == ByteMatrix::IDENTITY.to_quadword(),
    "the change of basis is not invertible"
);

/// The Cantor basis in this representation, carried over from the tower's own.
///
/// Both representations then span the same additive transform domain.
const CANTOR_BASIS: [u8; 8] = {
    let mut basis = [0u8; 8];
    let mut i = 0;
    while i < 8 {
        basis[i] = FROM_TOWER.apply(CANTOR_BASIS_128[i] as u8);
        i += 1;
    }
    basis
};

/// The binary field `GF(2^8)` modulo `x^8 + x^4 + x^3 + x + 1`.
///
/// Bit `i` stores the coefficient of `x^i`, and every byte is a distinct field element.
///
/// This is the modulus the byte-wise field instructions are defined for.
///
/// That is what makes it worth carrying beside the tower's own byte field.
///
/// Arithmetic here is table-free: nothing indexes a table by an operand.
#[derive(Copy, Clone, Default, PartialEq, Eq, Hash, PartialOrd, Ord, Serialize, Deserialize)]
#[serde(transparent)]
#[repr(transparent)]
#[must_use]
pub struct Rijndael8b(u8);

impl Rijndael8b {
    /// The number of bits of an element.
    pub(crate) const BITS: usize = 8;

    /// Construct a field element from its byte representation.
    ///
    /// Every byte is a valid element.
    #[inline]
    pub const fn from_byte(byte: u8) -> Self {
        Self(byte)
    }

    /// The byte representation of this element.
    #[must_use]
    #[inline]
    pub const fn to_byte(self) -> u8 {
        self.0
    }

    /// The inverse of this element, with zero sent to zero.
    ///
    /// The chain runs whatever the operand is, so its cost says nothing about the value.
    #[inline]
    pub const fn invert_or_zero(self) -> Self {
        Self(invert_byte(self.0))
    }

    /// Multiplication by this element, as a map ready for the byte-wise engine.
    ///
    /// Scaling is `F_2`-linear, so a fixed factor is a matrix like any other.
    pub const fn scaling_matrix(self) -> ByteMatrix {
        let mut images = [0u8; 8];
        let mut b = 0;
        while b < 8 {
            images[b] = mul_bytes(self.0, 1 << b);
            b += 1;
        }
        ByteMatrix::from_images(images)
    }
}

impl Packable for Rijndael8b {}

impl Display for Rijndael8b {
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        Display::fmt(&self.0, f)
    }
}

impl Debug for Rijndael8b {
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        Debug::fmt(&self.0, f)
    }
}

impl Distribution<Rijndael8b> for StandardUniform {
    #[inline]
    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> Rijndael8b {
        let mut byte = [0u8; 1];
        rng.fill_bytes(&mut byte);
        Rijndael8b(byte[0])
    }
}

impl PrimeCharacteristicRing for Rijndael8b {
    type PrimeSubfield = Gf2;

    const ZERO: Self = Self(0);
    const ONE: Self = Self(1);
    // The characteristic is 2, so `TWO = ONE + ONE = ZERO`.
    const TWO: Self = Self(0);
    // The characteristic is 2, so `NEG_ONE = ONE`.
    const NEG_ONE: Self = Self(1);

    #[inline]
    fn from_prime_subfield(f: Self::PrimeSubfield) -> Self {
        Self::from_bool(f.is_one())
    }

    #[inline]
    fn from_bool(b: bool) -> Self {
        Self(u8::from(b))
    }

    characteristic_two_methods!();

    #[inline]
    fn square(&self) -> Self {
        Self(Self::frobenius_map(1).apply(self.0))
    }

    /// Raising to `2^k` is one tabulated map, whatever the exponent.
    #[inline]
    fn exp_power_of_2(&self, power_log: usize) -> Self {
        Self(Self::frobenius_map(power_log).apply(self.0))
    }
}

impl Field for Rijndael8b {
    type Packing = PackedRijndael8b<{ packed::PACKING_WIDTH }>;

    const GENERATOR: Self = Self(FROM_TOWER.apply(TOWER_GENERATOR));

    #[inline]
    fn try_inverse(&self) -> Option<Self> {
        // The chain runs whatever the operand, so its cost says nothing about the value.
        //
        // Zero is outside the multiplicative group and the chain already sends it to itself.
        let inverse = self.invert_or_zero();
        (self.0 != 0).then_some(inverse)
    }

    #[inline]
    fn try_sqrt(&self) -> Option<Self> {
        // Squaring has order 8 as a map, so the seventh power of it squares back.
        Some(self.exp_power_of_2(Self::BITS - 1))
    }

    #[inline]
    fn order() -> BigUint {
        BigUint::from(1u16 << Self::BITS)
    }

    /// An element of `GF(2^n)` is exactly `n` bits wide.
    #[inline]
    fn bits() -> usize {
        Self::BITS
    }

    /// The enumeration by bit pattern.
    ///
    /// The coordinates of the returned element over `GF(2)` are the bits of the index.
    ///
    /// # Panics
    /// Panics if the index does not fit in a byte, which would otherwise alias an earlier node.
    #[inline]
    fn interpolation_node(i: usize) -> Self {
        assert!(i < 1 << Self::BITS, "interpolation node index out of range");
        Self(i as u8)
    }
}

impl RawDataSerializable for Rijndael8b {
    const NUM_BYTES: usize = 1;

    #[inline]
    fn into_bytes(self) -> impl IntoIterator<Item = u8> {
        [self.0]
    }
}

impl crate::tower::private::Sealed for Rijndael8b {}

impl TowerLevel for Rijndael8b {
    type Repr = u8;

    const LOG_BITS: usize = 3;

    #[inline]
    fn from_repr(r: Self::Repr) -> Self {
        Self(r)
    }

    #[inline]
    fn to_repr(self) -> Self::Repr {
        self.0
    }

    /// The tower scales by a basis vector, which costs it a shift.
    ///
    /// Here the same element is an arbitrary one, so it costs a full product.
    #[inline]
    fn mul_alpha(self) -> Self {
        Self(mul_bytes(self.0, XI[2]))
    }

    /// # Panics
    /// Panics if the stream ends before a whole element has been read.
    #[inline]
    fn from_le_byte_iter(mut bytes: impl Iterator<Item = u8>) -> Self {
        Self(
            bytes
                .next()
                .expect("byte stream ended before a whole element was read"),
        )
    }

    /// # Panics
    /// Panics if the index is at least the bit width of the field.
    #[inline]
    fn cantor_basis(i: usize) -> Self {
        assert!(i < Self::BITS, "Cantor basis index out of range");
        Self(CANTOR_BASIS[i])
    }
}

impl From<BinaryField8> for Rijndael8b {
    /// The same field element, seen in the polynomial basis.
    #[inline]
    fn from(x: BinaryField8) -> Self {
        Self(FROM_TOWER.apply(x.to_repr()))
    }
}

impl From<Rijndael8b> for BinaryField8 {
    /// The same field element, seen in the tower basis.
    #[inline]
    fn from(x: Rijndael8b) -> Self {
        Self::from_repr(INTO_TOWER.apply(x.0))
    }
}

impl Add for Rijndael8b {
    type Output = Self;

    #[inline]
    #[allow(clippy::suspicious_arithmetic_impl)]
    fn add(self, rhs: Self) -> Self {
        // Addition in characteristic 2 is `XOR`.
        Self(self.0 ^ rhs.0)
    }
}

impl Sub for Rijndael8b {
    type Output = Self;

    #[inline]
    #[allow(clippy::suspicious_arithmetic_impl)]
    fn sub(self, rhs: Self) -> Self {
        // Subtraction coincides with addition in characteristic 2.
        self + rhs
    }
}

impl Neg for Rijndael8b {
    type Output = Self;

    #[inline]
    fn neg(self) -> Self {
        // `-x = x` in characteristic 2.
        self
    }
}

impl Mul for Rijndael8b {
    type Output = Self;

    #[inline]
    fn mul(self, rhs: Self) -> Self {
        Self(mul_bytes(self.0, rhs.0))
    }
}

impl_add_assign!(Rijndael8b);
impl_sub_assign!(Rijndael8b);
impl_mul_methods!(Rijndael8b);
impl_div_methods!(Rijndael8b, Rijndael8b);
ring_sum!(Rijndael8b);

impl From<Gf2> for Rijndael8b {
    /// `GF(2)` is the prime subfield, embedded as `{ZERO, ONE}`.
    #[inline]
    fn from(x: Gf2) -> Self {
        Self::from_prime_subfield(x)
    }
}

impl_add_base_field!(Rijndael8b, Gf2);
impl_sub_base_field!(Rijndael8b, Gf2);
impl_mul_base_field!(Rijndael8b, Gf2);

impl Algebra<Gf2> for Rijndael8b {}

#[cfg(test)]
mod tests {
    extern crate std;

    use p3_field::{Field, PrimeCharacteristicRing};
    use proptest::prelude::*;

    use super::{CANTOR_BASIS, Rijndael8b, XI, mul_bytes};
    use crate::tower::TowerLevel;
    use crate::{BinaryField8, ByteMatrix};

    /// Multiplication straight from the modulus, one coefficient of the multiplier at a time.
    ///
    /// The routine under test packs both folds into shifts of a whole byte.
    ///
    /// This loop shares none of that structure.
    fn schoolbook(a: u8, b: u8) -> u8 {
        let mut acc = 0u8;
        for i in (0..8).rev() {
            // Double, rewriting the overflow as the modulus tail.
            let overflowed = acc & 0x80 != 0;
            acc <<= 1;
            if overflowed {
                acc ^= 0x1b;
            }
            if (b >> i) & 1 == 1 {
                acc ^= a;
            }
        }
        acc
    }

    #[test]
    fn the_modulus_reduces_the_way_the_polynomial_says() {
        // x^7 * x = x^8 = x^4 + x^3 + x + 1, the tail spelled 0x1b.
        assert_eq!(mul_bytes(0x80, 0x02), 0x1b);

        // The worked example from the block cipher's own column mixing.
        assert_eq!(mul_bytes(0xd4, 0x02), 0xb3);
        assert_eq!(mul_bytes(0xbf, 0x03), 0xda);

        // Far below the modulus, so nothing folds at all.
        assert_eq!(mul_bytes(0x02, 0x04), 0x08);
    }

    #[test]
    fn the_published_substitution_box_entries_come_out() {
        // The cipher's substitution box is inversion followed by a fixed affine map.
        //
        // Its published entries therefore pin both the inverse and the map engine.
        //
        //     s(x) = A(x^254) + 0x63,   A the map with the circulant rows of 0x1f
        let affine = ByteMatrix::from_images([0x1f, 0x3e, 0x7c, 0xf8, 0xf1, 0xe3, 0xc7, 0x8f]);

        // The quadword published for the inverting affine-byte instruction's matrix operand.
        //
        // This pins the row layout to the hardware on every host, not only on those that run it.
        assert_eq!(affine.to_quadword(), 0xf1e3_c78f_1f3e_7cf8);

        let sbox = |x: u8| affine.apply(super::invert_byte(x)) ^ 0x63;

        // Entries 0x00, 0x01, 0x53 and 0xff of the published table.
        assert_eq!(sbox(0x00), 0x63);
        assert_eq!(sbox(0x01), 0x7c);
        assert_eq!(sbox(0x53), 0xed);
        assert_eq!(sbox(0xff), 0x16);
    }

    #[test]
    fn every_nonzero_byte_inverts() {
        // Fixture state: 255 nonzero bytes, each with exactly one inverse.
        for byte in 1..=u8::MAX {
            assert_eq!(mul_bytes(byte, super::invert_byte(byte)), 1, "{byte:#x}");
        }
        // Zero is outside the group and the exponentiation leaves it alone.
        assert_eq!(super::invert_byte(0), 0);
        assert_eq!(Rijndael8b::ZERO.try_inverse(), None);
    }

    #[test]
    fn the_generator_is_the_image_of_the_tower_generator() {
        // Both representations are the same field, so one generator maps onto the other.
        //
        // This is what pins the transcribed bit pattern the constant is built from.
        assert_eq!(
            Rijndael8b::GENERATOR,
            Rijndael8b::from(BinaryField8::GENERATOR)
        );
    }

    #[test]
    fn the_generator_has_the_full_order() {
        // Invariant: the group has order 255 = 3 * 5 * 17.
        //
        // A generator survives all three proper-divisor powers.
        let g = Rijndael8b::GENERATOR;
        for divisor in [3u64, 5, 17] {
            assert_ne!(g.exp_u64(255 / divisor), Rijndael8b::ONE, "{divisor}");
        }
        assert_eq!(g.exp_u64(255), Rijndael8b::ONE);
    }

    #[test]
    fn the_generator_images_satisfy_the_tower_relations() {
        // Invariant: xi_k^2 + xi_{k-1} xi_k + 1 = 0, with xi_-1 = 1.
        //
        // This is what makes the change of basis multiplicative rather than merely invertible.
        let mut previous = 1u8;
        for (k, &xi) in XI.iter().enumerate() {
            assert_eq!(
                mul_bytes(xi, xi) ^ mul_bytes(previous, xi) ^ 1,
                0,
                "generator {k}"
            );
            previous = xi;
        }
    }

    #[test]
    fn the_change_of_basis_is_a_ring_isomorphism() {
        // Exhaustive over both operands: 65536 pairs settles the whole field.
        for a in 0..=u8::MAX {
            for b in 0..=u8::MAX {
                let (x, y) = (BinaryField8::from_repr(a), BinaryField8::from_repr(b));

                // Additive, which any change of coordinates would give.
                assert_eq!(
                    Rijndael8b::from(x + y),
                    Rijndael8b::from(x) + Rijndael8b::from(y)
                );

                // Multiplicative, which is the part that pins the generator images.
                assert_eq!(
                    Rijndael8b::from(x.reference_mul(y)),
                    Rijndael8b::from(x) * Rijndael8b::from(y)
                );
            }
            // And back again, so neither direction can be a different map.
            let x = Rijndael8b::from_byte(a);
            assert_eq!(Rijndael8b::from(BinaryField8::from(x)), x);
        }
    }

    #[test]
    fn the_cantor_basis_satisfies_its_recurrence() {
        // Invariant: v_0 = 1 and v_i^2 + v_i = v_{i-1}.
        assert_eq!(Rijndael8b::cantor_basis(0), Rijndael8b::ONE);
        for i in 1..Rijndael8b::BITS {
            let v = Rijndael8b::cantor_basis(i);
            assert_eq!(
                v.square() + v,
                Rijndael8b::cantor_basis(i - 1),
                "vector {i}"
            );
        }

        // And it is the image of the tower's own, so both span the same transform domain.
        for i in 0..Rijndael8b::BITS {
            assert_eq!(
                Rijndael8b::cantor_basis(i),
                Rijndael8b::from(BinaryField8::cantor_basis(i)),
                "vector {i}"
            );
        }
        assert_eq!(CANTOR_BASIS[0], 1);
    }

    #[test]
    fn scaling_by_the_top_generator_agrees_with_the_tower() {
        // The tower scales by a basis vector; here the same element is an arbitrary one.
        for byte in 0..=u8::MAX {
            let x = BinaryField8::from_repr(byte);
            assert_eq!(
                Rijndael8b::from(x).mul_alpha(),
                Rijndael8b::from(x.mul_alpha()),
                "{byte:#x}"
            );
        }
    }

    #[test]
    #[should_panic = "Cantor basis index out of range"]
    fn the_cantor_basis_rejects_an_index_beyond_the_field() {
        let _vector = Rijndael8b::cantor_basis(8);
    }

    #[test]
    #[should_panic = "byte stream ended before a whole element was read"]
    fn a_truncated_byte_stream_is_rejected() {
        let _element = Rijndael8b::from_le_byte_iter(core::iter::empty());
    }

    proptest! {
        #![proptest_config(ProptestConfig::with_cases(2000))]

        /// The shift-and-fold product must agree with reduction straight from the modulus.
        #[test]
        fn the_product_matches_the_schoolbook_reduction(a: u8, b: u8) {
            prop_assert_eq!(mul_bytes(a, b), schoolbook(a, b));
        }

        /// A fixed factor tabulated as a map must be the same product.
        #[test]
        fn the_scaling_matrix_is_the_product(a: u8, b: u8) {
            prop_assert_eq!(Rijndael8b::from_byte(a).scaling_matrix().apply(b), mul_bytes(a, b));
        }

        /// The square root must square back, on every byte.
        #[test]
        fn the_square_root_squares_back(a: u8) {
            let x = Rijndael8b::from_byte(a);
            let root = x.try_sqrt().expect("every element of a binary field is a square");
            prop_assert_eq!(root.square(), x);
        }
    }

    #[test]
    fn serde_accepts_every_byte() {
        // No bit is masked off, so nothing is out of range.
        for byte in [0u8, 1, 0x1b, u8::MAX] {
            let x = Rijndael8b::from_byte(byte);
            let encoded = serde_json::to_string(&x).unwrap();
            assert_eq!(serde_json::from_str::<Rijndael8b>(&encoded).unwrap(), x);
        }
    }
}