-----------------------------------------------------------------------------
-- |
-- Module    : Documentation.SBV.Examples.Crypto.AES
-- Copyright : (c) Levent Erkok
-- License   : BSD3
-- Maintainer: erkokl@gmail.com
-- Stability : experimental
--
-- An implementation of AES (Advanced Encryption Standard), using SBV.
-- For details on AES, see <http://en.wikipedia.org/wiki/Advanced_Encryption_Standard>.
--
-- We do a T-box implementation, which leads to good C code as we can take
-- advantage of look-up tables. Note that we make virtually no attempt to
-- optimize our Haskell code. The concern here is not with getting Haskell running
-- fast at all. The idea is to program the T-Box implementation as naturally and clearly
-- as possible in Haskell, and have SBV's code-generator generate fast C code automatically.
-- Therefore, we merely use ordinary Haskell lists as our data-structures, and do not
-- bother with any unboxing or strictness annotations. Thus, we achieve the separation
-- of concerns: Correctness via clarity and simplicity and proofs on the Haskell side,
-- performance by relying on SBV's code generator. If necessary, the generated code
-- can be FFI'd back into Haskell to complete the loop.
--
-- All 3 valid key sizes (128, 192, and 256) as required by the FIPS-197 standard
-- are supported.
-----------------------------------------------------------------------------

