1	# AEGIS-128 (H. Wu, B. Preneel, "AEGIS: A Fast Authenticated Encryption
2	# Algorithm (v1.1)", CAESAR final portfolio, 2016, section 2.3): a 128-bit
3	# key, a 128-bit nonce, a 128-bit tag, and an 80-byte state of five 16-byte
4	# blocks S0 .. S4 updated by the AES round function (aesround.bend):
5	#   StateUpdate128(S, m) = (AESRound(S4, S0 xor m), AESRound(S0, S1),
6	#                           AESRound(S1, S2), AESRound(S2, S3),
7	#                           AESRound(S3, S4))
8	#   initialization: S = (K xor IV, const1, const0, K xor const0,
9	#     K xor const1), then ten updates with K, K xor IV, K, K xor IV, ..
10	#   associated data: an update with each 16-byte block (the last zero-padded)
11	#   encryption: C_i = P_i xor S1 xor S4 xor (S2 and S3), then an update with
12	#     P_i (a short last block: only its bytes encrypted, the zero-padded
13	#     block used for the update)
14	#   finalization: tmp = S3 xor (adlen || msglen), the bit lengths as 64-bit
15	#     little-endian numbers; seven updates with tmp; the tag is
16	#     S0 xor S1 xor S2 xor S3 xor S4
17	# As an AEAD (crypto's aead.bend): ~AdAegisKey, ~AdAegis128.key,
18	# ~AdAegis128.seal, ~AdAegis128.open; the ciphertext is C followed by the
19	# 16-byte tag.
20	#
21	# The message is processed a byte at a time: each byte is XORed with the
22	# next keystream byte and joins the block that updates the state when it is
23	# whole. Decryption recovers each plaintext byte the same way and so runs
24	# through the same states, which is why open accepts what seal produced
25	# (LAWS.bend: aegis128_round_trip, for every key, nonce, associated data and
26	# message).
27	import Base
28	import bend-osdev-bytes@0.1.0.0/bytes.bend as B
29	import bend-osdev-word64@0.1.0.0/word64.bend as W
30	import bend-osdev-crypto@0.1.0.0/crypto.bend as Cr
31	import bend-osdev-crypto@0.1.0.0/octets.bend as O
32	import ./aead.bend as Ad
33	import ./aesround.bend as R
34	
35	# The state (internal)
36	type AdAegisSt is Data:
37	  AdAegisSt{s0: R.AdBlk, s1: R.AdBlk, s2: R.AdBlk, s3: R.AdBlk, s4: R.AdBlk}
38	
39	def ad.ae.update(+sb: R.AdSbox, st: AdAegisSt, +m: R.AdBlk) -> AdAegisSt:
40	  match st:
41	    case AdAegisSt{+a0, +a1, +a2, +a3, +a4}:
42	      AdAegisSt{R.AdAes.round(sb, a4, R.AdBlk.xor(a0, m)), R.AdAes.round(sb, a0, a1), R.AdAes.round(sb, a1, a2), R.AdAes.round(sb, a2, a3), R.AdAes.round(sb, a3, a4)}
43	
44	# the keystream block S1 xor S4 xor (S2 and S3)
45	def ad.ae.z(st: AdAegisSt) -> R.AdBlk:
46	  match st:
47	    case AdAegisSt{a0, a1, a2, a3, a4}:
48	      R.AdBlk.xor(R.AdBlk.xor(a1, a4), R.AdBlk.and(a2, a3))
49	
50	# const0 and const1 (the Fibonacci sequence mod 256)
51	def ad.ae.c0() -> R.AdBlk:
52	  R.AdBlk{65794, 50661389, 354563929, 2431220066}
53	
54	def ad.ae.c1() -> R.AdBlk:
55	  R.AdBlk{3678214229, 1841442801, 537997634, 1941252317}
56	
57	def ad.ae.init.go(n: Nat, +sb: R.AdSbox, +k: R.AdBlk, +kiv: R.AdBlk, st: AdAegisSt) -> AdAegisSt:
58	  match n:
59	    case 0n:
60	      st
61	    case 1n+p:
62	      ad.ae.init.go(p, sb, k, kiv, ad.ae.update(sb, ad.ae.update(sb, st, k), kiv))
63	
64	def ad.ae.init(+sb: R.AdSbox, +k: R.AdBlk, +iv: R.AdBlk) -> AdAegisSt:
65	  +kiv = R.AdBlk.xor(k, iv)
66	  ad.ae.init.go(5n, sb, k, kiv, AdAegisSt{kiv, ad.ae.c1(), ad.ae.c0(), R.AdBlk.xor(k, ad.ae.c0()), R.AdBlk.xor(k, ad.ae.c1())})
67	
68	# the state after the associated data's 16-byte pieces (the last
69	# zero-padded)
70	def ad.ae.absorb(+sb: R.AdSbox, ps: List<&2, List<&2, U32>>, st: AdAegisSt) -> AdAegisSt:
71	  match ps:
72	    case Nil{}:
73	      st
74	    case Con{p, t}:
75	      ad.ae.absorb(sb, t, ad.ae.update(sb, st, R.AdBlk.of_list(p)))
76	
77	# The message, a byte at a time (internal): the state, the keystream bytes
78	# left of the current block, the current block's plaintext bytes so far
79	# (last first) and how many
80	type AdAegisRun is Data:
81	  AdAegisRun{st: AdAegisSt, zs: List<&2, U32>, buf: List<&2, U32>, n: Nat}
82	
83	def ad.ae.run0(+st: AdAegisSt) -> AdAegisRun:
84	  AdAegisRun{st, R.AdBlk.to_list(ad.ae.z(st)), Nil{}, 0n}
85	
86	def ad.ae.rev(xs: List<&2, U32>, acc: List<&2, U32>) -> List<&2, U32>:
87	  match xs:
88	    case Nil{}:
89	      acc
90	    case Con{h, t}:
91	      ad.ae.rev(t, Con{h, acc})
92	
93	def ad.ae.tail(xs: List<&2, U32>) -> List<&2, U32>:
94	  match xs:
95	    case Nil{}:
96	      Nil{}
97	    case Con{h, t}:
98	      t
99	
100	def ad.ae.head(xs: List<&2, U32>) -> U32:
101	  match xs:
102	    case Nil{}:
103	      0
104	    case Con{h, t}:
105	      h
106	
107	def ad.ae.step.pick(full: Bool, +sb: R.AdSbox, st: AdAegisSt, zs: List<&2, U32>, +buf: List<&2, U32>, +n: Nat) -> AdAegisRun:
108	  match full:
109	    case True{}:
110	      ad.ae.run0(ad.ae.update(sb, st, R.AdBlk.of_list(ad.ae.rev(buf, Nil{}))))
111	    case False{}:
112	      AdAegisRun{st, ad.ae.tail(zs), buf, 1n+n}
113	
114	# the run after the plaintext byte x
115	def ad.ae.step(+sb: R.AdSbox, r: AdAegisRun, x: U32) -> AdAegisRun:
116	  match r:
117	    case AdAegisRun{st, zs, buf, +n}:
118	      ad.ae.step.pick(Nat.is_eq(n, 15n), sb, st, zs, Con{x, buf}, n)
119	
120	# the next keystream byte
121	def ad.ae.zhead(r: AdAegisRun) -> U32:
122	  match r:
123	    case AdAegisRun{st, zs, buf, n}:
124	      ad.ae.head(zs)
125	
126	def ad.ae.flush.pick(empty: Bool, +sb: R.AdSbox, st: AdAegisSt, buf: List<&2, U32>) -> AdAegisSt:
127	  match empty:
128	    case True{}:
129	      st
130	    case False{}:
131	      ad.ae.update(sb, st, R.AdBlk.of_list(ad.ae.rev(buf, Nil{})))
132	
133	# the state after the last, short block (zero-padded)
134	def ad.ae.flush(+sb: R.AdSbox, r: AdAegisRun) -> AdAegisSt:
135	  match r:
136	    case AdAegisRun{st, zs, buf, n}:
137	      ad.ae.flush.pick(Nat.is_eq(n, 0n), sb, st, buf)
138	
139	# Encryption and decryption, structurally (the proofs' reference): the
140	# output bytes, and the run after them
141	
142	def ad.ae.eout(xs: List<&2, U32>, +sb: R.AdSbox, +r: AdAegisRun) -> List<&2, U32>:
143	  match xs:
144	    case Nil{}:
145	      Nil{}
146	    case Con{+x, t}:
147	      Con{B.byte(U32.xor(x, ad.ae.zhead(r))), ad.ae.eout(t, sb, ad.ae.step(sb, r, B.byte(x)))}
148	
149	def ad.ae.eafter(xs: List<&2, U32>, +sb: R.AdSbox, +r: AdAegisRun) -> AdAegisRun:
150	  match xs:
151	    case Nil{}:
152	      r
153	    case Con{+x, t}:
154	      ad.ae.eafter(t, sb, ad.ae.step(sb, r, B.byte(x)))
155	
156	def ad.ae.dout(ys: List<&2, U32>, +sb: R.AdSbox, +r: AdAegisRun) -> List<&2, U32>:
157	  match ys:
158	    case Nil{}:
159	      Nil{}
160	    case Con{+y, t}:
161	      Con{B.byte(U32.xor(y, ad.ae.zhead(r))), ad.ae.dout(t, sb, ad.ae.step(sb, r, B.byte(U32.xor(y, ad.ae.zhead(r)))))}
162	
163	def ad.ae.dafter(ys: List<&2, U32>, +sb: R.AdSbox, +r: AdAegisRun) -> AdAegisRun:
164	  match ys:
165	    case Nil{}:
166	      r
167	    case Con{+y, t}:
168	      ad.ae.dafter(t, sb, ad.ae.step(sb, r, B.byte(U32.xor(y, ad.ae.zhead(r)))))
169	
170	# The loops that run (tail calls): the run after the bytes, and the output
171	type AdAegisOut is Data:
172	  AdAegisOut{r: AdAegisRun, out: List<&2, U32>}
173	
174	def ad.ae.ego(xs: List<&2, U32>, +sb: R.AdSbox, +r: AdAegisRun, acc: List<&2, U32>) -> AdAegisOut:
175	  match xs:
176	    case Nil{}:
177	      AdAegisOut{r, ad.ae.rev(acc, Nil{})}
178	    case Con{+x, t}:
179	      ad.ae.ego(t, sb, ad.ae.step(sb, r, B.byte(x)), Con{B.byte(U32.xor(x, ad.ae.zhead(r))), acc})
180	
181	def ad.ae.dgo(ys: List<&2, U32>, +sb: R.AdSbox, +r: AdAegisRun, acc: List<&2, U32>) -> AdAegisOut:
182	  match ys:
183	    case Nil{}:
184	      AdAegisOut{r, ad.ae.rev(acc, Nil{})}
185	    case Con{+y, t}:
186	      +p = B.byte(U32.xor(y, ad.ae.zhead(r)))
187	      ad.ae.dgo(t, sb, ad.ae.step(sb, r, p), Con{p, acc})
188	
189	# Finalization: the tag from the state after the message, and the lengths
190	def ad.ae.fin.go(n: Nat, +sb: R.AdSbox, +tmp: R.AdBlk, st: AdAegisSt) -> AdAegisSt:
191	  match n:
192	    case 0n:
193	      st
194	    case 1n+p:
195	      ad.ae.fin.go(p, sb, tmp, ad.ae.update(sb, st, tmp))
196	
197	def ad.ae.tag(st: AdAegisSt) -> B.Octets:
198	  match st:
199	    case AdAegisSt{a0, a1, a2, a3, a4}:
200	      R.AdBlk.to_octets(R.AdBlk.xor(R.AdBlk.xor(R.AdBlk.xor(a0, a1), R.AdBlk.xor(a2, a3)), a4))
201	
202	# n as a 64-bit word, counted up (W64.from_nat divides by 2^32, a literal
203	# the checker would expand when a proof meets it)
204	def ad.ae.w64(n: Nat, acc: W.W64) -> W.W64:
205	  match n:
206	    case 0n:
207	      acc
208	    case 1n+p:
209	      ad.ae.w64(p, W.W64.add(acc, W.W64.one()))
210	
211	# the bit length of n bytes as 8 little-endian bytes
212	def ad.ae.bits(n: Nat) -> List<&2, U32>:
213	  B.le64(W.W64.shl(ad.ae.w64(n, W.W64.zero()), 3n))
214	
215	def ad.ae.s3(st: AdAegisSt) -> R.AdBlk:
216	  match st:
217	    case AdAegisSt{a0, a1, a2, a3, a4}:
218	      a3
219	
220	def ad.ae.final(+sb: R.AdSbox, +st: AdAegisSt, +alen: Nat, +mlen: Nat) -> B.Octets:
221	  ad.ae.tag(ad.ae.fin.go(7n, sb, R.AdBlk.xor(ad.ae.s3(st), R.AdBlk.of_list(B.append(ad.ae.bits(alen), ad.ae.bits(mlen)))), st))
222	
223	# An AEGIS-128 key: the key block and the AES S-box (opaque)
224	type AdAegisKey is Data:
225	  AdAegisKey{k: R.AdBlk, sb: R.AdSbox}
226	
227	def ad.ae.key.pick(ok: Bool, +o: B.Octets) -> Result<&2, &2, Cr.CrError, AdAegisKey>:
228	  match ok:
229	    case True{}:
230	      Done{AdAegisKey{R.AdBlk.of_octets(o), R.AdAes.sbox()}}
231	    case False{}:
232	      Fail{Cr.CrBadKeyLength{B.Octets.len(o)}}
233	
234	def ad.ae.key(+o: B.Octets) -> Result<&2, &2, Cr.CrError, AdAegisKey>:
235	  ad.ae.key.pick(Nat.is_eq(B.Octets.len(o), 16n), o)
236	
237	# a 16-byte key
238	def AdAegis128.key(o: B.Octets) -> Result<&2, &2, Cr.CrError, AdAegisKey>:
239	  ad.ae.key(o)
240	
241	# the run before the message: the state after the key, nonce and associated
242	# data
243	def ad.ae.start(+sb: R.AdSbox, +k: R.AdBlk, +nonce: B.Octets, +aad: B.Octets) -> AdAegisRun:
244	  ad.ae.run0(ad.ae.absorb(sb, O.CrOct.pieces(16n, B.Octets.to_list(aad)), ad.ae.init(sb, k, R.AdBlk.of_octets(nonce))))
245	
246	def ad.ae.sealed(+sb: R.AdSbox, +aad: B.Octets, +p: B.Octets, o: AdAegisOut) -> B.Octets:
247	  match o:
248	    case AdAegisOut{r, out}:
249	      Ad.AdTag.join(B.Octets.of_list(out), ad.ae.final(sb, ad.ae.flush(sb, r), B.Octets.len(aad), B.Octets.len(p)))
250	
251	def ad.ae.seal.go(kk: AdAegisKey, +nonce: B.Octets, +aad: B.Octets, +p: B.Octets) -> B.Octets:
252	  match kk:
253	    case AdAegisKey{+k, +sb}:
254	      ad.ae.sealed(sb, aad, p, ad.ae.ego(B.Octets.to_list(p), sb, ad.ae.start(sb, k, nonce, aad), Nil{}))
255	
256	def ad.ae.nonce(+nonce: B.Octets) -> Result<&2, &2, Cr.CrError, Unit>:
257	  Cr.CrError.check(Nat.is_eq(B.Octets.len(nonce), 16n), Cr.CrBadNonceLength{B.Octets.len(nonce)})
258	
259	def ad.ae.seal(+kk: AdAegisKey, +nonce: B.Octets, +aad: B.Octets, +p: B.Octets) -> Result<&2, &2, Cr.CrError, B.Octets>:
260	  Result.map(&2, &2, Cr.CrError, Unit, B.Octets, u => ad.ae.seal.go(kk, nonce, aad, p), ad.ae.nonce(nonce))
261	
262	# p sealed under a 16-byte nonce: the ciphertext and the 16-byte tag
263	def AdAegis128.seal(kk: AdAegisKey, nonce: B.Octets, aad: B.Octets, p: B.Octets) -> Result<&2, &2, Cr.CrError, B.Octets>:
264	  ad.ae.seal(kk, nonce, aad, p)
265	
266	def ad.ae.tag_of.of(+sb: R.AdSbox, +aad: B.Octets, +c: B.Octets, o: AdAegisOut) -> B.Octets:
267	  match o:
268	    case AdAegisOut{r, out}:
269	      ad.ae.final(sb, ad.ae.flush(sb, r), B.Octets.len(aad), B.Octets.len(c))
270	
271	def ad.ae.tag_of.go(kk: AdAegisKey, +nonce: B.Octets, +aad: B.Octets, +c: B.Octets) -> B.Octets:
272	  match kk:
273	    case AdAegisKey{+k, +sb}:
274	      ad.ae.tag_of.of(sb, aad, c, ad.ae.dgo(B.Octets.to_list(c), sb, ad.ae.start(sb, k, nonce, aad), Nil{}))
275	
276	# The tag open expects for the ciphertext c (without its tag)
277	def AdAegis128.tag_of(kk: AdAegisKey, +nonce: B.Octets, +aad: B.Octets, +c: B.Octets) -> B.Octets:
278	  ad.ae.tag_of.go(kk, nonce, aad, c)
279	
280	def ad.ae.plain.of(o: AdAegisOut) -> B.Octets:
281	  match o:
282	    case AdAegisOut{r, out}:
283	      B.Octets.of_list(out)
284	
285	def ad.ae.plain(kk: AdAegisKey, +nonce: B.Octets, +aad: B.Octets, +c: B.Octets) -> B.Octets:
286	  match kk:
287	    case AdAegisKey{+k, +sb}:
288	      ad.ae.plain.of(ad.ae.dgo(B.Octets.to_list(c), sb, ad.ae.start(sb, k, nonce, aad), Nil{}))
289	
290	def ad.ae.check(ok: Bool, +p: B.Octets) -> Result<&2, &2, Cr.CrError, B.Octets>:
291	  match ok:
292	    case True{}:
293	      Done{p}
294	    case False{}:
295	      Fail{Cr.CrTagMismatch{}}
296	
297	def ad.ae.open.go(+kk: AdAegisKey, +nonce: B.Octets, +aad: B.Octets, s: Ad.AdSplit) -> Result<&2, &2, Cr.CrError, B.Octets>:
298	  match s:
299	    case Ad.AdSplit{+c, t}:
300	      ad.ae.check(Ad.AdTag.is_eq(AdAegis128.tag_of(kk, nonce, aad, c), t), ad.ae.plain(kk, nonce, aad, c))
301	
302	def ad.ae.open(+kk: AdAegisKey, +nonce: B.Octets, +aad: B.Octets, +c: B.Octets) -> Result<&2, &2, Cr.CrError, B.Octets>:
303	  Result.bind(&2, &2, Cr.CrError, Unit, B.Octets, ad.ae.nonce(nonce), u => Result.bind(&2, &2, Cr.CrError, Ad.AdSplit, B.Octets, Ad.AdTag.split(16n, c), s => ad.ae.open.go(kk, nonce, aad, s)))
304	
305	# c (the ciphertext with its tag) opened: the plaintext, or CrTagMismatch
306	# (no plaintext released)
307	def AdAegis128.open(kk: AdAegisKey, nonce: B.Octets, aad: B.Octets, c: B.Octets) -> Result<&2, &2, Cr.CrError, B.Octets>:
308	  ad.ae.open(kk, nonce, aad, c)