spec/crypto/aes/aes.bend

bend-collections-laws-crypto@1.0.0.0 · 291 lines · 43 definitions

1import Base
2import bend-collections-laws-crypto@1.0.0.0/src/crypto/aes/types.bend as T
3import bend-collections-laws-crypto@1.0.0.0/spec/crypto/aes/poly.bend as P
4
5# Executable specification of AES (FIPS 197), transcribed from the standard.
6# Only the neutral Quad/State representation is shared with the
7# implementation. A byte is a U32; the finite-field operations read its low
8# eight bits b7..b0 (section 3.2), the bytes of the API are all below 256.
9
10# Section 3.2/4: the byte b7 b6 ... b0 is the polynomial b7 x^7 + ... + b0.
11def bits(x: U32) -> List<&2, Bool>:
12  match x:
13    case U32{WCon{b0, WCon{b1, WCon{b2, WCon{b3, WCon{b4, WCon{b5, WCon{b6, WCon{b7, rest}}}}}}}}}:
14      [b0, b1, b2, b3, b4, b5, b6, b7]
15
16def byte(w: Word(8n)) -> U32:
17  match w:
18    case WCon{b0, WCon{b1, WCon{b2, WCon{b3, WCon{b4, WCon{b5, WCon{b6, WCon{b7, WNil{}}}}}}}}}:
19      U32{WCon{b0, WCon{b1, WCon{b2, WCon{b3, WCon{b4, WCon{b5, WCon{b6, WCon{b7, Word.zero(24n)}}}}}}}}}
20
21# Section 4.2: m(x) = x^8 + x^4 + x^3 + x + 1, so x^8 = x^4 + x^3 + x + 1.
22def low() -> Word(8n):
23  WCon{True{}, WCon{True{}, WCon{False{}, WCon{True{}, WCon{True{}, WCon{False{}, WCon{False{}, WCon{False{}, WNil{}}}}}}}}}
24
25# Section 4.2: the product of two bytes is their polynomial product modulo m(x).
26def mul(+a: U32, +b: U32) -> U32:
27  byte(P.reduce(8n, low(), P.mul(bits(a), bits(b))))
28
29# Section 4.4 / 5.1.1: the multiplicative inverse of a nonzero a in GF(2^8),
30# {00} mapped to itself. The multiplicative group of GF(2^8) has 255
31# elements, so a^255 = {01} and a^-1 = a^254 for a != {00} (and
32# {00}^254 = {00}): computed by squaring, a^254 = a^2 a^4 a^8 ... a^128.
33# (tools/check_aes.py checks this against the defining property a * b = {01}
34# for all 256 bytes.)
35def inverse(+a: U32) -> U32:
36  +a2 = mul(a, a)
37  +a4 = mul(a2, a2)
38  +a8 = mul(a4, a4)
39  +a16 = mul(a8, a8)
40  +a32 = mul(a16, a16)
41  +a64 = mul(a32, a32)
42  +a128 = mul(a64, a64)
43  mul(a2, mul(a4, mul(a8, mul(a16, mul(a32, mul(a64, a128))))))
44
45def nth_bit(xs: List<&2, Bool>, n: Nat) -> Bool:
46  match xs n:
47    case Nil{} _:
48      False{}
49    case x <> rest 0n:
50      x
51    case x <> rest 1n+p:
52      nth_bit(rest, p)
53
54# Section 5.1.1, equation (5.1): b'_i = b_i + b_(i+4 mod 8) + b_(i+5 mod 8)
55#   + b_(i+6 mod 8) + b_(i+7 mod 8) + c_i, with c = {63}.
56def affine_bit(+b: List<&2, Bool>, +c: List<&2, Bool>, +i: Nat) -> Bool:
57  Bool.xor(Bool.xor(Bool.xor(Bool.xor(Bool.xor(nth_bit(b, i),
58    nth_bit(b, Nat.mod(Nat.add(i, 4n), 8n))), nth_bit(b, Nat.mod(Nat.add(i, 5n), 8n))),
59    nth_bit(b, Nat.mod(Nat.add(i, 6n), 8n))), nth_bit(b, Nat.mod(Nat.add(i, 7n), 8n))),
60    nth_bit(c, i))
61
62def affine(+x: U32) -> U32:
63  +b = bits(x)
64  +c = bits(99)
65  byte(WCon{affine_bit(b, c, 0n), WCon{affine_bit(b, c, 1n), WCon{affine_bit(b, c, 2n),
66    WCon{affine_bit(b, c, 3n), WCon{affine_bit(b, c, 4n), WCon{affine_bit(b, c, 5n),
67    WCon{affine_bit(b, c, 6n), WCon{affine_bit(b, c, 7n), WNil{}}}}}}}}})
68
69# Section 5.1.1: the S-box is the inverse followed by the affine transformation.
70def sbox_byte(+x: U32) -> U32:
71  affine(inverse(x))
72
73# The byte c7..c0 (the S-box reads only the low eight bits of a U32).
74def byte8(c0: Bool, c1: Bool, c2: Bool, c3: Bool, c4: Bool, c5: Bool, c6: Bool, c7: Bool) -> U32:
75  U32{WCon{c0, WCon{c1, WCon{c2, WCon{c3, WCon{c4, WCon{c5, WCon{c6, WCon{c7, Word.zero(24n)}}}}}}}}}
76
77# Both arms are the S-box of the byte: the match on bit 7 only keeps a
78# proof about a byte whose bits are unknown from expanding the field
79# arithmetic (the specification is not run in constant time).
80def sbox_at(+c0: Bool, +c1: Bool, +c2: Bool, +c3: Bool, +c4: Bool, +c5: Bool, +c6: Bool, c7: Bool) -> U32:
81  match c7:
82    case False{}: sbox_byte(byte8(c0, c1, c2, c3, c4, c5, c6, False{}))
83    case True{}: sbox_byte(byte8(c0, c1, c2, c3, c4, c5, c6, True{}))
84
85def sbox(x: U32) -> U32:
86  match x:
87    case U32{WCon{c0, WCon{c1, WCon{c2, WCon{c3, WCon{c4, WCon{c5, WCon{c6, WCon{c7, rest}}}}}}}}}:
88      sbox_at(c0, c1, c2, c3, c4, c5, c6, c7)
89
90# s[r,c]: byte r of column c.
91def byte_at(w: T.Quad, r: Nat) -> U32:
92  match w r:
93    case T.W{a, b, c, d} 0n:
94      a
95    case T.W{a, b, c, d} 1n:
96      b
97    case T.W{a, b, c, d} 2n:
98      c
99    case T.W{a, b, c, d} 3n+p:
100      d
101
102def column(s: T.State, c: Nat) -> T.Quad:
103  match s c:
104    case T.S{c0, c1, c2, c3} 0n:
105      c0
106    case T.S{c0, c1, c2, c3} 1n:
107      c1
108    case T.S{c0, c1, c2, c3} 2n:
109      c2
110    case T.S{c0, c1, c2, c3} 3n+p:
111      c3
112
113def at(+s: T.State, +r: Nat, +c: Nat) -> U32:
114  byte_at(column(s, c), r)
115
116# Section 5.1.1: SubBytes applies the S-box to every byte of the state.
117def sub_word(w: T.Quad) -> T.Quad:
118  match w:
119    case T.W{a, b, c, d}:
120      T.W{sbox(a), sbox(b), sbox(c), sbox(d)}
121
122def sub_bytes(s: T.State) -> T.State:
123  match s:
124    case T.S{c0, c1, c2, c3}:
125      T.S{sub_word(c0), sub_word(c1), sub_word(c2), sub_word(c3)}
126
127# Section 5.1.2: s'[r,c] = s[r, (c + shift(r, 4)) mod 4], shift(r, 4) = r.
128def shifted(+s: T.State, +c: Nat) -> T.Quad:
129  T.W{at(s, 0n, c), at(s, 1n, Nat.mod(Nat.add(c, 1n), 4n)),
130    at(s, 2n, Nat.mod(Nat.add(c, 2n), 4n)), at(s, 3n, Nat.mod(Nat.add(c, 3n), 4n))}
131
132def shift_rows(+s: T.State) -> T.State:
133  T.S{shifted(s, 0n), shifted(s, 1n), shifted(s, 2n), shifted(s, 3n)}
134
135# Section 5.1.3, equation (5.6).
136def mix_column(w: T.Quad) -> T.Quad:
137  match w:
138    case T.W{+a, +b, +c, +d}:
139      T.W{U32.xor(U32.xor(U32.xor(mul(2, a), mul(3, b)), c), d),
140        U32.xor(U32.xor(U32.xor(a, mul(2, b)), mul(3, c)), d),
141        U32.xor(U32.xor(U32.xor(a, b), mul(2, c)), mul(3, d)),
142        U32.xor(U32.xor(U32.xor(mul(3, a), b), c), mul(2, d))}
143
144def mix_columns(s: T.State) -> T.State:
145  match s:
146    case T.S{c0, c1, c2, c3}:
147      T.S{mix_column(c0), mix_column(c1), mix_column(c2), mix_column(c3)}
148
149def xor_word(a: T.Quad, b: T.Quad) -> T.Quad:
150  match a b:
151    case T.W{a0, a1, a2, a3} T.W{b0, b1, b2, b3}:
152      T.W{U32.xor(a0, b0), U32.xor(a1, b1), U32.xor(a2, b2), U32.xor(a3, b3)}
153
154# Section 5.1.4: column c is added to the key word w[round*4 + c].
155def add_round_key(s: T.State, w0: T.Quad, w1: T.Quad, w2: T.Quad, w3: T.Quad) -> T.State:
156  match s:
157    case T.S{c0, c1, c2, c3}:
158      T.S{xor_word(c0, w0), xor_word(c1, w1), xor_word(c2, w2), xor_word(c3, w3)}
159
160# Section 5.1, Figure 5. After the initial AddRoundKey, n more full rounds
161# (SubBytes, ShiftRows, MixColumns, AddRoundKey) and the final round (no
162# MixColumns), each taking the next four words of the key schedule.
163def rounds(n: Nat, s: T.State, w: List<&2, T.Quad>) -> T.State:
164  match n w:
165    case 0n w0 <> w1 <> w2 <> w3 <> rest:
166      add_round_key(shift_rows(sub_bytes(s)), w0, w1, w2, w3)
167    case 1n+p w0 <> w1 <> w2 <> w3 <> rest:
168      rounds(p, add_round_key(mix_columns(shift_rows(sub_bytes(s))), w0, w1, w2, w3), rest)
169    case _ _:
170      s
171
172def cipher(+nr: Nat, s: T.State, w: List<&2, T.Quad>) -> T.State:
173  match w:
174    case w0 <> w1 <> w2 <> w3 <> rest:
175      rounds(Nat.sub(nr, 1n), add_round_key(s, w0, w1, w2, w3), rest)
176    case _:
177      s
178
179# Section 5.2: RotWord, SubWord (sub_word above), and Rcon[j] = [x^(j-1), 0, 0, 0].
180def rot_word(w: T.Quad) -> T.Quad:
181  match w:
182    case T.W{a, b, c, d}:
183      T.W{b, c, d, a}
184
185def xpow(n: Nat) -> U32:
186  match n:
187    case 0n:
188      1
189    case 1n+p:
190      mul(2, xpow(p))
191
192def rcon(j: Nat) -> T.Quad:
193  T.W{xpow(Nat.sub(j, 1n)), 0, 0, 0}
194
195# w[i] = [key[4i], key[4i+1], key[4i+2], key[4i+3]] for i < Nk.
196def key_words(key: List<&2, U32>) -> List<&2, T.Quad>:
197  match key:
198    case a <> b <> c <> d <> rest:
199      T.W{a, b, c, d} <> key_words(rest)
200    case _:
201      Nil{}
202
203def nth_word(ws: List<&2, T.Quad>, n: Nat) -> T.Quad:
204  match ws n:
205    case Nil{} _:
206      T.W{0, 0, 0, 0}
207    case w <> rest 0n:
208      w
209    case w <> rest 1n+p:
210      nth_word(rest, p)
211
212# The transformation of temp = w[i-1] in Figure 11, by k = i mod Nk,
213# big = (Nk > 6) and q = i / Nk.
214def transform(k: Nat, big: Bool, +q: Nat, temp: T.Quad) -> T.Quad:
215  match k:
216    case 0n:
217      xor_word(sub_word(rot_word(temp)), rcon(q))
218    case 1n:
219      temp
220    case 2n:
221      temp
222    case 3n:
223      temp
224    case 4n:
225      match big:
226        case True{}:
227          sub_word(temp)
228        case False{}:
229          temp
230    case 5n+p:
231      temp
232
233# w[i] = w[i-Nk] xor transform(w[i-1]); history holds w[i-1], w[i-2], ...
234def next_word(+nk: Nat, +i: Nat, +history: List<&2, T.Quad>) -> T.Quad:
235  xor_word(nth_word(history, Nat.sub(nk, 1n)),
236    transform(Nat.mod(i, nk), Nat.is_lt(6n, nk), Nat.div(i, nk), nth_word(history, 0n)))
237
238def expand(n: Nat, +nk: Nat, +i: Nat, +history: List<&2, T.Quad>) -> List<&2, T.Quad>:
239  match n:
240    case 0n:
241      history
242    case 1n+p:
243      expand(p, nk, 1n+i, next_word(nk, i, history) <> history)
244
245# Figure 11: Nb*(Nr+1) = 4*(Nr+1) words, the first Nk from the key.
246def key_expansion(+nk: Nat, +nr: Nat, key: List<&2, U32>) -> List<&2, T.Quad>:
247  List.reverse(&2, T.Quad, expand(Nat.sub(Nat.mul(4n, 1n+nr), nk), nk, nk,
248    List.reverse(&2, T.Quad, key_words(key))))
249
250def encrypt(+nk: Nat, +nr: Nat, key: List<&2, U32>, block: T.State) -> T.State:
251  cipher(nr, block, key_expansion(nk, nr, key))
252
253# Figure 4: Nk is the key length in 32-bit words and Nr = Nk + 6
254# (AES-128: 4, 10; AES-192: 6, 12; AES-256: 8, 14).
255def nk(+key: List<&2, U32>) -> Nat:
256  Nat.div(List.length(&2, U32, key), 4n)
257
258def nr(+key: List<&2, U32>) -> Nat:
259  Nat.add(nk(key), 6n)
260
261# AES with the key's own Nk and Nr.
262def aes(+key: List<&2, U32>, block: T.State) -> T.State:
263  encrypt(nk(key), nr(key), key, block)
264
265def aes128(key: List<&2, U32>, block: T.State) -> T.State:
266  encrypt(4n, 10n, key, block)
267
268def aes256(key: List<&2, U32>, block: T.State) -> T.State:
269  encrypt(8n, 14n, key, block)
270
271# Section 3.4: s[r,c] = in[r + 4c], out[r + 4c] = s[r,c]. Missing input
272# bytes read as {00}.
273def nth_byte(xs: List<&2, U32>, n: Nat) -> U32:
274  match xs n:
275    case Nil{} _:
276      0
277    case x <> rest 0n:
278      x
279    case x <> rest 1n+p:
280      nth_byte(rest, p)
281
282def state_of(+input: List<&2, U32>) -> T.State:
283  T.S{T.W{nth_byte(input, 0n), nth_byte(input, 1n), nth_byte(input, 2n), nth_byte(input, 3n)},
284    T.W{nth_byte(input, 4n), nth_byte(input, 5n), nth_byte(input, 6n), nth_byte(input, 7n)},
285    T.W{nth_byte(input, 8n), nth_byte(input, 9n), nth_byte(input, 10n), nth_byte(input, 11n)},
286    T.W{nth_byte(input, 12n), nth_byte(input, 13n), nth_byte(input, 14n), nth_byte(input, 15n)}}
287
288def bytes_of(s: T.State) -> List<&2, U32>:
289  match s:
290    case T.S{T.W{a0, a1, a2, a3}, T.W{b0, b1, b2, b3}, T.W{c0, c1, c2, c3}, T.W{d0, d1, d2, d3}}:
291      [a0, a1, a2, a3, b0, b1, b2, b3, c0, c1, c2, c3, d0, d1, d2, d3]