aes.bend

bend-osdev-cipher-modes@0.1.0.0 · 291 lines · 46 definitions

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.
11import Base
12import bend-osdev-bytes@0.1.0.0/bytes.bend as B
13import bend-osdev-cipher-modes@0.1.0.0/cipher_modes.bend as M
14import 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
18def 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
21def 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
28def 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
32def 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)
41def aes.inv(+x: U32) -> U32:
42  +a = aes.inv.go(6n, x, x)
43  aes.gmul(a, a)
44
45def 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
49def 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
54def 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
60type 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
63def aes.byte(+w: U32, n: Nat) -> U32:
64  U32.and(255, U32.shrn(w, n))
65
66def 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
71def 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
76def 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)
83def 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
89def aes.m9(+a: U32) -> U32:
90  U32.xor(aes.xt(aes.xt(aes.xt(a))), a)
91
92def aes.m11(+a: U32) -> U32:
93  +a2 = aes.xt(a)
94  U32.xor(U32.xor(aes.xt(aes.xt(a2)), a2), a)
95
96def aes.m13(+a: U32) -> U32:
97  +a4 = aes.xt(aes.xt(a))
98  U32.xor(U32.xor(aes.xt(a4), a4), a)
99
100def 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
105def 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
109def 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)
117def 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
124def 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
135def 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
138def aes.rot_word(+w: U32) -> U32:
139  U32.or(U32.shln(w, 8n), U32.shrn(w, 24n))
140
141def 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
149def 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
157def 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
165def 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
172def aes.head(xs: List<&2, U32>) -> U32:
173  match xs:
174    case Nil{}:
175      0
176    case Con{h, t}:
177      h
178
179def 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
186def 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
195def 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
204def 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
211def 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
219def 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.
229type CmAesKey is Data:
230  CmAesKey{sched: A.Schedule, raw: List<&2, U32>, dkeys: List<&2, CmAesState>}
231
232def 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
239def 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
247def 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
255def 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
260def 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
264def 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
268def 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)
276def 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
282def 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)
288def 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)))