1	import Base
2	import ../../../src/crypto/aes/types.bend as T
3	import ../../../src/crypto/aes/sbox.bend as B
4	import ../../../src/crypto/aes/aes.bend as I
5	import ../../../spec/crypto/aes/aes.bend as S
6	import ./sbox.bend as L
7	
8	# The AES implementation equals the FIPS 197 specification, layer by layer:
9	# the byte functions (./sbox.bend), then SubBytes, ShiftRows, MixColumns,
10	# AddRoundKey, the rounds, the key expansion and the cipher, for every key
11	# (of any length and any Nk, Nr) and every state.
12	
13	# Congruence of the record constructors (one rewrite per field, on variables).
14	law quad_cong:
15	  for +a: U32
16	  for +b: U32
17	  for +c: U32
18	  for +d: U32
19	  for +a2: U32
20	  for +b2: U32
21	  for +c2: U32
22	  for +d2: U32
23	  for +ea: {a == a2 : U32}
24	  for +eb: {b == b2 : U32}
25	  for +ec: {c == c2 : U32}
26	  for +ed: {d == d2 : U32}
27	  {T.W{a, b, c, d} == T.W{a2, b2, c2, d2} : T.Quad}
28	
29	def quad_cong(a, b, c, d, a2, b2, c2, d2, ea, eb, ec, ed):
30	  %ea : {T.W{a, b, c, d} == T.W{_, b2, c2, d2} : T.Quad}
31	  %eb : {T.W{a, b, c, d} == T.W{a, _, c2, d2} : T.Quad}
32	  %ec : {T.W{a, b, c, d} == T.W{a, b, _, d2} : T.Quad}
33	  %ed : {T.W{a, b, c, d} == T.W{a, b, c, _} : T.Quad}
34	  {==}
35	
36	law state_cong:
37	  for +a: T.Quad
38	  for +b: T.Quad
39	  for +c: T.Quad
40	  for +d: T.Quad
41	  for +a2: T.Quad
42	  for +b2: T.Quad
43	  for +c2: T.Quad
44	  for +d2: T.Quad
45	  for +ea: {a == a2 : T.Quad}
46	  for +eb: {b == b2 : T.Quad}
47	  for +ec: {c == c2 : T.Quad}
48	  for +ed: {d == d2 : T.Quad}
49	  {T.S{a, b, c, d} == T.S{a2, b2, c2, d2} : T.State}
50	
51	def state_cong(a, b, c, d, a2, b2, c2, d2, ea, eb, ec, ed):
52	  %ea : {T.S{a, b, c, d} == T.S{_, b2, c2, d2} : T.State}
53	  %eb : {T.S{a, b, c, d} == T.S{a, _, c2, d2} : T.State}
54	  %ec : {T.S{a, b, c, d} == T.S{a, b, _, d2} : T.State}
55	  %ed : {T.S{a, b, c, d} == T.S{a, b, c, _} : T.State}
56	  {==}
57	
58	law sub_word_ok:
59	  for +w: T.Quad
60	  {I.sub_word(w) == S.sub_word(w) : T.Quad}
61	
62	def sub_word_ok(w):
63	  match w:
64	    case T.W{a, b, c, d}:
65	      quad_cong(B.sbox(a), B.sbox(b), B.sbox(c), B.sbox(d), S.sbox(a), S.sbox(b), S.sbox(c), S.sbox(d), L.sbox_ok(a), L.sbox_ok(b), L.sbox_ok(c), L.sbox_ok(d))
66	
67	law sub_bytes_ok:
68	  for +s: T.State
69	  {I.sub_bytes(s) == S.sub_bytes(s) : T.State}
70	
71	def sub_bytes_ok(s):
72	  match s:
73	    case T.S{c0, c1, c2, c3}:
74	      state_cong(I.sub_word(c0), I.sub_word(c1), I.sub_word(c2), I.sub_word(c3), S.sub_word(c0), S.sub_word(c1), S.sub_word(c2), S.sub_word(c3), sub_word_ok(c0), sub_word_ok(c1), sub_word_ok(c2), sub_word_ok(c3))
75	
76	law shift_rows_ok:
77	  for +s: T.State
78	  {I.shift_rows(s) == S.shift_rows(s) : T.State}
79	
80	def shift_rows_ok(s):
81	  match s:
82	    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}}:
83	      {==}
84	
85	law xor4_cong:
86	  for +p: U32
87	  for +q: U32
88	  for +r: U32
89	  for +s: U32
90	  for +p2: U32
91	  for +q2: U32
92	  for +r2: U32
93	  for +s2: U32
94	  for +ep: {p == p2 : U32}
95	  for +eq: {q == q2 : U32}
96	  for +er: {r == r2 : U32}
97	  for +es: {s == s2 : U32}
98	  {U32.xor(U32.xor(U32.xor(p, q), r), s) == U32.xor(U32.xor(U32.xor(p2, q2), r2), s2) : U32}
99	
100	def xor4_cong(p, q, r, s, p2, q2, r2, s2, ep, eq, er, es):
101	  %ep : {U32.xor(U32.xor(U32.xor(p, q), r), s) == U32.xor(U32.xor(U32.xor(_, q2), r2), s2) : U32}
102	  %eq : {U32.xor(U32.xor(U32.xor(p, q), r), s) == U32.xor(U32.xor(U32.xor(p, _), r2), s2) : U32}
103	  %er : {U32.xor(U32.xor(U32.xor(p, q), r), s) == U32.xor(U32.xor(U32.xor(p, q), _), s2) : U32}
104	  %es : {U32.xor(U32.xor(U32.xor(p, q), r), s) == U32.xor(U32.xor(U32.xor(p, q), r), _) : U32}
105	  {==}
106	
107	law mix_column_ok:
108	  for +w: T.Quad
109	  {I.mix_column(w) == S.mix_column(w) : T.Quad}
110	
111	def mix_column_ok(w):
112	  match w:
113	    case T.W{a, b, c, d}:
114	      quad_cong(U32.xor(U32.xor(U32.xor(B.xtime(a), B.mul3(b)), c), d), U32.xor(U32.xor(U32.xor(a, B.xtime(b)), B.mul3(c)), d), U32.xor(U32.xor(U32.xor(a, b), B.xtime(c)), B.mul3(d)), U32.xor(U32.xor(U32.xor(B.mul3(a), b), c), B.xtime(d)), U32.xor(U32.xor(U32.xor(S.mul(2, a), S.mul(3, b)), c), d), U32.xor(U32.xor(U32.xor(a, S.mul(2, b)), S.mul(3, c)), d), U32.xor(U32.xor(U32.xor(a, b), S.mul(2, c)), S.mul(3, d)), U32.xor(U32.xor(U32.xor(S.mul(3, a), b), c), S.mul(2, d)),
115	        xor4_cong(B.xtime(a), B.mul3(b), c, d, S.mul(2, a), S.mul(3, b), c, d, L.xtime_ok(a), L.mul3_ok(b), {==}, {==}), xor4_cong(a, B.xtime(b), B.mul3(c), d, a, S.mul(2, b), S.mul(3, c), d, {==}, L.xtime_ok(b), L.mul3_ok(c), {==}), xor4_cong(a, b, B.xtime(c), B.mul3(d), a, b, S.mul(2, c), S.mul(3, d), {==}, {==}, L.xtime_ok(c), L.mul3_ok(d)), xor4_cong(B.mul3(a), b, c, B.xtime(d), S.mul(3, a), b, c, S.mul(2, d), L.mul3_ok(a), {==}, {==}, L.xtime_ok(d)))
116	
117	law mix_columns_ok:
118	  for +s: T.State
119	  {I.mix_columns(s) == S.mix_columns(s) : T.State}
120	
121	def mix_columns_ok(s):
122	  match s:
123	    case T.S{c0, c1, c2, c3}:
124	      state_cong(I.mix_column(c0), I.mix_column(c1), I.mix_column(c2), I.mix_column(c3), S.mix_column(c0), S.mix_column(c1), S.mix_column(c2), S.mix_column(c3), mix_column_ok(c0), mix_column_ok(c1), mix_column_ok(c2), mix_column_ok(c3))
125	
126	law xor_word_ok:
127	  for +a: T.Quad
128	  for +b: T.Quad
129	  {I.xor_word(a, b) == S.xor_word(a, b) : T.Quad}
130	
131	def xor_word_ok(a, b):
132	  match a b:
133	    case T.W{a0, a1, a2, a3} T.W{b0, b1, b2, b3}:
134	      {==}
135	
136	law ark_ok:
137	  for +s: T.State
138	  for +w0: T.Quad
139	  for +w1: T.Quad
140	  for +w2: T.Quad
141	  for +w3: T.Quad
142	  {I.add_round_key(s, T.S{w0, w1, w2, w3}) == S.add_round_key(s, w0, w1, w2, w3) : T.State}
143	
144	def ark_ok(s, w0, w1, w2, w3):
145	  match s:
146	    case T.S{c0, c1, c2, c3}:
147	      %xor_word_ok(c0, w0) : {T.S{I.xor_word(c0, w0), I.xor_word(c1, w1), I.xor_word(c2, w2), I.xor_word(c3, w3)} == T.S{_, S.xor_word(c1, w1), S.xor_word(c2, w2), S.xor_word(c3, w3)} : T.State}
148	      %xor_word_ok(c1, w1) : {T.S{I.xor_word(c0, w0), I.xor_word(c1, w1), I.xor_word(c2, w2), I.xor_word(c3, w3)} == T.S{I.xor_word(c0, w0), _, S.xor_word(c2, w2), S.xor_word(c3, w3)} : T.State}
149	      %xor_word_ok(c2, w2) : {T.S{I.xor_word(c0, w0), I.xor_word(c1, w1), I.xor_word(c2, w2), I.xor_word(c3, w3)} == T.S{I.xor_word(c0, w0), I.xor_word(c1, w1), _, S.xor_word(c3, w3)} : T.State}
150	      %xor_word_ok(c3, w3) : {T.S{I.xor_word(c0, w0), I.xor_word(c1, w1), I.xor_word(c2, w2), I.xor_word(c3, w3)} == T.S{I.xor_word(c0, w0), I.xor_word(c1, w1), I.xor_word(c2, w2), _} : T.State}
151	      {==}
152	
153	law round_ok:
154	  for +s: T.State
155	  for +w0: T.Quad
156	  for +w1: T.Quad
157	  for +w2: T.Quad
158	  for +w3: T.Quad
159	  {I.round(s, T.S{w0, w1, w2, w3}) == S.add_round_key(S.mix_columns(S.shift_rows(S.sub_bytes(s))), w0, w1, w2, w3) : T.State}
160	
161	def round_ok(s, w0, w1, w2, w3):
162	  %sub_bytes_ok(s) : {I.round(s, T.S{w0, w1, w2, w3}) == S.add_round_key(S.mix_columns(S.shift_rows(_)), w0, w1, w2, w3) : T.State}
163	  %shift_rows_ok(I.sub_bytes(s)) : {I.round(s, T.S{w0, w1, w2, w3}) == S.add_round_key(S.mix_columns(_), w0, w1, w2, w3) : T.State}
164	  %mix_columns_ok(I.shift_rows(I.sub_bytes(s))) : {I.round(s, T.S{w0, w1, w2, w3}) == S.add_round_key(_, w0, w1, w2, w3) : T.State}
165	  ark_ok(I.mix_columns(I.shift_rows(I.sub_bytes(s))), w0, w1, w2, w3)
166	
167	law last_round_ok:
168	  for +s: T.State
169	  for +w0: T.Quad
170	  for +w1: T.Quad
171	  for +w2: T.Quad
172	  for +w3: T.Quad
173	  {I.last_round(s, T.S{w0, w1, w2, w3}) == S.add_round_key(S.shift_rows(S.sub_bytes(s)), w0, w1, w2, w3) : T.State}
174	
175	def last_round_ok(s, w0, w1, w2, w3):
176	  %sub_bytes_ok(s) : {I.last_round(s, T.S{w0, w1, w2, w3}) == S.add_round_key(S.shift_rows(_), w0, w1, w2, w3) : T.State}
177	  %shift_rows_ok(I.sub_bytes(s)) : {I.last_round(s, T.S{w0, w1, w2, w3}) == S.add_round_key(_, w0, w1, w2, w3) : T.State}
178	  ark_ok(I.shift_rows(I.sub_bytes(s)), w0, w1, w2, w3)
179	
180	law rounds_ok:
181	  for +n: Nat
182	  for +s: T.State
183	  for +w: List<&2, T.Quad>
184	  {I.rounds(n, s, I.round_keys(w)) == S.rounds(n, s, w) : T.State}
185	
186	def rounds_ok(n, s, w):
187	  match n w:
188	    case 0n w0 <> w1 <> w2 <> w3 <> rest:
189	      last_round_ok(s, w0, w1, w2, w3)
190	    case 1n+p w0 <> w1 <> w2 <> w3 <> rest:
191	      %round_ok(s, w0, w1, w2, w3) : {I.rounds(p, I.round(s, T.S{w0, w1, w2, w3}), I.round_keys(rest)) == S.rounds(p, _, rest) : T.State}
192	      rounds_ok(p, I.round(s, T.S{w0, w1, w2, w3}), rest)
193	    case 0n Nil{}:
194	      {==}
195	    case 0n w0 <> Nil{}:
196	      {==}
197	    case 0n w0 <> w1 <> Nil{}:
198	      {==}
199	    case 0n w0 <> w1 <> w2 <> Nil{}:
200	      {==}
201	    case 1n+p Nil{}:
202	      {==}
203	    case 1n+p w0 <> Nil{}:
204	      {==}
205	    case 1n+p w0 <> w1 <> Nil{}:
206	      {==}
207	    case 1n+p w0 <> w1 <> w2 <> Nil{}:
208	      {==}
209	
210	law cipher_ok:
211	  for +nr: Nat
212	  for +s: T.State
213	  for +w: List<&2, T.Quad>
214	  {I.cipher(nr, s, I.round_keys(w)) == S.cipher(nr, s, w) : T.State}
215	
216	def cipher_ok(nr, s, w):
217	  match w:
218	    case w0 <> w1 <> w2 <> w3 <> rest:
219	      %ark_ok(s, w0, w1, w2, w3) : {I.rounds(Nat.sub(nr, 1n), I.add_round_key(s, T.S{w0, w1, w2, w3}), I.round_keys(rest)) == S.rounds(Nat.sub(nr, 1n), _, rest) : T.State}
220	      rounds_ok(Nat.sub(nr, 1n), I.add_round_key(s, T.S{w0, w1, w2, w3}), rest)
221	    case Nil{}:
222	      {==}
223	    case w0 <> Nil{}:
224	      {==}
225	    case w0 <> w1 <> Nil{}:
226	      {==}
227	    case w0 <> w1 <> w2 <> Nil{}:
228	      {==}
229	
230	# ---- key expansion ----
231	
232	law rot_word_ok:
233	  for +w: T.Quad
234	  {I.rot_word(w) == S.rot_word(w) : T.Quad}
235	
236	def rot_word_ok(w):
237	  match w:
238	    case T.W{a, b, c, d}:
239	      {==}
240	
241	law xpow_ok:
242	  for +n: Nat
243	  {I.xpow(n) == S.xpow(n) : U32}
244	
245	def xpow_ok(n):
246	  match n:
247	    case 0n:
248	      {==}
249	    case 1n+p:
250	      %xpow_ok(p) : {B.xtime(I.xpow(p)) == S.mul(2, _) : U32}
251	      L.xtime_ok(I.xpow(p))
252	
253	law rcon_ok:
254	  for +j: Nat
255	  {I.rcon(j) == S.rcon(j) : T.Quad}
256	
257	def rcon_ok(j):
258	  %xpow_ok(Nat.sub(j, 1n)) : {T.W{I.xpow(Nat.sub(j, 1n)), 0, 0, 0} == T.W{_, 0, 0, 0} : T.Quad}
259	  {==}
260	
261	law key_words_ok:
262	  for +key: List<&2, U32>
263	  {I.key_words(key) == S.key_words(key) : List<&2, T.Quad>}
264	
265	def key_words_ok(key):
266	  match key:
267	    case a <> b <> c <> d <> rest:
268	      %key_words_ok(rest) : {T.W{a, b, c, d} <> I.key_words(rest) == T.W{a, b, c, d} <> _ : List<&2, T.Quad>}
269	      {==}
270	    case Nil{}:
271	      {==}
272	    case a <> Nil{}:
273	      {==}
274	    case a <> b <> Nil{}:
275	      {==}
276	    case a <> b <> c <> Nil{}:
277	      {==}
278	
279	law get_word_ok:
280	  for +ws: List<&2, T.Quad>
281	  for +n: Nat
282	  {I.get_word(ws, n) == S.nth_word(ws, n) : T.Quad}
283	
284	def get_word_ok(ws, n):
285	  match ws n:
286	    case Nil{} _:
287	      {==}
288	    case w <> rest 0n:
289	      {==}
290	    case w <> rest 1n+p:
291	      get_word_ok(rest, p)
292	
293	law core_ok:
294	  for +k: Nat
295	  for +big: Bool
296	  for +q: Nat
297	  for +temp: T.Quad
298	  {I.schedule_core(k, big, q, temp) == S.transform(k, big, q, temp) : T.Quad}
299	
300	def core_ok(k, big, q, temp):
301	  match k:
302	    case 0n:
303	      %rcon_ok(q) : {I.xor_word(I.sub_word(I.rot_word(temp)), I.rcon(q)) == S.xor_word(S.sub_word(S.rot_word(temp)), _) : T.Quad}
304	      %rot_word_ok(temp) : {I.xor_word(I.sub_word(I.rot_word(temp)), I.rcon(q)) == S.xor_word(S.sub_word(_), I.rcon(q)) : T.Quad}
305	      %sub_word_ok(I.rot_word(temp)) : {I.xor_word(I.sub_word(I.rot_word(temp)), I.rcon(q)) == S.xor_word(_, I.rcon(q)) : T.Quad}
306	      xor_word_ok(I.sub_word(I.rot_word(temp)), I.rcon(q))
307	    case 1n:
308	      {==}
309	    case 2n:
310	      {==}
311	    case 3n:
312	      {==}
313	    case 4n:
314	      match big:
315	        case True{}:
316	          sub_word_ok(temp)
317	        case False{}:
318	          {==}
319	    case 5n+p:
320	      {==}
321	
322	law next_ok:
323	  for +nk: Nat
324	  for +i: Nat
325	  for +h: List<&2, T.Quad>
326	  {I.next_word(nk, i, h) == S.next_word(nk, i, h) : T.Quad}
327	
328	def next_ok(nk, i, h):
329	  %get_word_ok(h, 0n) : {I.next_word(nk, i, h) == S.xor_word(S.nth_word(h, Nat.sub(nk, 1n)), S.transform(Nat.mod(i, nk), Nat.is_lt(6n, nk), Nat.div(i, nk), _)) : T.Quad}
330	  %core_ok(Nat.mod(i, nk), Nat.is_lt(6n, nk), Nat.div(i, nk), I.get_word(h, 0n)) : {I.next_word(nk, i, h) == S.xor_word(S.nth_word(h, Nat.sub(nk, 1n)), _) : T.Quad}
331	  %get_word_ok(h, Nat.sub(nk, 1n)) : {I.next_word(nk, i, h) == S.xor_word(_, I.schedule_core(Nat.mod(i, nk), Nat.is_lt(6n, nk), Nat.div(i, nk), I.get_word(h, 0n))) : T.Quad}
332	  xor_word_ok(I.get_word(h, Nat.sub(nk, 1n)), I.schedule_core(Nat.mod(i, nk), Nat.is_lt(6n, nk), Nat.div(i, nk), I.get_word(h, 0n)))
333	
334	law grow_ok:
335	  for +n: Nat
336	  for +nk: Nat
337	  for +i: Nat
338	  for +h: List<&2, T.Quad>
339	  {I.grow(n, nk, i, h) == S.expand(n, nk, i, h) : List<&2, T.Quad>}
340	
341	def grow_ok(n, nk, i, h):
342	  match n:
343	    case 0n:
344	      {==}
345	    case 1n+p:
346	      %next_ok(nk, i, h) : {I.grow(p, nk, 1n+i, I.next_word(nk, i, h) <> h) == S.expand(p, nk, 1n+i, _ <> h) : List<&2, T.Quad>}
347	      grow_ok(p, nk, 1n+i, I.next_word(nk, i, h) <> h)
348	
349	law expand_ok:
350	  for +nk: Nat
351	  for +nr: Nat
352	  for +key: List<&2, U32>
353	  {I.expand(nk, nr, key) == I.round_keys(S.key_expansion(nk, nr, key)) : List<&2, T.State>}
354	
355	def expand_ok(nk, nr, key):
356	  %key_words_ok(key) : {I.expand(nk, nr, key) == I.round_keys(List.reverse(&2, T.Quad, S.expand(Nat.sub(Nat.mul(4n, 1n+nr), nk), nk, nk, List.reverse(&2, T.Quad, _)))) : List<&2, T.State>}
357	  %grow_ok(Nat.sub(Nat.mul(4n, 1n+nr), nk), nk, nk, List.reverse(&2, T.Quad, I.key_words(key))) : {I.expand(nk, nr, key) == I.round_keys(List.reverse(&2, T.Quad, _)) : List<&2, T.State>}
358	  {==}
359	
360	# The implementation's cipher with its expanded key is FIPS 197's Cipher
361	# with KeyExpansion, for every Nk, Nr, key and block.
362	law encrypt_ok:
363	  for +nk: Nat
364	  for +nr: Nat
365	  for +key: List<&2, U32>
366	  for +s: T.State
367	  {I.encrypt(nk, nr, key, s) == S.encrypt(nk, nr, key, s) : T.State}
368	
369	def encrypt_ok(nk, nr, key, s):
370	  %Equal.sym(List<&2, T.State>, I.expand(nk, nr, key), I.round_keys(S.key_expansion(nk, nr, key)), expand_ok(nk, nr, key)) : {I.cipher(nr, s, _) == S.cipher(nr, s, S.key_expansion(nk, nr, key)) : T.State}
371	  cipher_ok(nr, s, S.key_expansion(nk, nr, key))