1	import Base
2	import ../../../src/crypto/aes/types.bend as T
3	import ./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.
11	def 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	
16	def 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.
22	def 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).
26	def 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.)
35	def 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	
45	def 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}.
56	def 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	
62	def 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.
70	def 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).
74	def 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).
80	def 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	
85	def 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.
91	def 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	
102	def 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	
113	def 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.
117	def 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	
122	def 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.
128	def 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	
132	def 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).
136	def 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	
144	def 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	
149	def 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].
155	def 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.
163	def 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	
172	def 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].
180	def 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	
185	def xpow(n: Nat) -> U32:
186	  match n:
187	    case 0n:
188	      1
189	    case 1n+p:
190	      mul(2, xpow(p))
191	
192	def 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.
196	def 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	
203	def 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.
214	def 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], ...
234	def 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	
238	def 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.
246	def 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	
250	def 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).
255	def nk(+key: List<&2, U32>) -> Nat:
256	  Nat.div(List.length(&2, U32, key), 4n)
257	
258	def nr(+key: List<&2, U32>) -> Nat:
259	  Nat.add(nk(key), 6n)
260	
261	# AES with the key's own Nk and Nr.
262	def aes(+key: List<&2, U32>, block: T.State) -> T.State:
263	  encrypt(nk(key), nr(key), key, block)
264	
265	def aes128(key: List<&2, U32>, block: T.State) -> T.State:
266	  encrypt(4n, 10n, key, block)
267	
268	def 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}.
273	def 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	
282	def 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	
288	def 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]