aes.bend

bend-osdev-aes@0.1.0.0 · 262 lines · 44 definitions

1# AES-128, AES-192 and AES-256 (FIPS 197) as block ciphers for the modes:
2# the named 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): it
7# undoes bend-collections' cipher round by round, with that cipher's own
8# round keys (its Schedule), the last one first. The inverse S-box is
9# computed, not looked up: the inverse affine map, then the inverse in
10# GF(2^8) as x^254, with masks instead of branches, so no branch or table
11# index depends on a key or data byte; InvMixColumns' products by {09},
12# {0b}, {0d} and {0e} are sums of doublings (xtime), likewise branch-free.
13#
14# Both directions use the round keys with each byte masked to its low 8
15# bits (aes.sched), which changes nothing for a key made by CmAes.key (its
16# bytes are below 256) and lets LAWS.bend state the round trip for every
17# value of CmAesKey.
18import Base
19import bend-osdev-bytes@0.1.1.0/bytes.bend as B
20import bend-osdev-mode-core@0.1.0.0/mode_core.bend as M
21import bend-collections@1.0.0.0/src/crypto/aes/aes.bend as A
22import bend-collections@1.0.0.0/src/crypto/aes/types.bend as T
23
24# GF(2^8) modulo x^8 + x^4 + x^3 + x + 1
25
26def aes.xt(+x: U32) -> U32:
27  U32.and(255, U32.xor(U32.shl(x), U32.and(27, M.cm.mask(U32.shrn(x, 7n)))))
28
29def aes.gmul.go(n: Nat, +a: U32, +b: U32, p: U32) -> U32:
30  match n:
31    case 0n:
32      p
33    case 1n+q:
34      aes.gmul.go(q, aes.xt(a), U32.shr(b), U32.xor(p, U32.and(a, M.cm.mask(b))))
35
36def aes.gmul(+a: U32, +b: U32) -> U32:
37  aes.gmul.go(8n, a, b, 0)
38
39# a := a^2 * x, n times: from a = x, six times gives x^127
40def aes.inv.go(n: Nat, a: U32, +x: U32) -> U32:
41  match n:
42    case 0n:
43      a
44    case 1n+q:
45      +s = a
46      aes.inv.go(q, aes.gmul(aes.gmul(s, s), x), x)
47
48# the multiplicative inverse, x^254 (0 for 0)
49def aes.inv(+x: U32) -> U32:
50  +a = aes.inv.go(6n, x, x)
51  aes.gmul(a, a)
52
53def aes.rotl(+b: U32, +n: Nat) -> U32:
54  U32.and(255, U32.or(U32.shln(b, n), U32.shrn(b, Nat.sub(8n, n))))
55
56# the inverse S-box: the inverse affine map, then the inverse
57def aes.isbox(x: U32) -> U32:
58  +b = U32.and(255, x)
59  aes.inv(U32.xor(U32.xor(U32.xor(aes.rotl(b, 1n), aes.rotl(b, 3n)), aes.rotl(b, 6n)), 5))
60
61# The state is bend-collections' (T.State: four columns T.W{r0, r1, r2, r3},
62# byte 4c + r of the block in row r of column c)
63
64# each byte's low 8 bits
65def aes.mask_w(w: T.Quad) -> T.Quad:
66  match w:
67    case T.W{a, b, c, d}:
68      T.W{U32.and(255, a), U32.and(255, b), U32.and(255, c), U32.and(255, d)}
69
70def aes.mask(s: T.State) -> T.State:
71  match s:
72    case T.S{c0, c1, c2, c3}:
73      T.S{aes.mask_w(c0), aes.mask_w(c1), aes.mask_w(c2), aes.mask_w(c3)}
74
75def aes.mask_keys(ks: List<&2, T.State>) -> List<&2, T.State>:
76  match ks:
77    case Nil{}:
78      Nil{}
79    case Con{k, rest}:
80      Con{aes.mask(k), aes.mask_keys(rest)}
81
82# a schedule with every round key byte masked
83def aes.sched(s: A.Schedule) -> A.Schedule:
84  match s:
85    case A.Schedule{nr, ks}:
86      A.Schedule{nr, aes.mask_keys(ks)}
87
88# InvSubBytes: the inverse S-box on every byte
89def aes.isub_w(w: T.Quad) -> T.Quad:
90  match w:
91    case T.W{a, b, c, d}:
92      T.W{aes.isbox(a), aes.isbox(b), aes.isbox(c), aes.isbox(d)}
93
94def aes.inv_sub(s: T.State) -> T.State:
95  match s:
96    case T.S{c0, c1, c2, c3}:
97      T.S{aes.isub_w(c0), aes.isub_w(c1), aes.isub_w(c2), aes.isub_w(c3)}
98
99# InvShiftRows: row r moves r columns to the right
100def aes.inv_shift(s: T.State) -> T.State:
101  match s:
102    case T.S{T.W{p0, p1, p2, p3}, T.W{q0, q1, q2, q3}, T.W{r0, r1, r2, r3}, T.W{s0, s1, s2, s3}}:
103      T.S{T.W{p0, s1, r2, q3}, T.W{q0, p1, s2, r3}, T.W{r0, q1, p2, s3}, T.W{s0, r1, q2, p3}}
104
105# InvMixColumns: each column times {0e, 0b, 0d, 09}, rotated per row
106# a byte times {09}, {0b}, {0d}, {0e}: x^3 + 1, x^3 + x + 1, x^3 + x^2 + 1,
107# x^3 + x^2 + x, as sums of doublings
108def aes.c9(+b: U32) -> U32:
109  U32.xor(aes.xt(aes.xt(aes.xt(b))), b)
110
111def aes.c11(+b: U32) -> U32:
112  +b2 = aes.xt(b)
113  U32.xor(U32.xor(aes.xt(aes.xt(b2)), b2), b)
114
115def aes.c13(+b: U32) -> U32:
116  +b4 = aes.xt(aes.xt(b))
117  U32.xor(U32.xor(aes.xt(b4), b4), b)
118
119def aes.c14(+b: U32) -> U32:
120  +b2 = aes.xt(b)
121  +b4 = aes.xt(b2)
122  U32.xor(U32.xor(aes.xt(b4), b4), b2)
123
124# the same products of a word's low byte
125def aes.m9(a: U32) -> U32:
126  aes.c9(U32.and(255, a))
127
128def aes.m11(a: U32) -> U32:
129  aes.c11(U32.and(255, a))
130
131def aes.m13(a: U32) -> U32:
132  aes.c13(U32.and(255, a))
133
134def aes.m14(a: U32) -> U32:
135  aes.c14(U32.and(255, a))
136
137def aes.imix0(a: U32, b: U32, c: U32, d: U32) -> U32:
138  U32.xor(U32.xor(aes.m14(a), aes.m11(b)), U32.xor(aes.m13(c), aes.m9(d)))
139
140def aes.imix_w(w: T.Quad) -> T.Quad:
141  match w:
142    case T.W{+a, +b, +c, +d}:
143      T.W{aes.imix0(a, b, c, d), aes.imix0(b, c, d, a), aes.imix0(c, d, a, b), aes.imix0(d, a, b, c)}
144
145def aes.inv_mix(s: T.State) -> T.State:
146  match s:
147    case T.S{c0, c1, c2, c3}:
148      T.S{aes.imix_w(c0), aes.imix_w(c1), aes.imix_w(c2), aes.imix_w(c3)}
149
150# A round undone (round key k): AddRoundKey, InvMixColumns, InvShiftRows,
151# InvSubBytes; the last round has no MixColumns to undo
152def aes.inv_round(c: T.State, k: T.State) -> T.State:
153  aes.inv_sub(aes.inv_shift(aes.inv_mix(A.add_round_key(c, k))))
154
155def aes.inv_last(c: T.State, k: T.State) -> T.State:
156  aes.inv_sub(aes.inv_shift(A.add_round_key(c, k)))
157
158# bend-collections' rounds (n full rounds, then the last, a round key each)
159# undone, the last round key first
160def aes.inv_rounds(n: Nat, ks: List<&2, T.State>, c: T.State) -> T.State:
161  match n ks:
162    case 0n Con{k, rest}:
163      aes.inv_last(c, k)
164    case 1n+p Con{k, rest}:
165      aes.inv_round(aes.inv_rounds(p, rest, c), k)
166    case _ _:
167      aes.mask(c)
168
169# bend-collections' cipher (AddRoundKey with the first round key, then the
170# rounds) undone
171def aes.inv_cipher(+nr: Nat, ks: List<&2, T.State>, c: T.State) -> T.State:
172  match ks:
173    case Con{k, rest}:
174      aes.mask(A.add_round_key(aes.inv_rounds(Nat.sub(nr, 1n), rest, c), k))
175    case Nil{}:
176      aes.mask(c)
177
178# a block's 16 bytes as a state, and back
179def aes.state16(xs: List<&2, U32>) -> T.State:
180  match xs:
181    case Con{a0, Con{a1, Con{a2, Con{a3, Con{b0, Con{b1, Con{b2, Con{b3, Con{c0, Con{c1, Con{c2, Con{c3, Con{d0, Con{d1, Con{d2, Con{d3, t}}}}}}}}}}}}}}}}:
182      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}}
183    case _:
184      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}}
185
186def aes.state_in(b: M.CmBlock) -> T.State:
187  aes.state16(M.cm.bytes_of(b))
188
189def aes.out(s: T.State) -> M.CmBlock:
190  match s:
191    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}}:
192      M.cm.block16(a0, a1, a2, a3, b0, b1, b2, b3, c0, c1, c2, c3, d0, d1, d2, d3)
193
194# The cipher
195
196# An expanded AES key: bend-collections' schedule (the number of rounds and
197# the round keys), for encryption and decryption alike. Internal; built by
198# CmAes.key.
199type CmAesKey is Data:
200  CmAesKey{sched: A.Schedule}
201
202def aes.key.go(m: Maybe<&2, A.Schedule>, +xs: List<&2, U32>) -> Result<&2, &2, M.CmError, CmAesKey>:
203  match m:
204    case None{}:
205      Fail{M.CmBadKeyLength{B.len(xs)}}
206    case Some{s}:
207      Done{CmAesKey{s}}
208
209def aes.key.pick(ok: Bool, +xs: List<&2, U32>) -> Result<&2, &2, M.CmError, CmAesKey>:
210  match ok:
211    case True{}:
212      aes.key.go(A.expand_key(xs), xs)
213    case False{}:
214      Fail{M.CmBadKeyLength{B.len(xs)}}
215
216# a 16-byte (AES-128) or 32-byte (AES-256) key
217def CmAes.key(o: B.Octets) -> Result<&2, &2, M.CmError, CmAesKey>:
218  +xs = B.Octets.to_list(o)
219  +n = B.len(xs)
220  aes.key.pick(Bool.or(Nat.is_eq(n, 16n), Nat.is_eq(n, 32n)), xs)
221
222def CmAes.key128(o: B.Octets) -> Result<&2, &2, M.CmError, CmAesKey>:
223  +xs = B.Octets.to_list(o)
224  aes.key.pick(Nat.is_eq(B.len(xs), 16n), xs)
225
226def CmAes.key192(o: B.Octets) -> Result<&2, &2, M.CmError, CmAesKey>:
227  +xs = B.Octets.to_list(o)
228  aes.key.pick(Nat.is_eq(B.len(xs), 24n), xs)
229
230def CmAes.key256(o: B.Octets) -> Result<&2, &2, M.CmError, CmAesKey>:
231  +xs = B.Octets.to_list(o)
232  aes.key.pick(Nat.is_eq(B.len(xs), 32n), xs)
233
234def aes.of_maybe(m: Maybe<&2, List<&2, U32>>) -> M.CmBlock:
235  match m:
236    case None{}:
237      M.CmBlock.zero()
238    case Some{xs}:
239      M.cm.block_of(xs)
240
241# one block, encrypted (bend-collections' encrypt_block)
242def CmAes.encrypt(k: CmAesKey, b: M.CmBlock) -> M.CmBlock:
243  match k:
244    case CmAesKey{sched}:
245      aes.of_maybe(A.encrypt_block(aes.sched(sched), M.cm.bytes_of(b)))
246
247# the key as it is: decryption needs nothing beyond the schedule (kept so
248# that callers written for an earlier version, which computed decryption
249# round keys here, still build)
250def CmAes.prepare(k: CmAesKey) -> CmAesKey:
251  k
252
253def aes.decrypt_with(s: A.Schedule, b: M.CmBlock) -> M.CmBlock:
254  match s:
255    case A.Schedule{nr, ks}:
256      aes.out(aes.inv_cipher(nr, ks, aes.state_in(b)))
257
258# one block, decrypted (the inverse cipher above)
259def CmAes.decrypt(k: CmAesKey, b: M.CmBlock) -> M.CmBlock:
260  match k:
261    case CmAesKey{sched}:
262      aes.decrypt_with(aes.sched(sched), b)