1	import Base
2	import ./types.bend as T
3	import ./sbox.bend as B
4	
5	# AES (FIPS 197) block encryption for 128-, 192- and 256-bit keys. The
6	# state is four columns of four bytes (U32 values below 256); the S-box and
7	# the field doublings are the constant-time circuits of sbox.bend. The key
8	# schedule is expanded once into the list of round keys.
9	
10	def sub_word(w: T.Quad) -> T.Quad:
11	  match w:
12	    case T.W{a, b, c, d}:
13	      T.W{B.sbox(a), B.sbox(b), B.sbox(c), B.sbox(d)}
14	
15	def sub_bytes(s: T.State) -> T.State:
16	  match s:
17	    case T.S{c0, c1, c2, c3}:
18	      T.S{sub_word(c0), sub_word(c1), sub_word(c2), sub_word(c3)}
19	
20	# Row r moves r columns to the left.
21	def shift_rows(s: T.State) -> T.State:
22	  match s:
23	    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}}:
24	      T.S{T.W{a0, b1, c2, d3}, T.W{b0, c1, d2, a3}, T.W{c0, d1, a2, b3}, T.W{d0, a1, b2, c3}}
25	
26	def mix_column(w: T.Quad) -> T.Quad:
27	  match w:
28	    case T.W{+a, +b, +c, +d}:
29	      T.W{U32.xor(U32.xor(U32.xor(B.xtime(a), B.mul3(b)), c), d),
30	        U32.xor(U32.xor(U32.xor(a, B.xtime(b)), B.mul3(c)), d),
31	        U32.xor(U32.xor(U32.xor(a, b), B.xtime(c)), B.mul3(d)),
32	        U32.xor(U32.xor(U32.xor(B.mul3(a), b), c), B.xtime(d))}
33	
34	def mix_columns(s: T.State) -> T.State:
35	  match s:
36	    case T.S{c0, c1, c2, c3}:
37	      T.S{mix_column(c0), mix_column(c1), mix_column(c2), mix_column(c3)}
38	
39	def xor_word(a: T.Quad, b: T.Quad) -> T.Quad:
40	  match a b:
41	    case T.W{a0, a1, a2, a3} T.W{b0, b1, b2, b3}:
42	      T.W{U32.xor(a0, b0), U32.xor(a1, b1), U32.xor(a2, b2), U32.xor(a3, b3)}
43	
44	def add_round_key(s: T.State, k: T.State) -> T.State:
45	  match s k:
46	    case T.S{c0, c1, c2, c3} T.S{k0, k1, k2, k3}:
47	      T.S{xor_word(c0, k0), xor_word(c1, k1), xor_word(c2, k2), xor_word(c3, k3)}
48	
49	def round(s: T.State, k: T.State) -> T.State:
50	  add_round_key(mix_columns(shift_rows(sub_bytes(s))), k)
51	
52	def last_round(s: T.State, k: T.State) -> T.State:
53	  add_round_key(shift_rows(sub_bytes(s)), k)
54	
55	# n full rounds, then the final round.
56	def rounds(n: Nat, s: T.State, ks: List<&2, T.State>) -> T.State:
57	  match n ks:
58	    case 0n k <> rest:
59	      last_round(s, k)
60	    case 1n+p k <> rest:
61	      rounds(p, round(s, k), rest)
62	    case _ _:
63	      s
64	
65	def cipher(+nr: Nat, s: T.State, ks: List<&2, T.State>) -> T.State:
66	  match ks:
67	    case k <> rest:
68	      rounds(Nat.sub(nr, 1n), add_round_key(s, k), rest)
69	    case Nil{}:
70	      s
71	
72	# ---- key schedule ----
73	
74	def rot_word(w: T.Quad) -> T.Quad:
75	  match w:
76	    case T.W{a, b, c, d}:
77	      T.W{b, c, d, a}
78	
79	# x^n in GF(2^8), by doubling.
80	def xpow(n: Nat) -> U32:
81	  match n:
82	    case 0n:
83	      1
84	    case 1n+p:
85	      B.xtime(xpow(p))
86	
87	def rcon(j: Nat) -> T.Quad:
88	  T.W{xpow(Nat.sub(j, 1n)), 0, 0, 0}
89	
90	def key_words(key: List<&2, U32>) -> List<&2, T.Quad>:
91	  match key:
92	    case a <> b <> c <> d <> rest:
93	      T.W{a, b, c, d} <> key_words(rest)
94	    case _:
95	      Nil{}
96	
97	def get_word(ws: List<&2, T.Quad>, n: Nat) -> T.Quad:
98	  match ws n:
99	    case Nil{} _:
100	      T.W{0, 0, 0, 0}
101	    case w <> rest 0n:
102	      w
103	    case w <> rest 1n+p:
104	      get_word(rest, p)
105	
106	def schedule_core(k: Nat, big: Bool, +q: Nat, temp: T.Quad) -> T.Quad:
107	  match k:
108	    case 0n:
109	      xor_word(sub_word(rot_word(temp)), rcon(q))
110	    case 1n:
111	      temp
112	    case 2n:
113	      temp
114	    case 3n:
115	      temp
116	    case 4n:
117	      match big:
118	        case True{}:
119	          sub_word(temp)
120	        case False{}:
121	          temp
122	    case 5n+p:
123	      temp
124	
125	# history is the words so far, the newest first.
126	def next_word(+nk: Nat, +i: Nat, +history: List<&2, T.Quad>) -> T.Quad:
127	  xor_word(get_word(history, Nat.sub(nk, 1n)),
128	    schedule_core(Nat.mod(i, nk), Nat.is_lt(6n, nk), Nat.div(i, nk), get_word(history, 0n)))
129	
130	def grow(n: Nat, +nk: Nat, +i: Nat, +history: List<&2, T.Quad>) -> List<&2, T.Quad>:
131	  match n:
132	    case 0n:
133	      history
134	    case 1n+p:
135	      grow(p, nk, 1n+i, next_word(nk, i, history) <> history)
136	
137	# Four consecutive words make a round key.
138	def round_keys(ws: List<&2, T.Quad>) -> List<&2, T.State>:
139	  match ws:
140	    case w0 <> w1 <> w2 <> w3 <> rest:
141	      T.S{w0, w1, w2, w3} <> round_keys(rest)
142	    case _:
143	      Nil{}
144	
145	def expand(+nk: Nat, +nr: Nat, key: List<&2, U32>) -> List<&2, T.State>:
146	  round_keys(List.reverse(&2, T.Quad, grow(Nat.sub(Nat.mul(4n, 1n+nr), nk), nk, nk,
147	    List.reverse(&2, T.Quad, key_words(key)))))
148	
149	def encrypt(+nk: Nat, +nr: Nat, key: List<&2, U32>, block: T.State) -> T.State:
150	  cipher(nr, block, expand(nk, nr, key))
151	
152	# ---- byte API ----
153	
154	# An expanded key: the number of rounds and the round keys.
155	type Schedule is Data:
156	  Schedule{rounds: Nat, keys: List<&2, T.State>}
157	
158	# FIPS 197 Figure 4: a key of 16, 24 or 32 bytes; Nk = its length in words,
159	# Nr = Nk + 6.
160	def valid_length(+n: Nat) -> Bool:
161	  Bool.or(Nat.is_eq(n, 16n), Bool.or(Nat.is_eq(n, 24n), Nat.is_eq(n, 32n)))
162	
163	def schedule_if(ok: Bool, +nk: Nat, +key: List<&2, U32>) -> Maybe<&2, Schedule>:
164	  match ok:
165	    case True{}:
166	      Some{Schedule{Nat.add(nk, 6n), expand(nk, Nat.add(nk, 6n), key)}}
167	    case False{}:
168	      None{}
169	
170	# A 16-, 24- or 32-byte key (AES-128, AES-192, AES-256); None otherwise.
171	def expand_key(+key: List<&2, U32>) -> Maybe<&2, Schedule>:
172	  +n = List.length(&2, U32, key)
173	  schedule_if(valid_length(n), Nat.div(n, 4n), key)
174	
175	def encrypt_state(sched: Schedule, s: T.State) -> T.State:
176	  match sched:
177	    case Schedule{nr, ks}:
178	      cipher(nr, s, ks)
179	
180	def state_of(bytes: List<&2, U32>) -> Maybe<&2, T.State>:
181	  match bytes:
182	    case a0 <> a1 <> a2 <> a3 <> b0 <> b1 <> b2 <> b3 <> c0 <> c1 <> c2 <> c3 <> d0 <> d1 <> d2 <> d3 <> Nil{}:
183	      Some{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}}}
184	    case _:
185	      None{}
186	
187	def bytes_of(s: T.State) -> List<&2, U32>:
188	  match s:
189	    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}}:
190	      [a0, a1, a2, a3, b0, b1, b2, b3, c0, c1, c2, c3, d0, d1, d2, d3]
191	
192	def encrypt_bytes(sched: Schedule, m: Maybe<&2, T.State>) -> Maybe<&2, List<&2, U32>>:
193	  match m:
194	    case None{}:
195	      None{}
196	    case Some{s}:
197	      Some{bytes_of(encrypt_state(sched, s))}
198	
199	# Encrypts one 16-byte block; None for another length.
200	def encrypt_block(sched: Schedule, block: List<&2, U32>) -> Maybe<&2, List<&2, U32>>:
201	  encrypt_bytes(sched, state_of(block))
202	
203	def hex_digit_if(x: U32, small: Bool) -> Char:
204	  match small:
205	    case True{}:
206	      Chr{U32.add(48, x)}
207	    case False{}:
208	      Chr{U32.add(87, x)}
209	
210	def hex_digit(+x: U32) -> Char:
211	  hex_digit_if(x, U32.is_lt(x, 10))
212	
213	# Lowercase hexadecimal of a byte list.
214	def hex(bytes: List<&2, U32>) -> String:
215	  match bytes:
216	    case Nil{}:
217	      ""
218	    case +b <> rest:
219	      SCon{hex_digit(U32.and(15, U32.shrn(b, 4n))), SCon{hex_digit(U32.and(15, b)), hex(rest)}}