src/crypto/aes/aes.bend

bend-collections@1.0.0.0 · 219 lines · 34 definitions

1import Base
2import bend-collections@1.0.0.0/src/crypto/aes/types.bend as T
3import bend-collections@1.0.0.0/src/crypto/aes/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
10def 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
15def 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.
21def 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
26def 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
34def 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
39def 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
44def 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
49def round(s: T.State, k: T.State) -> T.State:
50  add_round_key(mix_columns(shift_rows(sub_bytes(s))), k)
51
52def 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.
56def 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
65def 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
74def 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.
80def xpow(n: Nat) -> U32:
81  match n:
82    case 0n:
83      1
84    case 1n+p:
85      B.xtime(xpow(p))
86
87def rcon(j: Nat) -> T.Quad:
88  T.W{xpow(Nat.sub(j, 1n)), 0, 0, 0}
89
90def 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
97def 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
106def 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.
126def 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
130def 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.
138def 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
145def 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
149def 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.
155type 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.
160def 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
163def 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.
171def 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
175def encrypt_state(sched: Schedule, s: T.State) -> T.State:
176  match sched:
177    case Schedule{nr, ks}:
178      cipher(nr, s, ks)
179
180def 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
187def 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
192def 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.
200def encrypt_block(sched: Schedule, block: List<&2, U32>) -> Maybe<&2, List<&2, U32>>:
201  encrypt_bytes(sched, state_of(block))
202
203def 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
210def hex_digit(+x: U32) -> Char:
211  hex_digit_if(x, U32.is_lt(x, 10))
212
213# Lowercase hexadecimal of a byte list.
214def 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)}}