1	import Base
2	import ./types.bend as T
3	import ./aes.bend as A
4	import ../subtle.bend as Subtle
5	
6	# AES-GCM (NIST SP 800-38D) with 96-bit nonces and 128-bit tags.
7	#
8	# GHASH works on 128-bit blocks held as four big-endian U32 words: bit 0 of
9	# a block (the coefficient of x^0 in GCM's bit order, the leftmost bit) is
10	# bit 31 of the first word. The product X * H is Algorithm 1 of SP 800-38D,
11	# one bit of X at a time, with masks instead of branches: Z ^= V & (0 - x_i),
12	# then V = V * x (a right shift that folds the dropped bit back in as
13	# R = 0xe1 || 0^120, again through a mask). No table is indexed by secret
14	# data. The tag comparison is subtle.eq.
15	
16	type Block is Data:
17	  B{w0: U32, w1: U32, w2: U32, w3: U32}
18	
19	type Acc is Data:
20	  Acc{z: Block, v: Block}
21	
22	def zero() -> Block:
23	  B{0, 0, 0, 0}
24	
25	# V * x: the block shifted one bit toward x^127, x^128 = R.
26	def mulx(v: Block) -> Block:
27	  match v:
28	    case B{+a, +b, +c, +d}:
29	      B{U32.xor(U32.and(3774873600, U32.sub(0, U32.and(1, d))), U32.shrn(a, 1n)),
30	        U32.or(U32.shln(a, 31n), U32.shrn(b, 1n)),
31	        U32.or(U32.shln(b, 31n), U32.shrn(c, 1n)),
32	        U32.or(U32.shln(c, 31n), U32.shrn(d, 1n))}
33	
34	# Z xor (V and (0 - bit)), bit being 0 or 1.
35	def add_masked(+bit: U32, z: Block, v: Block) -> Block:
36	  match z v:
37	    case B{z0, z1, z2, z3} B{v0, v1, v2, v3}:
38	      +m = U32.sub(0, bit)
39	      B{U32.xor(z0, U32.and(m, v0)), U32.xor(z1, U32.and(m, v1)),
40	        U32.xor(z2, U32.and(m, v2)), U32.xor(z3, U32.and(m, v3))}
41	
42	def step(+bit: U32, acc: Acc) -> Acc:
43	  match acc:
44	    case Acc{z, +v}:
45	      Acc{add_masked(bit, z, v), mulx(v)}
46	
47	# The eight bits of a byte, the most significant first.
48	def mul_byte(+a: U32, acc: Acc) -> Acc:
49	  step(U32.and(1, a), step(U32.and(1, U32.shrn(a, 1n)), step(U32.and(1, U32.shrn(a, 2n)),
50	    step(U32.and(1, U32.shrn(a, 3n)), step(U32.and(1, U32.shrn(a, 4n)),
51	    step(U32.and(1, U32.shrn(a, 5n)), step(U32.and(1, U32.shrn(a, 6n)),
52	    step(U32.and(1, U32.shrn(a, 7n)), acc))))))))
53	
54	def mul_bytes(xs: List<&2, U32>, acc: Acc) -> Acc:
55	  match xs:
56	    case Nil{}:
57	      acc
58	    case x <> rest:
59	      mul_bytes(rest, mul_byte(x, acc))
60	
61	def acc_z(acc: Acc) -> Block:
62	  match acc:
63	    case Acc{z, v}:
64	      z
65	
66	# The 16 bytes of a block, and the block of 16 bytes (big-endian words).
67	def unpack(z: Block) -> List<&2, U32>:
68	  match z:
69	    case B{+a, +b, +c, +d}:
70	      [U32.and(255, U32.shrn(a, 24n)), U32.and(255, U32.shrn(a, 16n)), U32.and(255, U32.shrn(a, 8n)), U32.and(255, a),
71	       U32.and(255, U32.shrn(b, 24n)), U32.and(255, U32.shrn(b, 16n)), U32.and(255, U32.shrn(b, 8n)), U32.and(255, b),
72	       U32.and(255, U32.shrn(c, 24n)), U32.and(255, U32.shrn(c, 16n)), U32.and(255, U32.shrn(c, 8n)), U32.and(255, c),
73	       U32.and(255, U32.shrn(d, 24n)), U32.and(255, U32.shrn(d, 16n)), U32.and(255, U32.shrn(d, 8n)), U32.and(255, d)]
74	
75	def word(q: T.Quad) -> U32:
76	  match q:
77	    case T.W{a, b, c, d}:
78	      U32.or(U32.shln(U32.and(255, a), 24n), U32.or(U32.shln(U32.and(255, b), 16n),
79	        U32.or(U32.shln(U32.and(255, c), 8n), U32.and(255, d))))
80	
81	# An AES state as a block: its columns are the words.
82	def pack(s: T.State) -> Block:
83	  match s:
84	    case T.S{c0, c1, c2, c3}:
85	      B{word(c0), word(c1), word(c2), word(c3)}
86	
87	def xor_bytes(xs: List<&2, U32>, ks: List<&2, U32>) -> List<&2, U32>:
88	  match xs ks:
89	    case x <> xt k <> kt:
90	      U32.xor(x, k) <> xor_bytes(xt, kt)
91	    case _ _:
92	      Nil{}
93	
94	# (Y xor X) * H for a 16-byte X.
95	def absorb(+h: Block, y: Block, x: List<&2, U32>) -> Block:
96	  acc_z(mul_bytes(xor_bytes(unpack(y), x), Acc{zero(), h}))
97	
98	def zeros(n: Nat) -> List<&2, U32>:
99	  List.replicate(U32, n, 0)
100	
101	# GHASH_H over the bytes, a last partial block padded with zero bytes.
102	def ghash(+h: Block, +xs: List<&2, U32>, y: Block) -> Block:
103	  match xs:
104	    case x0 <> x1 <> x2 <> x3 <> x4 <> x5 <> x6 <> x7 <> x8 <> x9 <> x10 <> x11 <> x12 <> x13 <> x14 <> x15 <> rest:
105	      ghash(h, rest, absorb(h, y, [x0, x1, x2, x3, x4, x5, x6, x7, x8, x9, x10, x11, x12, x13, x14, x15]))
106	    case Nil{}:
107	      y
108	    case _:
109	      absorb(h, y, List.append(&2, U32, xs, zeros(Nat.sub(16n, List.length(&2, U32, xs)))))
110	
111	# [n]_(8*width), most significant byte first.
112	def octets(width: Nat, +n: Nat) -> List<&2, U32>:
113	  match width:
114	    case 0n:
115	      Nil{}
116	    case 1n+p:
117	      List.append(&2, U32, octets(p, Nat.div(n, 256n)), [U32.from_nat(Nat.mod(n, 256n))])
118	
119	# [len(A)]_64 || [len(C)]_64, in bits.
120	def lengths(+aad: List<&2, U32>, +c: List<&2, U32>) -> List<&2, U32>:
121	  List.append(&2, U32, octets(8n, Nat.mul(8n, List.length(&2, U32, aad))),
122	    octets(8n, Nat.mul(8n, List.length(&2, U32, c))))
123	
124	# ---- counter mode ----
125	
126	# The counter block: the 96-bit nonce (three words) and a 32-bit counter.
127	def counter(+n0: T.Quad, +n1: T.Quad, +n2: T.Quad, +c: U32) -> T.State:
128	  T.S{n0, n1, n2, T.W{U32.and(255, U32.shrn(c, 24n)), U32.and(255, U32.shrn(c, 16n)), U32.and(255, U32.shrn(c, 8n)), U32.and(255, c)}}
129	
130	def keystream(+sched: A.Schedule, +n0: T.Quad, +n1: T.Quad, +n2: T.Quad, +c: U32) -> List<&2, U32>:
131	  A.bytes_of(A.encrypt_state(sched, counter(n0, n1, n2, c)))
132	
133	def gctr(+sched: A.Schedule, +xs: List<&2, U32>, +n0: T.Quad, +n1: T.Quad, +n2: T.Quad, +c: U32) -> List<&2, U32>:
134	  match xs:
135	    case x0 <> x1 <> x2 <> x3 <> x4 <> x5 <> x6 <> x7 <> x8 <> x9 <> x10 <> x11 <> x12 <> x13 <> x14 <> x15 <> rest:
136	      List.append(&2, U32, xor_bytes([x0, x1, x2, x3, x4, x5, x6, x7, x8, x9, x10, x11, x12, x13, x14, x15], keystream(sched, n0, n1, n2, c)),
137	        gctr(sched, rest, n0, n1, n2, U32.inc(c)))
138	    case _:
139	      xor_bytes(xs, keystream(sched, n0, n1, n2, c))
140	
141	# ---- GCM ----
142	
143	def hash_key(+sched: A.Schedule) -> Block:
144	  pack(A.encrypt_state(sched, T.S{T.W{0, 0, 0, 0}, T.W{0, 0, 0, 0}, T.W{0, 0, 0, 0}, T.W{0, 0, 0, 0}}))
145	
146	# The tag of ciphertext c: GHASH over A || 0^v || C || 0^u || lengths,
147	# encrypted with the counter block J0 = nonce || 1.
148	def tag(+sched: A.Schedule, +n0: T.Quad, +n1: T.Quad, +n2: T.Quad, +aad: List<&2, U32>, +c: List<&2, U32>) -> List<&2, U32>:
149	  +h = hash_key(sched)
150	  gctr(sched, unpack(ghash(h, lengths(aad, c), ghash(h, c, ghash(h, aad, zero())))), n0, n1, n2, 1)
151	
152	# C || T; the counter of the first data block is 2 (inc32(J0)).
153	def seal_core(+sched: A.Schedule, +n0: T.Quad, +n1: T.Quad, +n2: T.Quad, +aad: List<&2, U32>, pt: List<&2, U32>) -> List<&2, U32>:
154	  +c = gctr(sched, pt, n0, n1, n2, 2)
155	  List.append(&2, U32, c, tag(sched, n0, n1, n2, aad, c))
156	
157	def accept(ok: Bool, p: List<&2, U32>) -> Maybe<&2, List<&2, U32>>:
158	  match ok:
159	    case True{}:
160	      Some{p}
161	    case False{}:
162	      None{}
163	
164	# The plaintext when t is the tag of c; None otherwise (compared in constant time).
165	def open_checked(+sched: A.Schedule, +n0: T.Quad, +n1: T.Quad, +n2: T.Quad, +aad: List<&2, U32>, +c: List<&2, U32>, t: List<&2, U32>) -> Maybe<&2, List<&2, U32>>:
166	  accept(Subtle.eq(t, tag(sched, n0, n1, n2, aad, c)), gctr(sched, c, n0, n1, n2, 2))
167	
168	def open_split(+sched: A.Schedule, +n0: T.Quad, +n1: T.Quad, +n2: T.Quad, +aad: List<&2, U32>, +input: List<&2, U32>, +n: Nat, short: Bool) -> Maybe<&2, List<&2, U32>>:
169	  match short:
170	    case True{}:
171	      None{}
172	    case False{}:
173	      open_checked(sched, n0, n1, n2, aad, List.take(&2, U32, input, Nat.sub(n, 16n)), List.drop(&2, U32, input, Nat.sub(n, 16n)))
174	
175	def open_core(+sched: A.Schedule, +n0: T.Quad, +n1: T.Quad, +n2: T.Quad, +aad: List<&2, U32>, +input: List<&2, U32>) -> Maybe<&2, List<&2, U32>>:
176	  +n = List.length(&2, U32, input)
177	  open_split(sched, n0, n1, n2, aad, input, n, Nat.is_lt(n, 16n))
178	
179	# ---- byte API ----
180	
181	def seal_nonce(+sched: A.Schedule, nonce: List<&2, U32>, +aad: List<&2, U32>, +pt: List<&2, U32>) -> Maybe<&2, List<&2, U32>>:
182	  match nonce:
183	    case n0 <> n1 <> n2 <> n3 <> n4 <> n5 <> n6 <> n7 <> n8 <> n9 <> n10 <> n11 <> Nil{}:
184	      Some{seal_core(sched, T.W{n0, n1, n2, n3}, T.W{n4, n5, n6, n7}, T.W{n8, n9, n10, n11}, aad, pt)}
185	    case _:
186	      None{}
187	
188	def open_nonce(+sched: A.Schedule, nonce: List<&2, U32>, +aad: List<&2, U32>, +input: List<&2, U32>) -> Maybe<&2, List<&2, U32>>:
189	  match nonce:
190	    case n0 <> n1 <> n2 <> n3 <> n4 <> n5 <> n6 <> n7 <> n8 <> n9 <> n10 <> n11 <> Nil{}:
191	      open_core(sched, T.W{n0, n1, n2, n3}, T.W{n4, n5, n6, n7}, T.W{n8, n9, n10, n11}, aad, input)
192	    case _:
193	      None{}
194	
195	def seal_key(m: Maybe<&2, A.Schedule>, +nonce: List<&2, U32>, +aad: List<&2, U32>, +pt: List<&2, U32>) -> Maybe<&2, List<&2, U32>>:
196	  match m:
197	    case None{}:
198	      None{}
199	    case Some{sched}:
200	      seal_nonce(sched, nonce, aad, pt)
201	
202	def open_key(m: Maybe<&2, A.Schedule>, +nonce: List<&2, U32>, +aad: List<&2, U32>, +input: List<&2, U32>) -> Maybe<&2, List<&2, U32>>:
203	  match m:
204	    case None{}:
205	      None{}
206	    case Some{sched}:
207	      open_nonce(sched, nonce, aad, input)
208	
209	def schedule_if(ok: Bool, +key: List<&2, U32>) -> Maybe<&2, A.Schedule>:
210	  match ok:
211	    case True{}:
212	      A.expand_key(key)
213	    case False{}:
214	      None{}
215	
216	# The expanded key when the key has the given length.
217	def schedule_of(+width: Nat, +key: List<&2, U32>) -> Maybe<&2, A.Schedule>:
218	  schedule_if(Nat.is_eq(List.length(&2, U32, key), width), key)