{-# LANGUAGE DataKinds        #-}
{-# LANGUAGE ParallelListComp #-}
{-# LANGUAGE TypeApplications #-}

{-# OPTIONS_GHC -Wall -Werror #-}

module Documentation.SBV.Examples.Crypto.AES where

import Control.Monad (void)

import Data.SBV
import Data.SBV.Tools.CodeGen
import Data.SBV.Tools.Polynomial

import Data.List (transpose)
import Data.Maybe (fromJust)

import Numeric (showHex)

-----------------------------------------------------------------------------
-- * Formalizing GF(2^8)
-----------------------------------------------------------------------------

-- | An element of the Galois Field 2^8, which are essentially polynomials with
-- maximum degree 7. They are conveniently represented as values between 0 and 255.
type GF28 = SWord 8

-- | Multiplication in GF(2^8). This is simple polynomial multiplication, followed
-- by the irreducible polynomial @x^8+x^4+x^3+x^1+1@. We simply use the 'pMult'
-- function exported by SBV to do the operation. 
gf28Mult :: GF28 -> GF28 -> GF28
gf28Mult :: GF28 -> GF28 -> GF28
gf28Mult GF28
x GF28
y = (GF28, GF28, [Int]) -> GF28
forall a. Polynomial a => (a, a, [Int]) -> a
pMult (GF28
x, GF28
y, [Int
8, Int
4, Int
3, Int
1, Int
0])

-- | Exponentiation by a constant in GF(2^8). The implementation uses the usual
-- square-and-multiply trick to speed up the computation.
gf28Pow :: GF28 -> Int -> GF28
gf28Pow :: GF28 -> Int -> GF28
gf28Pow GF28
n = Int -> GF28
forall a. (Integral a, Bits a) => a -> GF28
pow
  where sq :: GF28 -> GF28
sq GF28
x  = GF28
x GF28 -> GF28 -> GF28
`gf28Mult` GF28
x
        pow :: a -> GF28
pow a
0    = GF28
1
        pow a
i
         | a -> Bool
forall a. Integral a => a -> Bool
odd a
i = GF28
n GF28 -> GF28 -> GF28
`gf28Mult` GF28 -> GF28
sq (a -> GF28
pow (a
i a -> Int -> a
forall a. Bits a => a -> Int -> a
`shiftR` Int
1))
         | Bool
True  = GF28 -> GF28
sq (a -> GF28
pow (a
i a -> Int -> a
forall a. Bits a => a -> Int -> a
`shiftR` Int
1))

-- | Computing inverses in GF(2^8). By the mathematical properties of GF(2^8)
-- and the particular irreducible polynomial used @x^8+x^5+x^3+x^1+1@, it
-- turns out that raising to the 254 power gives us the multiplicative inverse.
-- Of course, we can prove this using SBV:
--
-- >>> prove $ \x -> x ./= 0 .=> x `gf28Mult` gf28Inverse x .== 1
-- Q.E.D.
--
-- Note that we exclude @0@ in our theorem, as it does not have a
-- multiplicative inverse.
gf28Inverse :: GF28 -> GF28
gf28Inverse :: GF28 -> GF28
gf28Inverse GF28
x = GF28
x GF28 -> Int -> GF28
`gf28Pow` Int
254

-----------------------------------------------------------------------------
-- * Implementing AES
-----------------------------------------------------------------------------

-----------------------------------------------------------------------------
-- ** Types and basic operations
-----------------------------------------------------------------------------
-- | AES state. The state consists of four 32-bit words, each of which is in turn treated
-- as four GF28's, i.e., 4 bytes. The T-Box implementation keeps the four-bytes together
-- for efficient representation.
type State = [SWord 32]

-- | The key, which can be 128, 192, or 256 bits. Represented as a sequence of 32-bit words.
type Key = [SWord 32]

-- | The key schedule. AES executes in rounds, and it treats first and last round keys slightly
-- differently than the middle ones. We reflect that choice by being explicit about it in our type.
-- The length of the middle list of keys depends on the key-size, which in turn determines
-- the number of rounds.
type KS = (Key, [Key], Key)

-- | Rotating a state row by a fixed amount to the right.
rotR :: [GF28] -> Int -> [GF28]
rotR :: [GF28] -> Int -> [GF28]
rotR [GF28
a, GF28
b, GF28
c, GF28
d] Int
1 = [GF28
d, GF28
a, GF28
b, GF28
c]
rotR [GF28
a, GF28
b, GF28
c, GF28
d] Int
2 = [GF28
c, GF28
d, GF28
a, GF28
b]
rotR [GF28
a, GF28
b, GF28
c, GF28
d] Int
3 = [GF28
b, GF28
c, GF28
d, GF28
a]
rotR [GF28]
xs           Int
i = [Char] -> [GF28]
forall a. HasCallStack => [Char] -> a
error ([Char] -> [GF28]) -> [Char] -> [GF28]
forall a b. (a -> b) -> a -> b
$ [Char]
"rotR: Unexpected input: " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ ([GF28], Int) -> [Char]
forall a. Show a => a -> [Char]
show ([GF28]
xs, Int
i)

-----------------------------------------------------------------------------
-- ** The key schedule
-----------------------------------------------------------------------------

-- | Definition of round-constants, as specified in Section 5.2 of the AES standard.
roundConstants :: [GF28]
roundConstants :: [GF28]
roundConstants = GF28
0 GF28 -> [GF28] -> [GF28]
forall a. a -> [a] -> [a]
: [ GF28 -> Int -> GF28
gf28Pow GF28
2 (Int
kInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1) | Int
k <- [Int
1 .. ] ]

-- | The @InvMixColumns@ transformation, as described in Section 5.3.3 of the standard. Note
-- that this transformation is only used explicitly during key-expansion in the T-Box implementation
-- of AES.
invMixColumns :: State -> State
invMixColumns :: State -> State
invMixColumns State
state = ([GF28] -> SWord 32) -> [[GF28]] -> State
forall a b. (a -> b) -> [a] -> [b]
map [GF28] -> SWord 32
forall a. ByteConverter a => [GF28] -> a
fromBytes ([[GF28]] -> State) -> [[GF28]] -> State
forall a b. (a -> b) -> a -> b
$ [[GF28]] -> [[GF28]]
forall a. [[a]] -> [[a]]
transpose ([[GF28]] -> [[GF28]]) -> [[GF28]] -> [[GF28]]
forall a b. (a -> b) -> a -> b
$ [[GF28]] -> [[GF28]]
mmult ((SWord 32 -> [GF28]) -> State -> [[GF28]]
forall a b. (a -> b) -> [a] -> [b]
map SWord 32 -> [GF28]
forall a. ByteConverter a => a -> [GF28]
toBytes State
state)
 where dot :: [a -> c] -> [a] -> c
dot [a -> c]
f   = (c -> c -> c) -> [c] -> c
forall (t :: * -> *) a. Foldable t => (a -> a -> a) -> t a -> a
foldr1 c -> c -> c
forall a. Bits a => a -> a -> a
xor ([c] -> c) -> ([a] -> [c]) -> [a] -> c
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((a -> c) -> a -> c) -> [a -> c] -> [a] -> [c]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (a -> c) -> a -> c
forall a b. (a -> b) -> a -> b
($) [a -> c]
f
       mmult :: [[GF28]] -> [[GF28]]
mmult [[GF28]]
n = [([GF28] -> GF28) -> [[GF28]] -> [GF28]
forall a b. (a -> b) -> [a] -> [b]
map ([GF28 -> GF28] -> [GF28] -> GF28
forall c a. Bits c => [a -> c] -> [a] -> c
dot [GF28 -> GF28]
r) [[GF28]]
n | [GF28 -> GF28]
r <- [ [GF28 -> GF28
mE, GF28 -> GF28
mB, GF28 -> GF28
mD, GF28 -> GF28
m9]
                                       , [GF28 -> GF28
m9, GF28 -> GF28
mE, GF28 -> GF28
mB, GF28 -> GF28
mD]
                                       , [GF28 -> GF28
mD, GF28 -> GF28
m9, GF28 -> GF28
mE, GF28 -> GF28
mB]
                                       , [GF28 -> GF28
mB, GF28 -> GF28
mD, GF28 -> GF28
m9, GF28 -> GF28
mE]
                                       ]]
       -- table-lookup versions of gf28Mult with the constants used in invMixColumns
       mE :: GF28 -> GF28
mE = [GF28] -> GF28 -> GF28 -> GF28
forall a b.
(Mergeable a, Ord b, SymVal b, Num b) =>
[a] -> a -> SBV b -> a
select [GF28]
mETable GF28
0
       mB :: GF28 -> GF28
mB = [GF28] -> GF28 -> GF28 -> GF28
forall a b.
(Mergeable a, Ord b, SymVal b, Num b) =>
[a] -> a -> SBV b -> a
select [GF28]
mBTable GF28
0
       mD :: GF28 -> GF28
mD = [GF28] -> GF28 -> GF28 -> GF28
forall a b.
(Mergeable a, Ord b, SymVal b, Num b) =>
[a] -> a -> SBV b -> a
select [GF28]
mDTable GF28
0
       m9 :: GF28 -> GF28
m9 = [GF28] -> GF28 -> GF28 -> GF28
forall a b.
(Mergeable a, Ord b, SymVal b, Num b) =>
[a] -> a -> SBV b -> a
select [GF28]
m9Table GF28
0
       mETable :: [GF28]
mETable = (GF28 -> GF28) -> [GF28] -> [GF28]
forall a b. (a -> b) -> [a] -> [b]
map (GF28 -> GF28 -> GF28
gf28Mult GF28
0xE) [GF28
0..GF28
255]
       mBTable :: [GF28]
mBTable = (GF28 -> GF28) -> [GF28] -> [GF28]
forall a b. (a -> b) -> [a] -> [b]
map (GF28 -> GF28 -> GF28
gf28Mult GF28
0xB) [GF28
0..GF28
255]
       mDTable :: [GF28]
mDTable = (GF28 -> GF28) -> [GF28] -> [GF28]
forall a b. (a -> b) -> [a] -> [b]
map (GF28 -> GF28 -> GF28
gf28Mult GF28
0xD) [GF28
0..GF28
255]
       m9Table :: [GF28]
m9Table = (GF28 -> GF28) -> [GF28] -> [GF28]
forall a b. (a -> b) -> [a] -> [b]
map (GF28 -> GF28 -> GF28
gf28Mult GF28
0x9) [GF28
0..GF28
255]

-- | Key expansion. Starting with the given key, returns an infinite sequence of
-- words, as described by the AES standard, Section 5.2, Figure 11.
keyExpansion :: Int -> Key -> [Key]
keyExpansion :: Int -> State -> [State]
keyExpansion Int
nk State
key = State -> [State]
forall a. [a] -> [[a]]
chop4 State
keys
   where keys :: [SWord 32]
         keys :: State
keys = State
key State -> State -> State
forall a. [a] -> [a] -> [a]
++ [Int -> SWord 32 -> SWord 32 -> SWord 32
nextWord Int
i SWord 32
prev SWord 32
old | Int
i <- [Int
nk ..] | SWord 32
prev <- Int -> State -> State
forall a. Int -> [a] -> [a]
drop (Int
nkInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1) State
keys | SWord 32
old <- State
keys]
         chop4 :: [a] -> [[a]]
         chop4 :: [a] -> [[a]]
chop4 [a]
xs = let ([a]
f, [a]
r) = Int -> [a] -> ([a], [a])
forall a. Int -> [a] -> ([a], [a])
splitAt Int
4 [a]
xs in [a]
f [a] -> [[a]] -> [[a]]
forall a. a -> [a] -> [a]
: [a] -> [[a]]
forall a. [a] -> [[a]]
chop4 [a]
r
         nextWord :: Int -> SWord 32 -> SWord 32 -> SWord 32
         nextWord :: Int -> SWord 32 -> SWord 32 -> SWord 32
nextWord Int
i SWord 32
prev SWord 32
old
           | Int
i Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` Int
nk Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0           = SWord 32
old SWord 32 -> SWord 32 -> SWord 32
forall a. Bits a => a -> a -> a
`xor` SWord 32 -> GF28 -> SWord 32
subWordRcon (SWord 32
prev SWord 32 -> Int -> SWord 32
forall a. Bits a => a -> Int -> a
`rotateL` Int
8) ([GF28]
roundConstants [GF28] -> Int -> GF28
forall a. [a] -> Int -> a
!! (Int
i Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
nk))
           | Int
i Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` Int
nk Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
4 Bool -> Bool -> Bool
&& Int
nk Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
6 = SWord 32
old SWord 32 -> SWord 32 -> SWord 32
forall a. Bits a => a -> a -> a
`xor` SWord 32 -> GF28 -> SWord 32
subWordRcon SWord 32
prev GF28
0
           | Bool
True                      = SWord 32
old SWord 32 -> SWord 32 -> SWord 32
forall a. Bits a => a -> a -> a
`xor` SWord 32
prev
         subWordRcon :: SWord 32 -> GF28 -> SWord 32
         subWordRcon :: SWord 32 -> GF28 -> SWord 32
subWordRcon SWord 32
w GF28
rc = [GF28] -> SWord 32
forall a. ByteConverter a => [GF28] -> a
fromBytes [GF28
a GF28 -> GF28 -> GF28
forall a. Bits a => a -> a -> a
`xor` GF28
rc, GF28
b, GF28
c, GF28
d]
            where [GF28
a, GF28
b, GF28
c, GF28
d] = (GF28 -> GF28) -> [GF28] -> [GF28]
forall a b. (a -> b) -> [a] -> [b]
map GF28 -> GF28
sbox ([GF28] -> [GF28]) -> [GF28] -> [GF28]
forall a b. (a -> b) -> a -> b
$ SWord 32 -> [GF28]
forall a. ByteConverter a => a -> [GF28]
toBytes SWord 32
w

-----------------------------------------------------------------------------
-- ** The S-box transformation
-----------------------------------------------------------------------------

-- | The values of the AES S-box table. Note that we describe the S-box programmatically
-- using the mathematical construction given in Section 5.1.1 of the standard. However,
-- the code-generation will turn this into a mere look-up table, as it is just a
-- constant table, all computation being done at \"compile-time\".
sboxTable :: [GF28]
sboxTable :: [GF28]
sboxTable = [GF28 -> GF28
xformByte (GF28 -> GF28
gf28Inverse GF28
b) | GF28
b <- [GF28
0 .. GF28
255]]
  where xformByte :: GF28 -> GF28
        xformByte :: GF28 -> GF28
xformByte GF28
b = (GF28 -> GF28 -> GF28) -> GF28 -> [GF28] -> GF28
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr GF28 -> GF28 -> GF28
forall a. Bits a => a -> a -> a
xor GF28
0x63 [GF28
b GF28 -> Int -> GF28
forall a. Bits a => a -> Int -> a
`rotateR` Int
i | Int
i <- [Int
0, Int
4, Int
5, Int
6, Int
7]]

-- | The sbox transformation. We simply select from the sbox table. Note that we
-- are obliged to give a default value (here @0@) to be used if the index is out-of-bounds
-- as required by SBV's 'select' function. However, that will never happen since
-- the table has all 256 elements in it.
sbox :: GF28 -> GF28
sbox :: GF28 -> GF28
sbox = [GF28] -> GF28 -> GF28 -> GF28
forall a b.
(Mergeable a, Ord b, SymVal b, Num b) =>
[a] -> a -> SBV b -> a
select [GF28]
sboxTable GF28
0

-----------------------------------------------------------------------------
-- ** The inverse S-box transformation
-----------------------------------------------------------------------------

-- | The values of the inverse S-box table. Again, the construction is programmatic.
unSBoxTable :: [GF28]
unSBoxTable :: [GF28]
unSBoxTable = [GF28 -> GF28
gf28Inverse (GF28 -> GF28
xformByte GF28
b) | GF28
b <- [GF28
0 .. GF28
255]]
  where xformByte :: GF28 -> GF28
        xformByte :: GF28 -> GF28
xformByte GF28
b = (GF28 -> GF28 -> GF28) -> GF28 -> [GF28] -> GF28
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr GF28 -> GF28 -> GF28
forall a. Bits a => a -> a -> a
xor GF28
0x05 [GF28
b GF28 -> Int -> GF28
forall a. Bits a => a -> Int -> a
`rotateR` Int
i | Int
i <- [Int
2, Int
5, Int
7]]

-- | The inverse s-box transformation.
unSBox :: GF28 -> GF28
unSBox :: GF28 -> GF28
unSBox = [GF28] -> GF28 -> GF28 -> GF28
forall a b.
(Mergeable a, Ord b, SymVal b, Num b) =>
[a] -> a -> SBV b -> a
select [GF28]
unSBoxTable GF28
0

-- | Prove that the 'sbox' and 'unSBox' are inverses. We have:
--
-- >>> prove sboxInverseCorrect
-- Q.E.D.
--
sboxInverseCorrect :: GF28 -> SBool
sboxInverseCorrect :: GF28 -> SBool
sboxInverseCorrect GF28
x = GF28 -> GF28
unSBox (GF28 -> GF28
sbox GF28
x) GF28 -> GF28 -> SBool
forall a. EqSymbolic a => a -> a -> SBool
.== GF28
x SBool -> SBool -> SBool
.&& GF28 -> GF28
sbox (GF28 -> GF28
unSBox GF28
x) GF28 -> GF28 -> SBool
forall a. EqSymbolic a => a -> a -> SBool
.== GF28
x

-----------------------------------------------------------------------------
-- ** AddRoundKey transformation
-----------------------------------------------------------------------------

-- | Adding the round-key to the current state. We simply exploit the fact
-- that addition is just xor in implementing this transformation.
addRoundKey :: Key -> State -> State
addRoundKey :: State -> State -> State
addRoundKey = (SWord 32 -> SWord 32 -> SWord 32) -> State -> State -> State
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith SWord 32 -> SWord 32 -> SWord 32
forall a. Bits a => a -> a -> a
xor

-----------------------------------------------------------------------------
-- ** Tables for T-Box encryption
-----------------------------------------------------------------------------

-- | T-box table generation function for encryption
t0Func :: GF28 -> [GF28]
t0Func :: GF28 -> [GF28]
t0Func GF28
a = [GF28
s GF28 -> GF28 -> GF28
`gf28Mult` GF28
2, GF28
s, GF28
s, GF28
s GF28 -> GF28 -> GF28
`gf28Mult` GF28
3] where s :: GF28
s = GF28 -> GF28
sbox GF28
a

-- | First look-up table used in encryption
t0 :: GF28 -> SWord 32
t0 :: GF28 -> SWord 32
t0 = State -> SWord 32 -> GF28 -> SWord 32
forall a b.
(Mergeable a, Ord b, SymVal b, Num b) =>
[a] -> a -> SBV b -> a
select State
t0Table SWord 32
0 where t0Table :: State
t0Table = [[GF28] -> SWord 32
forall a. ByteConverter a => [GF28] -> a
fromBytes (GF28 -> [GF28]
t0Func GF28
a)          | GF28
a <- [GF28
0..GF28
255]]

-- | Second look-up table used in encryption
t1 :: GF28 -> SWord 32
t1 :: GF28 -> SWord 32
t1 = State -> SWord 32 -> GF28 -> SWord 32
forall a b.
(Mergeable a, Ord b, SymVal b, Num b) =>
[a] -> a -> SBV b -> a
select State
t1Table SWord 32
0 where t1Table :: State
t1Table = [[GF28] -> SWord 32
forall a. ByteConverter a => [GF28] -> a
fromBytes (GF28 -> [GF28]
t0Func GF28
a [GF28] -> Int -> [GF28]
`rotR` Int
1) | GF28
a <- [GF28
0..GF28
255]]

-- | Third look-up table used in encryption
t2 :: GF28 -> SWord 32
t2 :: GF28 -> SWord 32
t2 = State -> SWord 32 -> GF28 -> SWord 32
forall a b.
(Mergeable a, Ord b, SymVal b, Num b) =>
[a] -> a -> SBV b -> a
select State
t2Table SWord 32
0 where t2Table :: State
t2Table = [[GF28] -> SWord 32
forall a. ByteConverter a => [GF28] -> a
fromBytes (GF28 -> [GF28]
t0Func GF28
a [GF28] -> Int -> [GF28]
`rotR` Int
2) | GF28
a <- [GF28
0..GF28
255]]

-- | Fourth look-up table used in encryption
t3 :: GF28 -> SWord 32
t3 :: GF28 -> SWord 32
t3 = State -> SWord 32 -> GF28 -> SWord 32
forall a b.
(Mergeable a, Ord b, SymVal b, Num b) =>
[a] -> a -> SBV b -> a
select State
t3Table SWord 32
0 where t3Table :: State
t3Table = [[GF28] -> SWord 32
forall a. ByteConverter a => [GF28] -> a
fromBytes (GF28 -> [GF28]
t0Func GF28
a [GF28] -> Int -> [GF28]
`rotR` Int
3) | GF28
a <- [GF28
0..GF28
255]]

-----------------------------------------------------------------------------
-- ** Tables for T-Box decryption
-----------------------------------------------------------------------------

-- | T-box table generating function for decryption
u0Func :: GF28 -> [GF28]
u0Func :: GF28 -> [GF28]
u0Func GF28
a = [GF28
s GF28 -> GF28 -> GF28
`gf28Mult` GF28
0xE, GF28
s GF28 -> GF28 -> GF28
`gf28Mult` GF28
0x9, GF28
s GF28 -> GF28 -> GF28
`gf28Mult` GF28
0xD, GF28
s GF28 -> GF28 -> GF28
`gf28Mult` GF28
0xB] where s :: GF28
s = GF28 -> GF28
unSBox GF28
a

-- | First look-up table used in decryption
u0 :: GF28 -> SWord 32
u0 :: GF28 -> SWord 32
u0 = State -> SWord 32 -> GF28 -> SWord 32
forall a b.
(Mergeable a, Ord b, SymVal b, Num b) =>
[a] -> a -> SBV b -> a
select State
t0Table SWord 32
0 where t0Table :: State
t0Table = [[GF28] -> SWord 32
forall a. ByteConverter a => [GF28] -> a
fromBytes (GF28 -> [GF28]
u0Func GF28
a)          | GF28
a <- [GF28
0..GF28
255]]

-- | Second look-up table used in decryption
u1 :: GF28 -> SWord 32
u1 :: GF28 -> SWord 32
u1 = State -> SWord 32 -> GF28 -> SWord 32
forall a b.
(Mergeable a, Ord b, SymVal b, Num b) =>
[a] -> a -> SBV b -> a
select State
t1Table SWord 32
0 where t1Table :: State
t1Table = [[GF28] -> SWord 32
forall a. ByteConverter a => [GF28] -> a
fromBytes (GF28 -> [GF28]
u0Func GF28
a [GF28] -> Int -> [GF28]
`rotR` Int
1) | GF28
a <- [GF28
0..GF28
255]]

-- | Third look-up table used in decryption
u2 :: GF28 -> SWord 32
u2 :: GF28 -> SWord 32
u2 = State -> SWord 32 -> GF28 -> SWord 32
forall a b.
(Mergeable a, Ord b, SymVal b, Num b) =>
[a] -> a -> SBV b -> a
select State
t2Table SWord 32
0 where t2Table :: State
t2Table = [[GF28] -> SWord 32
forall a. ByteConverter a => [GF28] -> a
fromBytes (GF28 -> [GF28]
u0Func GF28
a [GF28] -> Int -> [GF28]
`rotR` Int
2) | GF28
a <- [GF28
0..GF28
255]]

-- | Fourth look-up table used in decryption
u3 :: GF28 -> SWord 32
u3 :: GF28 -> SWord 32
u3 = State -> SWord 32 -> GF28 -> SWord 32
forall a b.
(Mergeable a, Ord b, SymVal b, Num b) =>
[a] -> a -> SBV b -> a
select State
t3Table SWord 32
0 where t3Table :: State
t3Table = [[GF28] -> SWord 32
forall a. ByteConverter a => [GF28] -> a
fromBytes (GF28 -> [GF28]
u0Func GF28
a [GF28] -> Int -> [GF28]
`rotR` Int
3) | GF28
a <- [GF28
0..GF28
255]]

-----------------------------------------------------------------------------
-- ** AES rounds
-----------------------------------------------------------------------------

-- | Generic round function. Given the function to perform one round, a key-schedule,
-- and a starting state, it performs the AES rounds.
doRounds :: (Bool -> State -> Key -> State) -> KS -> State -> State
doRounds :: (Bool -> State -> State -> State) -> KS -> State -> State
doRounds Bool -> State -> State -> State
rnd (State
ikey, [State]
rkeys, State
fkey) State
sIn = Bool -> State -> State -> State
rnd Bool
True ([State] -> State
forall a. [a] -> a
last [State]
rs) State
fkey
  where s0 :: State
s0 = State
ikey State -> State -> State
`addRoundKey` State
sIn
        rs :: [State]
rs = State
s0 State -> [State] -> [State]
forall a. a -> [a] -> [a]
: [Bool -> State -> State -> State
rnd Bool
False State
s State
k | State
s <- [State]
rs | State
k <- [State]
rkeys ]

-- | One encryption round. The first argument indicates whether this is the final round
-- or not, in which case the construction is slightly different.
aesRound :: Bool -> State -> Key -> State
aesRound :: Bool -> State -> State -> State
aesRound Bool
isFinal State
s State
key = State
d State -> State -> State
`addRoundKey` State
key
  where d :: State
d = (Int -> SWord 32) -> [Int] -> State
forall a b. (a -> b) -> [a] -> [b]
map (Bool -> Int -> SWord 32
f Bool
isFinal) [Int
0..Int
3]
        a :: [[GF28]]
a = (SWord 32 -> [GF28]) -> State -> [[GF28]]
forall a b. (a -> b) -> [a] -> [b]
map SWord 32 -> [GF28]
forall a. ByteConverter a => a -> [GF28]
toBytes State
s
        f :: Bool -> Int -> SWord 32
f Bool
True Int
j = [GF28] -> SWord 32
forall a. ByteConverter a => [GF28] -> a
fromBytes [ GF28 -> GF28
sbox ([[GF28]]
a [[GF28]] -> Int -> [GF28]
forall a. [a] -> Int -> a
!! ((Int
jInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
0) Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` Int
4) [GF28] -> Int -> GF28
forall a. [a] -> Int -> a
!! Int
0)
                             , GF28 -> GF28
sbox ([[GF28]]
a [[GF28]] -> Int -> [GF28]
forall a. [a] -> Int -> a
!! ((Int
jInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` Int
4) [GF28] -> Int -> GF28
forall a. [a] -> Int -> a
!! Int
1)
                             , GF28 -> GF28
sbox ([[GF28]]
a [[GF28]] -> Int -> [GF28]
forall a. [a] -> Int -> a
!! ((Int
jInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
2) Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` Int
4) [GF28] -> Int -> GF28
forall a. [a] -> Int -> a
!! Int
2)
                             , GF28 -> GF28
sbox ([[GF28]]
a [[GF28]] -> Int -> [GF28]
forall a. [a] -> Int -> a
!! ((Int
jInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
3) Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` Int
4) [GF28] -> Int -> GF28
forall a. [a] -> Int -> a
!! Int
3)
                             ]
        f Bool
False Int
j = SWord 32
e0 SWord 32 -> SWord 32 -> SWord 32
forall a. Bits a => a -> a -> a
`xor` SWord 32
e1 SWord 32 -> SWord 32 -> SWord 32
forall a. Bits a => a -> a -> a
`xor` SWord 32
e2 SWord 32 -> SWord 32 -> SWord 32
forall a. Bits a => a -> a -> a
`xor` SWord 32
e3
              where e0 :: SWord 32
e0 = GF28 -> SWord 32
t0 ([[GF28]]
a [[GF28]] -> Int -> [GF28]
forall a. [a] -> Int -> a
!! ((Int
jInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
0) Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` Int
4) [GF28] -> Int -> GF28
forall a. [a] -> Int -> a
!! Int
0)
                    e1 :: SWord 32
e1 = GF28 -> SWord 32
t1 ([[GF28]]
a [[GF28]] -> Int -> [GF28]
forall a. [a] -> Int -> a
!! ((Int
jInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` Int
4) [GF28] -> Int -> GF28
forall a. [a] -> Int -> a
!! Int
1)
                    e2 :: SWord 32
e2 = GF28 -> SWord 32
t2 ([[GF28]]
a [[GF28]] -> Int -> [GF28]
forall a. [a] -> Int -> a
!! ((Int
jInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
2) Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` Int
4) [GF28] -> Int -> GF28
forall a. [a] -> Int -> a
!! Int
2)
                    e3 :: SWord 32
e3 = GF28 -> SWord 32
t3 ([[GF28]]
a [[GF28]] -> Int -> [GF28]
forall a. [a] -> Int -> a
!! ((Int
jInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
3) Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` Int
4) [GF28] -> Int -> GF28
forall a. [a] -> Int -> a
!! Int
3)

-- | One decryption round. Similar to the encryption round, the first argument
-- indicates whether this is the final round or not.
aesInvRound :: Bool -> State -> Key -> State
aesInvRound :: Bool -> State -> State -> State
aesInvRound Bool
isFinal State
s State
key = State
d State -> State -> State
`addRoundKey` State
key
  where d :: State
d = (Int -> SWord 32) -> [Int] -> State
forall a b. (a -> b) -> [a] -> [b]
map (Bool -> Int -> SWord 32
f Bool
isFinal) [Int
0..Int
3]
        a :: [[GF28]]
a = (SWord 32 -> [GF28]) -> State -> [[GF28]]
forall a b. (a -> b) -> [a] -> [b]
map SWord 32 -> [GF28]
forall a. ByteConverter a => a -> [GF28]
toBytes State
s
        f :: Bool -> Int -> SWord 32
f Bool
True Int
j = [GF28] -> SWord 32
forall a. ByteConverter a => [GF28] -> a
fromBytes [ GF28 -> GF28
unSBox ([[GF28]]
a [[GF28]] -> Int -> [GF28]
forall a. [a] -> Int -> a
!! ((Int
jInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
0) Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` Int
4) [GF28] -> Int -> GF28
forall a. [a] -> Int -> a
!! Int
0)
                             , GF28 -> GF28
unSBox ([[GF28]]
a [[GF28]] -> Int -> [GF28]
forall a. [a] -> Int -> a
!! ((Int
jInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
3) Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` Int
4) [GF28] -> Int -> GF28
forall a. [a] -> Int -> a
!! Int
1)
                             , GF28 -> GF28
unSBox ([[GF28]]
a [[GF28]] -> Int -> [GF28]
forall a. [a] -> Int -> a
!! ((Int
jInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
2) Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` Int
4) [GF28] -> Int -> GF28
forall a. [a] -> Int -> a
!! Int
2)
                             , GF28 -> GF28
unSBox ([[GF28]]
a [[GF28]] -> Int -> [GF28]
forall a. [a] -> Int -> a
!! ((Int
jInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` Int
4) [GF28] -> Int -> GF28
forall a. [a] -> Int -> a
!! Int
3)
                             ]
        f Bool
False Int
j = SWord 32
e0 SWord 32 -> SWord 32 -> SWord 32
forall a. Bits a => a -> a -> a
`xor` SWord 32
e1 SWord 32 -> SWord 32 -> SWord 32
forall a. Bits a => a -> a -> a
`xor` SWord 32
e2 SWord 32 -> SWord 32 -> SWord 32
forall a. Bits a => a -> a -> a
`xor` SWord 32
e3
              where e0 :: SWord 32
e0 = GF28 -> SWord 32
u0 ([[GF28]]
a [[GF28]] -> Int -> [GF28]
forall a. [a] -> Int -> a
!! ((Int
jInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
0) Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` Int
4) [GF28] -> Int -> GF28
forall a. [a] -> Int -> a
!! Int
0)
                    e1 :: SWord 32
e1 = GF28 -> SWord 32
u1 ([[GF28]]
a [[GF28]] -> Int -> [GF28]
forall a. [a] -> Int -> a
!! ((Int
jInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
3) Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` Int
4) [GF28] -> Int -> GF28
forall a. [a] -> Int -> a
!! Int
1)
                    e2 :: SWord 32
e2 = GF28 -> SWord 32
u2 ([[GF28]]
a [[GF28]] -> Int -> [GF28]
forall a. [a] -> Int -> a
!! ((Int
jInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
2) Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` Int
4) [GF28] -> Int -> GF28
forall a. [a] -> Int -> a
!! Int
2)
                    e3 :: SWord 32
e3 = GF28 -> SWord 32
u3 ([[GF28]]
a [[GF28]] -> Int -> [GF28]
forall a. [a] -> Int -> a
!! ((Int
jInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` Int
4) [GF28] -> Int -> GF28
forall a. [a] -> Int -> a
!! Int
3)

-----------------------------------------------------------------------------
-- * AES API
-----------------------------------------------------------------------------

-- | Key schedule. Given a 128, 192, or 256 bit key, expand it to get key-schedules
-- for encryption and decryption. The key is given as a sequence of 32-bit words.
-- (4 elements for 128-bits, 6 for 192, and 8 for 256.)
aesKeySchedule :: Key -> (KS, KS)
aesKeySchedule :: State -> (KS, KS)
aesKeySchedule State
key
  | Int
nk Int -> [Int] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [Int
4, Int
6, Int
8]
  = (KS
encKS, KS
decKS)
  | Bool
True
  = [Char] -> (KS, KS)
forall a. HasCallStack => [Char] -> a
error [Char]
"aesKeySchedule: Invalid key size"
  where nk :: Int
nk = State -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length State
key
        nr :: Int
nr = Int
nk Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
6
        encKS :: KS
encKS@(State
f, [State]
m, State
l) = ([State] -> State
forall a. [a] -> a
head [State]
rKeys, Int -> [State] -> [State]
forall a. Int -> [a] -> [a]
take (Int
nrInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1) ([State] -> [State]
forall a. [a] -> [a]
tail [State]
rKeys), [State]
rKeys [State] -> Int -> State
forall a. [a] -> Int -> a
!! Int
nr)
        decKS :: KS
decKS = (State
l, (State -> State) -> [State] -> [State]
forall a b. (a -> b) -> [a] -> [b]
map State -> State
invMixColumns ([State] -> [State]
forall a. [a] -> [a]
reverse [State]
m), State
f)
        rKeys :: [State]
rKeys = Int -> State -> [State]
keyExpansion Int
nk State
key

-- | Block encryption. The first argument is the plain-text, which must have
-- precisely 4 elements, for a total of 128-bits of input. The second
-- argument is the key-schedule to be used, obtained by a call to 'aesKeySchedule'.
-- The output will always have 4 32-bit words, which is the cipher-text.
aesEncrypt :: [SWord 32] -> KS -> [SWord 32]
aesEncrypt :: State -> KS -> State
aesEncrypt State
pt KS
encKS
  | State -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length State
pt Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
4
  = (Bool -> State -> State -> State) -> KS -> State -> State
doRounds Bool -> State -> State -> State
aesRound KS
encKS State
pt
  | Bool
True
  = [Char] -> State
forall a. HasCallStack => [Char] -> a
error [Char]
"aesEncrypt: Invalid plain-text size"

-- | Block decryption. The arguments are the same as in 'aesEncrypt', except
-- the first argument is the cipher-text and the output is the corresponding
-- plain-text.
aesDecrypt :: [SWord 32] -> KS -> [SWord 32]
aesDecrypt :: State -> KS -> State
aesDecrypt State
ct KS
decKS
  | State -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length State
ct Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
4
  = (Bool -> State -> State -> State) -> KS -> State -> State
doRounds Bool -> State -> State -> State
aesInvRound KS
decKS State
ct
  | Bool
True
  = [Char] -> State
forall a. HasCallStack => [Char] -> a
error [Char]
"aesDecrypt: Invalid cipher-text size"

-----------------------------------------------------------------------------
-- * Test vectors
-----------------------------------------------------------------------------

-----------------------------------------------------------------------------
-- ** 128-bit enc/dec test
-----------------------------------------------------------------------------

-- | 128-bit encryption test, from Appendix C.1 of the AES standard:
--
-- >>> map hex8 t128Enc
-- ["69c4e0d8","6a7b0430","d8cdb780","70b4c55a"]
--
t128Enc :: [SWord 32]
t128Enc :: State
t128Enc = State -> KS -> State
aesEncrypt State
pt KS
ks
  where pt :: State
pt  = [SWord 32
0x00112233, SWord 32
0x44556677, SWord 32
0x8899aabb, SWord 32
0xccddeeff]
        key :: State
key = [SWord 32
0x00010203, SWord 32
0x04050607, SWord 32
0x08090a0b, SWord 32
0x0c0d0e0f]
        (KS
ks, KS
_) = State -> (KS, KS)
aesKeySchedule State
key

-- | 128-bit decryption test, from Appendix C.1 of the AES standard:
--
-- >>> map hex8 t128Dec
-- ["00112233","44556677","8899aabb","ccddeeff"]
--
t128Dec :: [SWord 32]
t128Dec :: State
t128Dec = State -> KS -> State
aesDecrypt State
ct KS
ks
  where ct :: State
ct  = [SWord 32
0x69c4e0d8, SWord 32
0x6a7b0430, SWord 32
0xd8cdb780, SWord 32
0x70b4c55a]
        key :: State
key = [SWord 32
0x00010203, SWord 32
0x04050607, SWord 32
0x08090a0b, SWord 32
0x0c0d0e0f]
        (KS
_, KS
ks) = State -> (KS, KS)
aesKeySchedule State
key

-----------------------------------------------------------------------------
-- ** 192-bit enc/dec test
-----------------------------------------------------------------------------

-- | 192-bit encryption test, from Appendix C.2 of the AES standard:
--
-- >>> map hex8 t192Enc
-- ["dda97ca4","864cdfe0","6eaf70a0","ec0d7191"]
--
t192Enc :: [SWord 32]
t192Enc :: State
t192Enc = State -> KS -> State
aesEncrypt State
pt KS
ks
  where pt :: State
pt  = [SWord 32
0x00112233, SWord 32
0x44556677, SWord 32
0x8899aabb, SWord 32
0xccddeeff]
        key :: State
key = [SWord 32
0x00010203, SWord 32
0x04050607, SWord 32
0x08090a0b, SWord 32
0x0c0d0e0f, SWord 32
0x10111213, SWord 32
0x14151617]
        (KS
ks, KS
_) = State -> (KS, KS)
aesKeySchedule State
key

-- | 192-bit decryption test, from Appendix C.2 of the AES standard:
--
-- >>> map hex8 t192Dec
-- ["00112233","44556677","8899aabb","ccddeeff"]
--
t192Dec :: [SWord 32]
t192Dec :: State
t192Dec = State -> KS -> State
aesDecrypt State
ct KS
ks
  where ct :: State
ct  = [SWord 32
0xdda97ca4, SWord 32
0x864cdfe0, SWord 32
0x6eaf70a0, SWord 32
0xec0d7191]
        key :: State
key = [SWord 32
0x00010203, SWord 32
0x04050607, SWord 32
0x08090a0b, SWord 32
0x0c0d0e0f, SWord 32
0x10111213, SWord 32
0x14151617]
        (KS
_, KS
ks) = State -> (KS, KS)
aesKeySchedule State
key

-----------------------------------------------------------------------------
-- ** 256-bit enc/dec test
-----------------------------------------------------------------------------

-- | 256-bit encryption, from Appendix C.3 of the AES standard:
--
-- >>> map hex8 t256Enc
-- ["8ea2b7ca","516745bf","eafc4990","4b496089"]
--
t256Enc :: [SWord 32]
t256Enc :: State
t256Enc = State -> KS -> State
aesEncrypt State
pt KS
ks
  where pt :: State
pt  = [SWord 32
0x00112233, SWord 32
0x44556677, SWord 32
0x8899aabb, SWord 32
0xccddeeff]
        key :: State
key = [SWord 32
0x00010203, SWord 32
0x04050607, SWord 32
0x08090a0b, SWord 32
0x0c0d0e0f, SWord 32
0x10111213, SWord 32
0x14151617, SWord 32
0x18191a1b, SWord 32
0x1c1d1e1f]
        (KS
ks, KS
_) = State -> (KS, KS)
aesKeySchedule State
key

-- | 256-bit decryption, from Appendix C.3 of the AES standard:
--
-- >>> map hex8 t256Dec
-- ["00112233","44556677","8899aabb","ccddeeff"]
--
t256Dec :: [SWord 32]
t256Dec :: State
t256Dec = State -> KS -> State
aesDecrypt State
ct KS
ks
  where ct :: State
ct  = [SWord 32
0x8ea2b7ca, SWord 32
0x516745bf, SWord 32
0xeafc4990, SWord 32
0x4b496089]
        key :: State
key = [SWord 32
0x00010203, SWord 32
0x04050607, SWord 32
0x08090a0b, SWord 32
0x0c0d0e0f, SWord 32
0x10111213, SWord 32
0x14151617, SWord 32
0x18191a1b, SWord 32
0x1c1d1e1f]
        (KS
_, KS
ks) = State -> (KS, KS)
aesKeySchedule State
key


-----------------------------------------------------------------------------
-- * Verification
-- ${verifIntro}
-----------------------------------------------------------------------------
{- $verifIntro
  While SMT based technologies can prove correct many small properties fairly quickly, it would
  be naive for them to automatically verify that our AES implementation is correct. (By correct,
  we mean decryption follewed by encryption yielding the same result.) However, we can state
  this property precisely using SBV, and use quick-check to gain some confidence.
-}

-- | Correctness theorem for 128-bit AES. Ideally, we would run:
--
-- @
--   prove aes128IsCorrect
-- @
--
-- to get a proof automatically. Unfortunately, while SBV will successfully generate the proof
-- obligation for this theorem and ship it to the SMT solver, it would be naive to expect the SMT-solver
-- to finish that proof in any reasonable time with the currently available SMT solving technologies.
-- Instead, we can issue:
--
-- @
--   quickCheck aes128IsCorrect
-- @
-- 
-- and get some degree of confidence in our code. Similar predicates can be easily constructed for 192, and
-- 256 bit cases as well.
aes128IsCorrect :: (SWord 32, SWord 32, SWord 32, SWord 32)  -- ^ plain-text words
                -> (SWord 32, SWord 32, SWord 32, SWord 32)  -- ^ key-words
                -> SBool                                 -- ^ True if round-trip gives us plain-text back
aes128IsCorrect :: (SWord 32, SWord 32, SWord 32, SWord 32)
-> (SWord 32, SWord 32, SWord 32, SWord 32) -> SBool
aes128IsCorrect (SWord 32
i0, SWord 32
i1, SWord 32
i2, SWord 32
i3) (SWord 32
k0, SWord 32
k1, SWord 32
k2, SWord 32
k3) = State
pt State -> State -> SBool
forall a. EqSymbolic a => a -> a -> SBool
.== State
pt'
   where pt :: State
pt  = [SWord 32
i0, SWord 32
i1, SWord 32
i2, SWord 32
i3]
         key :: State
key = [SWord 32
k0, SWord 32
k1, SWord 32
k2, SWord 32
k3]
         (KS
encKS, KS
decKS) = State -> (KS, KS)
aesKeySchedule State
key
         ct :: State
ct  = State -> KS -> State
aesEncrypt State
pt KS
encKS
         pt' :: State
pt' = State -> KS -> State
aesDecrypt State
ct KS
decKS

-----------------------------------------------------------------------------
-- * Code generation
-- ${codeGenIntro}
-----------------------------------------------------------------------------
{- $codeGenIntro
   We have emphasized that our T-Box implementation in Haskell was guided by clarity and correctness, not
   performance. Indeed, our implementation is hardly the fastest AES implementation in Haskell. However,
   we can use it to automatically generate straight-line C-code that can run fairly fast.

   For the purposes of illustration, we only show here how to generate code for a 128-bit AES block-encrypt
   function, that takes 8 32-bit words as an argument. The first 4 are the 128-bit input, and the final
   four are the 128-bit key. The impact of this is that the generated function would expand the key for
   each block of encryption, a needless task unless we change the key in every block. In a more serious application,
   we would instead generate code for both the 'aesKeySchedule' and the 'aesEncrypt' functions, thus reusing the
   key-schedule over many applications of the encryption call. (Unfortunately doing this is rather cumbersome right
   now, since Haskell does not support fixed-size lists.)
-}

-- | Code generation for 128-bit AES encryption.
--
-- The following sample from the generated code-lines show how T-Boxes are rendered as C arrays:
--
-- @
--   static const SWord32 table1[] = {
--       0xc66363a5UL, 0xf87c7c84UL, 0xee777799UL, 0xf67b7b8dUL,
--       0xfff2f20dUL, 0xd66b6bbdUL, 0xde6f6fb1UL, 0x91c5c554UL,
--       0x60303050UL, 0x02010103UL, 0xce6767a9UL, 0x562b2b7dUL,
--       0xe7fefe19UL, 0xb5d7d762UL, 0x4dababe6UL, 0xec76769aUL,
--       ...
--       }
-- @
--
-- The generated program has 5 tables (one sbox table, and 4-Tboxes), all converted to fast C arrays. Here
-- is a sample of the generated straightline C-code:
--
-- @
--   const SWord8  s1915 = (SWord8) s1912;
--   const SWord8  s1916 = table0[s1915];
--   const SWord16 s1917 = (((SWord16) s1914) << 8) | ((SWord16) s1916);
--   const SWord32 s1918 = (((SWord32) s1911) << 16) | ((SWord32) s1917);
--   const SWord32 s1919 = s1844 ^ s1918;
--   const SWord32 s1920 = s1903 ^ s1919;
-- @
--
-- The GNU C-compiler does a fine job of optimizing this straightline code to generate a fairly efficient C implementation.
cgAES128BlockEncrypt :: IO ()
cgAES128BlockEncrypt :: IO ()
cgAES128BlockEncrypt = Maybe [Char] -> [Char] -> SBVCodeGen () -> IO ()
forall a. Maybe [Char] -> [Char] -> SBVCodeGen a -> IO a
compileToC Maybe [Char]
forall a. Maybe a
Nothing [Char]
"aes128BlockEncrypt" (SBVCodeGen () -> IO ()) -> SBVCodeGen () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
        State
pt  <- Int -> [Char] -> SBVCodeGen State
forall a. SymVal a => Int -> [Char] -> SBVCodeGen [SBV a]
cgInputArr Int
4 [Char]
"pt"        -- plain-text as an array of 4 Word32's
        State
key <- Int -> [Char] -> SBVCodeGen State
forall a. SymVal a => Int -> [Char] -> SBVCodeGen [SBV a]
cgInputArr Int
4 [Char]
"key"       -- key as an array of 4 Word32s
        -- Use the test values from Appendix C.1 of the AES standard as the driver values
        [Integer] -> SBVCodeGen ()
cgSetDriverValues ([Integer] -> SBVCodeGen ()) -> [Integer] -> SBVCodeGen ()
forall a b. (a -> b) -> a -> b
$    [Integer
0x00112233, Integer
0x44556677, Integer
0x8899aabb, Integer
0xccddeeff]
                            [Integer] -> [Integer] -> [Integer]
forall a. [a] -> [a] -> [a]
++ [Integer
0x00010203, Integer
0x04050607, Integer
0x08090a0b, Integer
0x0c0d0e0f]
        let (KS
encKs, KS
_) = State -> (KS, KS)
aesKeySchedule State
key
        [Char] -> State -> SBVCodeGen ()
forall a. SymVal a => [Char] -> [SBV a] -> SBVCodeGen ()
cgOutputArr [Char]
"ct" (State -> SBVCodeGen ()) -> State -> SBVCodeGen ()
forall a b. (a -> b) -> a -> b
$ State -> KS -> State
aesEncrypt State
pt KS
encKs

-----------------------------------------------------------------------------
-- * C-library generation
-- ${libraryIntro}
-----------------------------------------------------------------------------
{- $libraryIntro
   The 'cgAES128BlockEncrypt' example shows how to generate code for 128-bit AES encryption. As the generated
   function performs encryption on a given block, it performs key expansion as necessary. However, this is
   not quite practical: We would like to expand the key only once, and encrypt the stream of plain-text blocks using
   the same expanded key (potentially using some crypto-mode), until we decide to change the key. In this
   section, we show how to use SBV to instead generate a library of functions that can be used in such a scenario.
   The generated library is a typical @.a@ archive, that can be linked using the C-compiler as usual.
-}

-- | Components of the AES implementation that the library is generated from
aesLibComponents :: Int -> [(String, SBVCodeGen ())]
aesLibComponents :: Int -> [([Char], SBVCodeGen ())]
aesLibComponents Int
sz = [ ([Char]
"aes" [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Int -> [Char]
forall a. Show a => a -> [Char]
show Int
sz [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
"KeySchedule",  SBVCodeGen ()
keySchedule)
                      , ([Char]
"aes" [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Int -> [Char]
forall a. Show a => a -> [Char]
show Int
sz [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
"BlockEncrypt", SBVCodeGen ()
enc)
                      , ([Char]
"aes" [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Int -> [Char]
forall a. Show a => a -> [Char]
show Int
sz [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
"BlockDecrypt", SBVCodeGen ()
dec)
                      ]
  where -- key-schedule
        nk :: Int
nk
         | Int
sz Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
128 = Int
4
         | Int
sz Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
192 = Int
6
         | Int
sz Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
256 = Int
8
         | Bool
True      = [Char] -> Int
forall a. HasCallStack => [Char] -> a
error ([Char] -> Int) -> [Char] -> Int
forall a b. (a -> b) -> a -> b
$ [Char]
"aesLibComponents: Size must be one of 128, 192, or 256; received: " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Int -> [Char]
forall a. Show a => a -> [Char]
show Int
sz
        -- We get 4*(nr+1) keys, where nr = nk + 6
        nr :: Int
nr = Int
nk Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
6
        xk :: Int
xk = Int
4 Int -> Int -> Int
forall a. Num a => a -> a -> a
* (Int
nr Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
        keySchedule :: SBVCodeGen ()
keySchedule = do State
key <- Int -> [Char] -> SBVCodeGen State
forall a. SymVal a => Int -> [Char] -> SBVCodeGen [SBV a]
cgInputArr Int
nk [Char]
"key"     -- key
                         let (KS
encKS, KS
decKS) = State -> (KS, KS)
aesKeySchedule State
key
                         [Char] -> State -> SBVCodeGen ()
forall a. SymVal a => [Char] -> [SBV a] -> SBVCodeGen ()
cgOutputArr [Char]
"encKS" (KS -> State
ksToXKey KS
encKS)
                         [Char] -> State -> SBVCodeGen ()
forall a. SymVal a => [Char] -> [SBV a] -> SBVCodeGen ()
cgOutputArr [Char]
"decKS" (KS -> State
ksToXKey KS
decKS)
        -- encryption
        enc :: SBVCodeGen ()
enc = do State
pt   <- Int -> [Char] -> SBVCodeGen State
forall a. SymVal a => Int -> [Char] -> SBVCodeGen [SBV a]
cgInputArr Int
4  [Char]
"pt"    -- plain-text
                 State
xkey <- Int -> [Char] -> SBVCodeGen State
forall a. SymVal a => Int -> [Char] -> SBVCodeGen [SBV a]
cgInputArr Int
xk [Char]
"xkey"  -- expanded key
                 [Char] -> State -> SBVCodeGen ()
forall a. SymVal a => [Char] -> [SBV a] -> SBVCodeGen ()
cgOutputArr [Char]
"ct" (State -> SBVCodeGen ()) -> State -> SBVCodeGen ()
forall a b. (a -> b) -> a -> b
$ State -> KS -> State
aesEncrypt State
pt (State -> KS
xkeyToKS State
xkey)
        -- decryption
        dec :: SBVCodeGen ()
dec = do State
pt   <- Int -> [Char] -> SBVCodeGen State
forall a. SymVal a => Int -> [Char] -> SBVCodeGen [SBV a]
cgInputArr Int
4  [Char]
"ct"    -- cipher-text
                 State
xkey <- Int -> [Char] -> SBVCodeGen State
forall a. SymVal a => Int -> [Char] -> SBVCodeGen [SBV a]
cgInputArr Int
xk [Char]
"xkey"  -- expanded key
                 [Char] -> State -> SBVCodeGen ()
forall a. SymVal a => [Char] -> [SBV a] -> SBVCodeGen ()
cgOutputArr [Char]
"pt" (State -> SBVCodeGen ()) -> State -> SBVCodeGen ()
forall a b. (a -> b) -> a -> b
$ State -> KS -> State
aesDecrypt State
pt (State -> KS
xkeyToKS State
xkey)
        -- Transforming back and forth from our KS type to a flat array used by the generated C code
        -- Turn a series of expanded keys to our internal KS type
        xkeyToKS :: [SWord 32] -> KS
        xkeyToKS :: State -> KS
xkeyToKS State
xs = (State
f, [State]
m, State
l)
           where f :: State
f  = Int -> State -> State
forall a. Int -> [a] -> [a]
take Int
4 State
xs                             -- first round key
                 m :: [State]
m  = State -> [State]
forall a. [a] -> [[a]]
chop4 (Int -> State -> State
forall a. Int -> [a] -> [a]
take (Int
xk Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
8) (Int -> State -> State
forall a. Int -> [a] -> [a]
drop Int
4 State
xs))     -- middle rounds
                 l :: State
l  = Int -> State -> State
forall a. Int -> [a] -> [a]
drop (Int
xk Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
4) State
xs                      -- last round key
        -- Turn a KS to a series of expanded key words
        ksToXKey :: KS -> [SWord 32]
        ksToXKey :: KS -> State
ksToXKey (State
f, [State]
m, State
l) = State
f State -> State -> State
forall a. [a] -> [a] -> [a]
++ [State] -> State
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [State]
m State -> State -> State
forall a. [a] -> [a] -> [a]
++ State
l
        -- chunk in fours. (This function must be in some standard library, where?)
        chop4 :: [a] -> [[a]]
        chop4 :: [a] -> [[a]]
chop4 [] = []
        chop4 [a]
xs = let ([a]
f, [a]
r) = Int -> [a] -> ([a], [a])
forall a. Int -> [a] -> ([a], [a])
splitAt Int
4 [a]
xs in [a]
f [a] -> [[a]] -> [[a]]
forall a. a -> [a] -> [a]
: [a] -> [[a]]
forall a. [a] -> [[a]]
chop4 [a]
r

-- | Generate code for AES functionality; given the key size.
cgAESLibrary :: Int -> Maybe FilePath -> IO ()
cgAESLibrary :: Int -> Maybe [Char] -> IO ()
cgAESLibrary Int
sz Maybe [Char]
mbd
  | Int
sz Int -> [Int] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [Int
128, Int
192, Int
256] = IO [()] -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO [()] -> IO ()) -> IO [()] -> IO ()
forall a b. (a -> b) -> a -> b
$ Maybe [Char] -> [Char] -> [([Char], SBVCodeGen ())] -> IO [()]
forall a.
Maybe [Char] -> [Char] -> [([Char], SBVCodeGen a)] -> IO [a]
compileToCLib Maybe [Char]
mbd [Char]
nm (Int -> [([Char], SBVCodeGen ())]
aesLibComponents Int
sz)
  | Bool
True                      = [Char] -> IO ()
forall a. HasCallStack => [Char] -> a
error ([Char] -> IO ()) -> [Char] -> IO ()
forall a b. (a -> b) -> a -> b
$ [Char]
"cgAESLibrary: Size must be one of 128, 192, or 256, received: " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Int -> [Char]
forall a. Show a => a -> [Char]
show Int
sz
  where nm :: [Char]
nm = [Char]
"aes" [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Int -> [Char]
forall a. Show a => a -> [Char]
show Int
sz [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
"Lib"

-- | Generate a C library, containing functions for performing 128-bit enc/dec/key-expansion.
-- A note on performance: In a very rough speed test, the generated code was able to do
-- 6.3 million block encryptions per second on a decent MacBook Pro. On the same machine, OpenSSL
-- reports 8.2 million block encryptions per second. So, the generated code is about 25% slower
-- as compared to the highly optimized OpenSSL implementation. (Note that the speed test was done
-- somewhat simplistically, so these numbers should be considered very rough estimates.)
cgAES128Library :: IO ()
cgAES128Library :: IO ()
cgAES128Library = Int -> Maybe [Char] -> IO ()
cgAESLibrary Int
128 Maybe [Char]
forall a. Maybe a
Nothing

--------------------------------------------------------------------------------------------
-- | For doctest purposes only
hex8 :: (SymVal a, Show a, Integral a) => SBV a -> String
hex8 :: SBV a -> [Char]
hex8 SBV a
v = Int -> Char -> [Char]
forall a. Int -> a -> [a]
replicate (Int
8 Int -> Int -> Int
forall a. Num a => a -> a -> a
- [Char] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Char]
s) Char
'0' [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
s
  where s :: [Char]
s = (a -> [Char] -> [Char]) -> [Char] -> a -> [Char]
forall a b c. (a -> b -> c) -> b -> a -> c
flip a -> [Char] -> [Char]
forall a. (Integral a, Show a) => a -> [Char] -> [Char]
showHex [Char]
"" (a -> [Char]) -> (SBV a -> a) -> SBV a -> [Char]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Maybe a -> a
forall a. HasCallStack => Maybe a -> a
fromJust (Maybe a -> a) -> (SBV a -> Maybe a) -> SBV a -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SBV a -> Maybe a
forall a. SymVal a => SBV a -> Maybe a
unliteral (SBV a -> [Char]) -> SBV a -> [Char]
forall a b. (a -> b) -> a -> b
$ SBV a
v

{-# ANN aesRound    ("HLint: ignore Use head" :: String) #-}
{-# ANN aesInvRound ("HLint: ignore Use head" :: String) #-}