aesround.bend

bend-osdev-aead@0.1.0.0 · 173 lines · 26 definitions

1# The AES encryption round function, AESRound(A, B) = MixColumns(ShiftRows(
2# SubBytes(A))) xor B (FIPS 197, sections 5.1.1 to 5.1.4; not the last
3# round), on 16-byte blocks: the building block of AEGIS (aegis128.bend). It
4# is the x86 instruction AESENC; RFC 10032 (the AEGIS family), Appendix A.1,
5# gives its test vector.
6#
7# A block is four big-endian 32-bit words, each one column of the AES state:
8# byte r of word c is the state's row r, column c (FIPS 197, 3.4: byte i of
9# the input is row i mod 4, column i / 4). The S-box is a table built once
10# (AdAes.sbox) and passed in, so a caller that runs many rounds (AEGIS keeps
11# it in its key) builds it once.
12import Base
13import bend-osdev-bytes@0.1.0.0/bytes.bend as B
14
15# A 16-byte block (open type)
16type AdBlk is Data:
17  AdBlk{w0: U32, w1: U32, w2: U32, w3: U32}
18
19def AdBlk.zero() -> AdBlk:
20  AdBlk{0, 0, 0, 0}
21
22def AdBlk.xor(a: AdBlk, b: AdBlk) -> AdBlk:
23  match a b:
24    case AdBlk{a0, a1, a2, a3} AdBlk{b0, b1, b2, b3}:
25      AdBlk{U32.xor(a0, b0), U32.xor(a1, b1), U32.xor(a2, b2), U32.xor(a3, b3)}
26
27def AdBlk.and(a: AdBlk, b: AdBlk) -> AdBlk:
28  match a b:
29    case AdBlk{a0, a1, a2, a3} AdBlk{b0, b1, b2, b3}:
30      AdBlk{U32.and(a0, b0), U32.and(a1, b1), U32.and(a2, b2), U32.and(a3, b3)}
31
32def ad.blk.at(xs: List<&2, U32>, i: Nat) -> U32:
33  match xs:
34    case Nil{}:
35      0
36    case Con{h, t}:
37      match i:
38        case 0n:
39          h
40        case 1n+p:
41          ad.blk.at(t, p)
42
43def ad.blk.word(+xs: List<&2, U32>, +i: Nat) -> U32:
44  U32.or(U32.or(U32.shln(B.byte(ad.blk.at(xs, i)), 24n), U32.shln(B.byte(ad.blk.at(xs, Nat.add(i, 1n))), 16n)), U32.or(U32.shln(B.byte(ad.blk.at(xs, Nat.add(i, 2n))), 8n), B.byte(ad.blk.at(xs, Nat.add(i, 3n)))))
45
46# the block of the first 16 bytes of xs (zeros past the end)
47def AdBlk.of_list(+xs: List<&2, U32>) -> AdBlk:
48  AdBlk{ad.blk.word(xs, 0n), ad.blk.word(xs, 4n), ad.blk.word(xs, 8n), ad.blk.word(xs, 12n)}
49
50# a block's 16 bytes
51def AdBlk.to_list(b: AdBlk) -> List<&2, U32>:
52  match b:
53    case AdBlk{w0, w1, w2, w3}:
54      B.append(B.be32(w0), B.append(B.be32(w1), B.append(B.be32(w2), B.be32(w3))))
55
56def AdBlk.of_octets(+o: B.Octets) -> AdBlk:
57  AdBlk.of_list(B.Octets.to_list(o))
58
59def AdBlk.to_octets(b: AdBlk) -> B.Octets:
60  B.Octets.of_list(AdBlk.to_list(b))
61
62def AdBlk.hex(b: AdBlk) -> String:
63  B.Octets.hex(AdBlk.to_octets(b))
64
65# The S-box
66
67# BEGIN GENERATED by gen/tables.py: sbox
68def ad.aes.sbox.list() -> List<&2, U32>:
69  [99, 124, 119, 123, 242, 107, 111, 197, 48, 1, 103, 43, 254, 215, 171, 118,
70    202, 130, 201, 125, 250, 89, 71, 240, 173, 212, 162, 175, 156, 164, 114, 192,
71    183, 253, 147, 38, 54, 63, 247, 204, 52, 165, 229, 241, 113, 216, 49, 21,
72    4, 199, 35, 195, 24, 150, 5, 154, 7, 18, 128, 226, 235, 39, 178, 117,
73    9, 131, 44, 26, 27, 110, 90, 160, 82, 59, 214, 179, 41, 227, 47, 132,
74    83, 209, 0, 237, 32, 252, 177, 91, 106, 203, 190, 57, 74, 76, 88, 207,
75    208, 239, 170, 251, 67, 77, 51, 133, 69, 249, 2, 127, 80, 60, 159, 168,
76    81, 163, 64, 143, 146, 157, 56, 245, 188, 182, 218, 33, 16, 255, 243, 210,
77    205, 12, 19, 236, 95, 151, 68, 23, 196, 167, 126, 61, 100, 93, 25, 115,
78    96, 129, 79, 220, 34, 42, 144, 136, 70, 238, 184, 20, 222, 94, 11, 219,
79    224, 50, 58, 10, 73, 6, 36, 92, 194, 211, 172, 98, 145, 149, 228, 121,
80    231, 200, 55, 109, 141, 213, 78, 169, 108, 86, 244, 234, 101, 122, 174, 8,
81    186, 120, 37, 46, 28, 166, 180, 198, 232, 221, 116, 31, 75, 189, 139, 138,
82    112, 62, 181, 102, 72, 3, 246, 14, 97, 53, 87, 185, 134, 193, 29, 158,
83    225, 248, 152, 17, 105, 217, 142, 148, 155, 30, 135, 233, 206, 85, 40, 223,
84    140, 161, 137, 13, 191, 230, 66, 104, 65, 153, 45, 15, 176, 84, 187, 22]
85# END GENERATED: sbox
86
87# a table of 2^d words, read by the bits of an index (internal)
88type AdTab is Data:
89  AdLeaf{v: U32}
90  AdNode{l: AdTab, r: AdTab}
91
92def ad.tab.leaves(xs: List<&2, U32>) -> List<&2, AdTab>:
93  match xs:
94    case Nil{}:
95      Nil{}
96    case Con{h, t}:
97      Con{AdLeaf{h}, ad.tab.leaves(t)}
98
99def ad.tab.pairs(ts: List<&2, AdTab>) -> List<&2, AdTab>:
100  match ts:
101    case Con{a, Con{b, t}}:
102      Con{AdNode{a, b}, ad.tab.pairs(t)}
103    case _:
104      Nil{}
105
106def ad.tab.up(n: Nat, ts: List<&2, AdTab>) -> List<&2, AdTab>:
107  match n:
108    case 0n:
109      ts
110    case 1n+p:
111      ad.tab.up(p, ad.tab.pairs(ts))
112
113def ad.tab.root(ts: List<&2, AdTab>) -> AdTab:
114  match ts:
115    case Nil{}:
116      AdLeaf{0}
117    case Con{h, t}:
118      h
119
120def ad.tab.pick(one: Bool, l: AdTab, r: AdTab) -> AdTab:
121  match one:
122    case True{}:
123      r
124    case False{}:
125      l
126
127def ad.tab.get(d: Nat, t: AdTab, +i: U32) -> U32:
128  match d:
129    case 0n:
130      match t:
131        case AdLeaf{v}:
132          v
133        case AdNode{l, r}:
134          0
135    case 1n+(+p):
136      match t:
137        case AdLeaf{v}:
138          v
139        case AdNode{l, r}:
140          ad.tab.get(p, ad.tab.pick(U32.is_eq(U32.and(1, U32.shrn(i, p)), 1), l, r), i)
141
142# The S-box as a table (opaque)
143type AdSbox is Data:
144  AdSbox{t: AdTab}
145
146def AdAes.sbox() -> AdSbox:
147  AdSbox{ad.tab.root(ad.tab.up(8n, ad.tab.leaves(ad.aes.sbox.list())))}
148
149def ad.aes.s(+t: AdTab, +w: U32, n: Nat) -> U32:
150  ad.tab.get(8n, t, U32.and(255, U32.shrn(w, n)))
151
152# x times 2 in GF(2^8) (xtime, FIPS 197 4.2.1)
153def ad.aes.x2(+a: U32) -> U32:
154  U32.xor(U32.and(255, U32.shl(a)), U32.mul(27, U32.shrn(a, 7n)))
155
156# MixColumns of one column a0 .. a3 (row 0 first), as a word
157def ad.aes.mix(+a0: U32, +a1: U32, +a2: U32, +a3: U32) -> U32:
158  +b0 = U32.xor(U32.xor(ad.aes.x2(a0), U32.xor(ad.aes.x2(a1), a1)), U32.xor(a2, a3))
159  +b1 = U32.xor(U32.xor(a0, ad.aes.x2(a1)), U32.xor(U32.xor(ad.aes.x2(a2), a2), a3))
160  +b2 = U32.xor(U32.xor(a0, a1), U32.xor(ad.aes.x2(a2), U32.xor(ad.aes.x2(a3), a3)))
161  +b3 = U32.xor(U32.xor(U32.xor(ad.aes.x2(a0), a0), a1), U32.xor(a2, ad.aes.x2(a3)))
162  U32.or(U32.or(U32.shln(b0, 24n), U32.shln(b1, 16n)), U32.or(U32.shln(b2, 8n), b3))
163
164# one column of the round: rows 0 .. 3 from columns c, c + 1, c + 2, c + 3
165# (ShiftRows), each through the S-box (SubBytes), then MixColumns
166def ad.aes.col(+t: AdTab, +x0: U32, +x1: U32, +x2: U32, +x3: U32) -> U32:
167  ad.aes.mix(ad.aes.s(t, x0, 24n), ad.aes.s(t, x1, 16n), ad.aes.s(t, x2, 8n), ad.aes.s(t, x3, 0n))
168
169# AESRound(a, k)
170def AdAes.round(s: AdSbox, a: AdBlk, k: AdBlk) -> AdBlk:
171  match s a k:
172    case AdSbox{+t} AdBlk{+a0, +a1, +a2, +a3} AdBlk{k0, k1, k2, k3}:
173      AdBlk{U32.xor(ad.aes.col(t, a0, a1, a2, a3), k0), U32.xor(ad.aes.col(t, a1, a2, a3, a0), k1), U32.xor(ad.aes.col(t, a2, a3, a0, a1), k2), U32.xor(ad.aes.col(t, a3, a0, a1, a2), k3)}