aes: simplify s-box implementationa #77

Primitive/Symmetric/Cipher/Block/AES/ExpandKey.cry

@@ -3,7 +3,7 @@ module Primitive::Symmetric::Cipher::Block::AES::ExpandKey where
 import Common::GF28 as GF28
 private type GF28 = GF28::GF28
 import Primitive::Symmetric::Cipher::Block::AES::State
-import Primitive::Symmetric::Cipher::Block::AES::SubByteSBox
+import Primitive::Symmetric::Cipher::Block::AES::SBox as Sbox
 
 parameter
   /** Number of 32 bit words in the key */
@@ -40,7 +40,7 @@ Rcon : [8] -> [4]GF28
 Rcon i = [ GF28::pow <| x |> (i-1), 0, 0, 0]
 
 SubWord : [4]GF28 -> [4]GF28
-SubWord bs = [ SubByte b | b <- bs ]
+SubWord bs = [ Sbox::sbox @ b | b <- bs ]
 
 RotWord : [4]GF28 -> [4]GF28
 RotWord [a0, a1, a2, a3] = [a1, a2, a3, a0]

Primitive/Symmetric/Cipher/Block/AES/Round.cry

@@ -2,7 +2,7 @@ module Primitive::Symmetric::Cipher::Block::AES::Round where
 
 import Common::GF28 as GF28
 import Primitive::Symmetric::Cipher::Block::AES::State
-import Primitive::Symmetric::Cipher::Block::AES::SubByteSBox
+import Primitive::Symmetric::Cipher::Block::AES::SBox
 
 /**
  * One round of AES.
@@ -14,9 +14,11 @@ AESRound rk s = rk ^ MixColumns (ShiftRows (SubBytes s))
 /**
  * SubBytes applies an invertible, non-linear transformation to the state.
  * [FIPS-197u1] Section 5.1.1.
+ *
+ * It does so by applying the AES S-box independently to each byte in the state.
  */
 SubBytes : State -> State
-SubBytes state = [ [ SubByte b | b <- row ] | row <- state ]
+SubBytes state = [ [ sbox @ b | b <- row ] | row <- state ]
 
 /**
  * ShiftRows transforms the state by cycling the last three rows.
@@ -58,7 +60,7 @@ InvShiftRows state = [ row >>> shiftAmount | row <- state
  * [FIPS-197u1] Section 5.3.3.
  */
 InvSubBytes : State -> State
-InvSubBytes state = [ [ InvSubByte b | b <- row ] | row <- state ]
+InvSubBytes state = [ [ sboxInv @ b | b <- row ] | row <- state ]
 
 /**
  * Inverts the `SubBytes` function.

Primitive/Symmetric/Cipher/Block/AES/SBox.cry

@@ -1,12 +1,31 @@
 module Primitive::Symmetric::Cipher::Block::AES::SBox where
 
-import Common::GF28
-import Primitive::Symmetric::Cipher::Block::AES::SubBytePlain
+import Common::GF28 as GF28
+private type GF28 = GF28::GF28
 
 type SBox = [256] GF28
 
+/**
+ * SBox: A non-linear substitution table for AES.
+ * [FIPS-197u1] Section 5.1.1.
+ *
+ * `GF28::inverse b` corresponds to Equation 5.2.
+ */
 sbox : SBox
-sbox = [ SubByte x | x <- [0 .. 255] ]
+sbox = [ transform (GF28::inverse b) | b <- [0 .. 255] ] where
+    // Equation 5.3.
+    transform b = GF28::add [b, (b >>> 4), (b >>> 5), (b >>> 6), (b >>> 7), c]
+    // The constant byte {01100011}.
+    c = 0x63
 
+/**
+ * Inverted substitution table for AES.
+ * [FIPS-197u1] Section 5.3.2.
+ */
 sboxInv : SBox
-sboxInv = [ InvSubByte x | x <- [0 .. 255] ]
+sboxInv = [ GF28::inverse (transformInv b) | b <- [0 .. 255] ] where
+    transformInv b = GF28::add [(b >>> 2), (b >>> 5), (b >>> 7), d]
+    d = 0x05
+
+// This `:prove`s efficiently.
+property sBoxInverts b = sboxInv @ ( sbox @ b) == b

Primitive/Symmetric/Cipher/Block/AES_specification.cry

@@ -19,7 +19,6 @@ module Primitive::Symmetric::Cipher::Block::AES_specification where
 
 import `Primitive::Symmetric::Cipher::Block::AES::Algorithm as AES
 import `Primitive::Symmetric::Cipher::Block::AES::ExpandKey
-import Primitive::Symmetric::Cipher::Block::AES::TBox
 
 parameter
    // This constraint enforces the standard key sizes of 128, 192, and