1	# AES-128 and AES-256 (FIPS 197) as block ciphers for the modes: the named
2	# default ciphers, used as
3	#   ~CmAesKey, ~CmAes.encrypt, ~CmAes.decrypt
4	# Encryption is bend-collections' AES (bend-collections@1.0.0.0,
5	# src/crypto/aes/aes.bend: expand_key and encrypt_block). That library has no
6	# decryption, so this file adds FIPS 197's inverse cipher (section 5.3) and the
7	# key expansion it needs (section 5.2), written here from the specification.
8	# The inverse S-box is computed, not looked up: the inverse affine map, then
9	# the inverse in GF(2^8) as x^254, with masks instead of branches, so no
10	# branch or table index depends on a key or data byte.
11	import Base
12	import bend-osdev-bytes@0.1.0.0/bytes.bend as B
13	import ./cipher_modes.bend as M
14	import bend-collections@1.0.0.0/src/crypto/aes/aes.bend as A
15	
16	# GF(2^8) modulo x^8 + x^4 + x^3 + x + 1
17	
18	def aes.xt(+x: U32) -> U32:
19	  U32.and(255, U32.xor(U32.shl(x), U32.and(27, M.cm.mask(U32.shrn(x, 7n)))))
20	
21	def aes.gmul.go(n: Nat, +a: U32, +b: U32, p: U32) -> U32:
22	  match n:
23	    case 0n:
24	      p
25	    case 1n+q:
26	      aes.gmul.go(q, aes.xt(a), U32.shr(b), U32.xor(p, U32.and(a, M.cm.mask(b))))
27	
28	def aes.gmul(+a: U32, +b: U32) -> U32:
29	  aes.gmul.go(8n, a, b, 0)
30	
31	# a := a^2 * x, n times: from a = x, six times gives x^127
32	def aes.inv.go(n: Nat, a: U32, +x: U32) -> U32:
33	  match n:
34	    case 0n:
35	      a
36	    case 1n+q:
37	      +s = a
38	      aes.inv.go(q, aes.gmul(aes.gmul(s, s), x), x)
39	
40	# the multiplicative inverse, x^254 (0 for 0)
41	def aes.inv(+x: U32) -> U32:
42	  +a = aes.inv.go(6n, x, x)
43	  aes.gmul(a, a)
44	
45	def aes.rotl(+b: U32, +n: Nat) -> U32:
46	  U32.and(255, U32.or(U32.shln(b, n), U32.shrn(b, Nat.sub(8n, n))))
47	
48	# the S-box: the inverse, then the affine map
49	def aes.sbox(+x: U32) -> U32:
50	  +b = aes.inv(U32.and(255, x))
51	  U32.xor(U32.xor(U32.xor(U32.xor(U32.xor(b, aes.rotl(b, 1n)), aes.rotl(b, 2n)), aes.rotl(b, 3n)), aes.rotl(b, 4n)), 99)
52	
53	# the inverse S-box: the inverse affine map, then the inverse
54	def aes.isbox(x: U32) -> U32:
55	  +b = U32.and(255, x)
56	  aes.inv(U32.xor(U32.xor(U32.xor(aes.rotl(b, 1n), aes.rotl(b, 3n)), aes.rotl(b, 6n)), 5))
57	
58	# The state: 16 bytes, column by column (byte 4c + r is row r of column c)
59	
60	type CmAesState is Data:
61	  CmAesState{s0: U32, s1: U32, s2: U32, s3: U32, s4: U32, s5: U32, s6: U32, s7: U32, s8: U32, s9: U32, s10: U32, s11: U32, s12: U32, s13: U32, s14: U32, s15: U32}
62	
63	def aes.byte(+w: U32, n: Nat) -> U32:
64	  U32.and(255, U32.shrn(w, n))
65	
66	def aes.state_of(b: M.CmBlock) -> CmAesState:
67	  match b:
68	    case M.CmBlock{+w0, +w1, +w2, +w3}:
69	      CmAesState{aes.byte(w0, 24n), aes.byte(w0, 16n), aes.byte(w0, 8n), aes.byte(w0, 0n), aes.byte(w1, 24n), aes.byte(w1, 16n), aes.byte(w1, 8n), aes.byte(w1, 0n), aes.byte(w2, 24n), aes.byte(w2, 16n), aes.byte(w2, 8n), aes.byte(w2, 0n), aes.byte(w3, 24n), aes.byte(w3, 16n), aes.byte(w3, 8n), aes.byte(w3, 0n)}
70	
71	def aes.block_of(s: CmAesState) -> M.CmBlock:
72	  match s:
73	    case CmAesState{a0, a1, a2, a3, a4, a5, a6, a7, a8, a9, a10, a11, a12, a13, a14, a15}:
74	      M.CmBlock{M.cm.word(a0, a1, a2, a3), M.cm.word(a4, a5, a6, a7), M.cm.word(a8, a9, a10, a11), M.cm.word(a12, a13, a14, a15)}
75	
76	def aes.add(s: CmAesState, k: CmAesState) -> CmAesState:
77	  match s k:
78	    case CmAesState{a0, a1, a2, a3, a4, a5, a6, a7, a8, a9, a10, a11, a12, a13, a14, a15} CmAesState{b0, b1, b2, b3, b4, b5, b6, b7, b8, b9, b10, b11, b12, b13, b14, b15}:
79	      CmAesState{U32.xor(a0, b0), U32.xor(a1, b1), U32.xor(a2, b2), U32.xor(a3, b3), U32.xor(a4, b4), U32.xor(a5, b5), U32.xor(a6, b6), U32.xor(a7, b7), U32.xor(a8, b8), U32.xor(a9, b9), U32.xor(a10, b10), U32.xor(a11, b11), U32.xor(a12, b12), U32.xor(a13, b13), U32.xor(a14, b14), U32.xor(a15, b15)}
80	
81	# InvShiftRows: row r moves r columns to the right
82	# InvSubBytes: the inverse S-box on every byte (both in one pass)
83	def aes.inv_shift_sub(s: CmAesState) -> CmAesState:
84	  match s:
85	    case CmAesState{a0, a1, a2, a3, a4, a5, a6, a7, a8, a9, a10, a11, a12, a13, a14, a15}:
86	      CmAesState{aes.isbox(a0), aes.isbox(a13), aes.isbox(a10), aes.isbox(a7), aes.isbox(a4), aes.isbox(a1), aes.isbox(a14), aes.isbox(a11), aes.isbox(a8), aes.isbox(a5), aes.isbox(a2), aes.isbox(a15), aes.isbox(a12), aes.isbox(a9), aes.isbox(a6), aes.isbox(a3)}
87	
88	# multiples of a byte by 9, 11, 13 and 14
89	def aes.m9(+a: U32) -> U32:
90	  U32.xor(aes.xt(aes.xt(aes.xt(a))), a)
91	
92	def aes.m11(+a: U32) -> U32:
93	  +a2 = aes.xt(a)
94	  U32.xor(U32.xor(aes.xt(aes.xt(a2)), a2), a)
95	
96	def aes.m13(+a: U32) -> U32:
97	  +a4 = aes.xt(aes.xt(a))
98	  U32.xor(U32.xor(aes.xt(a4), a4), a)
99	
100	def aes.m14(+a: U32) -> U32:
101	  +a2 = aes.xt(a)
102	  +a4 = aes.xt(a2)
103	  U32.xor(U32.xor(aes.xt(a4), a4), a2)
104	
105	def aes.mix0(+a: U32, +b: U32, +c: U32, +d: U32) -> U32:
106	  U32.xor(U32.xor(aes.m14(a), aes.m11(b)), U32.xor(aes.m13(c), aes.m9(d)))
107	
108	# InvMixColumns: each column times {0e, 0b, 0d, 09}, rotated per row
109	def aes.inv_mix(s: CmAesState) -> CmAesState:
110	  match s:
111	    case CmAesState{+a0, +a1, +a2, +a3, +a4, +a5, +a6, +a7, +a8, +a9, +a10, +a11, +a12, +a13, +a14, +a15}:
112	      CmAesState{aes.mix0(a0, a1, a2, a3), aes.mix0(a1, a2, a3, a0), aes.mix0(a2, a3, a0, a1), aes.mix0(a3, a0, a1, a2), aes.mix0(a4, a5, a6, a7), aes.mix0(a5, a6, a7, a4), aes.mix0(a6, a7, a4, a5), aes.mix0(a7, a4, a5, a6), aes.mix0(a8, a9, a10, a11), aes.mix0(a9, a10, a11, a8), aes.mix0(a10, a11, a8, a9), aes.mix0(a11, a8, a9, a10), aes.mix0(a12, a13, a14, a15), aes.mix0(a13, a14, a15, a12), aes.mix0(a14, a15, a12, a13), aes.mix0(a15, a12, a13, a14)}
113	
114	# The inverse cipher over the round keys, last round key first:
115	# AddRoundKey(Nr), then Nr - 1 rounds of InvShiftRows, InvSubBytes,
116	# AddRoundKey, InvMixColumns, then InvShiftRows, InvSubBytes, AddRoundKey(0)
117	def aes.rounds(rest: List<&2, CmAesState>, k: CmAesState, s: CmAesState) -> CmAesState:
118	  match rest:
119	    case Nil{}:
120	      aes.add(aes.inv_shift_sub(s), k)
121	    case Con{k2, rest2}:
122	      aes.rounds(rest2, k2, aes.inv_mix(aes.add(aes.inv_shift_sub(s), k)))
123	
124	def aes.decrypt(ks: List<&2, CmAesState>, s: CmAesState) -> CmAesState:
125	  match ks:
126	    case Nil{}:
127	      s
128	    case Con{k, Nil{}}:
129	      aes.add(s, k)
130	    case Con{k, Con{k2, rest}}:
131	      aes.rounds(rest, k2, aes.add(s, k))
132	
133	# Key expansion (FIPS 197, 5.2), on big-endian words
134	
135	def aes.sub_word(+w: U32) -> U32:
136	  M.cm.word(aes.sbox(U32.shrn(w, 24n)), aes.sbox(U32.shrn(w, 16n)), aes.sbox(U32.shrn(w, 8n)), aes.sbox(w))
137	
138	def aes.rot_word(+w: U32) -> U32:
139	  U32.or(U32.shln(w, 8n), U32.shrn(w, 24n))
140	
141	def aes.rcon.go(n: Nat, +x: U32) -> U32:
142	  match n:
143	    case 0n:
144	      x
145	    case 1n+p:
146	      aes.rcon.go(p, aes.xt(x))
147	
148	# x^(j - 1) in GF(2^8), the round constant of word j * Nk
149	def aes.rcon(j: Nat) -> U32:
150	  match j:
151	    case 0n:
152	      0
153	    case 1n+p:
154	      aes.rcon.go(p, 1)
155	
156	# Nk > 6 and i mod Nk = 4: SubWord alone
157	def aes.temp.big(big: Bool, +w: U32) -> U32:
158	  match big:
159	    case True{}:
160	      aes.sub_word(w)
161	    case False{}:
162	      w
163	
164	# temp from w[i - 1], for word i
165	def aes.temp(first: Bool, +i: Nat, +nk: Nat, +w: U32) -> U32:
166	  match first:
167	    case True{}:
168	      U32.xor(aes.sub_word(aes.rot_word(w)), U32.shln(aes.rcon(Nat.div(i, nk)), 24n))
169	    case False{}:
170	      aes.temp.big(Bool.and(Nat.is_lt(6n, nk), Nat.is_eq(Nat.mod(i, nk), 4n)), w)
171	
172	def aes.head(xs: List<&2, U32>) -> U32:
173	  match xs:
174	    case Nil{}:
175	      0
176	    case Con{h, t}:
177	      h
178	
179	def aes.tail(xs: List<&2, U32>) -> List<&2, U32>:
180	  match xs:
181	    case Nil{}:
182	      Nil{}
183	    case Con{h, t}:
184	      t
185	
186	def aes.last(xs: List<&2, U32>, d: U32) -> U32:
187	  match xs:
188	    case Nil{}:
189	      d
190	    case Con{h, t}:
191	      aes.last(t, h)
192	
193	# n more words from word i; win holds the last Nk words, oldest first, and acc
194	# every word so far, newest first
195	def aes.expand.go(n: Nat, +i: Nat, +nk: Nat, +win: List<&2, U32>, acc: List<&2, U32>) -> List<&2, U32>:
196	  match n:
197	    case 0n:
198	      List.reverse(&2, U32, acc)
199	    case 1n+p:
200	      +w = U32.xor(aes.head(win), aes.temp(Nat.is_eq(Nat.mod(i, nk), 0n), i, nk, aes.last(win, 0)))
201	      aes.expand.go(p, 1n+i, nk, B.append(aes.tail(win), Con{w, Nil{}}), Con{w, acc})
202	
203	# the key's bytes as big-endian words
204	def aes.words(xs: List<&2, U32>) -> List<&2, U32>:
205	  match xs:
206	    case Con{a, Con{b, Con{c, Con{d, t}}}}:
207	      Con{M.cm.word(a, b, c, d), aes.words(t)}
208	    case _:
209	      Nil{}
210	
211	def aes.round_keys(ws: List<&2, U32>) -> List<&2, CmAesState>:
212	  match ws:
213	    case Con{a, Con{b, Con{c, Con{d, t}}}}:
214	      Con{aes.state_of(M.CmBlock{a, b, c, d}), aes.round_keys(t)}
215	    case _:
216	      Nil{}
217	
218	# the 4 (Nr + 1) words of the schedule, Nr = Nk + 6
219	def aes.expand(+nk: Nat, xs: List<&2, U32>) -> List<&2, U32>:
220	  +ws = aes.words(xs)
221	  aes.expand.go(Nat.sub(Nat.mul(4n, Nat.add(nk, 7n)), nk), nk, nk, ws, List.reverse(&2, U32, ws))
222	
223	# The cipher
224	
225	# An expanded AES key: bend-collections' schedule for encryption, the key's
226	# bytes, and the round keys for decryption, last first, once CmAes.prepare
227	# has computed them (Nil until then: CmAes.decrypt then derives them for each
228	# block). Internal; built by CmAes.key.
229	type CmAesKey is Data:
230	  CmAesKey{sched: A.Schedule, raw: List<&2, U32>, dkeys: List<&2, CmAesState>}
231	
232	def aes.key.go(m: Maybe<&2, A.Schedule>, +xs: List<&2, U32>) -> Result<&2, &2, M.CmError, CmAesKey>:
233	  match m:
234	    case None{}:
235	      Fail{M.CmBadKeyLength{B.len(xs)}}
236	    case Some{s}:
237	      Done{CmAesKey{s, xs, Nil{}}}
238	
239	def aes.key.pick(ok: Bool, +xs: List<&2, U32>) -> Result<&2, &2, M.CmError, CmAesKey>:
240	  match ok:
241	    case True{}:
242	      aes.key.go(A.expand_key(xs), xs)
243	    case False{}:
244	      Fail{M.CmBadKeyLength{B.len(xs)}}
245	
246	# the decryption round keys: the ones computed already, or from the key bytes
247	def aes.dkeys(dk: List<&2, CmAesState>, +raw: List<&2, U32>) -> List<&2, CmAesState>:
248	  match dk:
249	    case Nil{}:
250	      List.reverse(&2, CmAesState, aes.round_keys(aes.expand(Nat.div(B.len(raw), 4n), raw)))
251	    case Con{h, t}:
252	      Con{h, t}
253	
254	# a 16-byte (AES-128) or 32-byte (AES-256) key
255	def CmAes.key(o: B.Octets) -> Result<&2, &2, M.CmError, CmAesKey>:
256	  +xs = B.Octets.to_list(o)
257	  +n = B.len(xs)
258	  aes.key.pick(Bool.or(Nat.is_eq(n, 16n), Nat.is_eq(n, 32n)), xs)
259	
260	def CmAes.key128(o: B.Octets) -> Result<&2, &2, M.CmError, CmAesKey>:
261	  +xs = B.Octets.to_list(o)
262	  aes.key.pick(Nat.is_eq(B.len(xs), 16n), xs)
263	
264	def CmAes.key256(o: B.Octets) -> Result<&2, &2, M.CmError, CmAesKey>:
265	  +xs = B.Octets.to_list(o)
266	  aes.key.pick(Nat.is_eq(B.len(xs), 32n), xs)
267	
268	def aes.of_maybe(m: Maybe<&2, List<&2, U32>>) -> M.CmBlock:
269	  match m:
270	    case None{}:
271	      M.CmBlock.zero()
272	    case Some{xs}:
273	      M.cm.block_of(xs)
274	
275	# one block, encrypted (bend-collections' encrypt_block)
276	def CmAes.encrypt(k: CmAesKey, b: M.CmBlock) -> M.CmBlock:
277	  match k:
278	    case CmAesKey{sched, raw, dk}:
279	      aes.of_maybe(A.encrypt_block(sched, M.cm.bytes_of(b)))
280	
281	# the key with its decryption round keys computed, for decrypting many blocks
282	def CmAes.prepare(k: CmAesKey) -> CmAesKey:
283	  match k:
284	    case CmAesKey{sched, +raw, dk}:
285	      CmAesKey{sched, raw, aes.dkeys(dk, raw)}
286	
287	# one block, decrypted (the inverse cipher above)
288	def CmAes.decrypt(k: CmAesKey, b: M.CmBlock) -> M.CmBlock:
289	  match k:
290	    case CmAesKey{sched, raw, dk}:
291	      aes.block_of(aes.decrypt(aes.dkeys(dk, raw), aes.state_of(b)))