1	import Base
2	import ./core.bend as C
3	
4	# ChaCha20 (RFC 8439 sections 2.3-2.4), HChaCha20 and XChaCha20
5	# (draft-irtf-cfrg-xchacha-03 sections 2.2-2.3) over byte lists.
6	#
7	# Bytes are U32 values (each < 256 in the byte convention); a key or nonce
8	# word is read from the low 8 bits of four bytes, and the keystream is XORed
9	# onto the whole U32 (so encrypting twice is the identity on every list).
10	# Keys are 32 bytes, ChaCha20 nonces 12, HChaCha20 nonces 16 and XChaCha20
11	# nonces 24; the checked functions return None for any other length. The
12	# *_rounds functions take the number of double rounds (ChaCha20: 10).
13	#
14	# Constant time: no branch or memory access depends on key, nonce, counter
15	# or data bytes; only the (public) lengths steer the loops. Bend has no
16	# timing model, so this is a property of the code's shape, not a proved fact.
17	
18	type Key is Data:
19	  K{k0: U32, k1: U32, k2: U32, k3: U32, k4: U32, k5: U32, k6: U32, k7: U32}
20	
21	type Nonce is Data:
22	  N{n0: U32, n1: U32, n2: U32}
23	
24	# ---------------------------------------------------------------- bytes to words
25	
26	def word(b0: U32, b1: U32, b2: U32, b3: U32) -> U32:
27	  U32.or(U32.or(U32.or(U32.and(b0, 255), U32.shln(U32.and(b1, 255), 8n)),
28	    U32.shln(U32.and(b2, 255), 16n)), U32.shln(U32.and(b3, 255), 24n))
29	
30	def words(bytes: List<&2, U32>) -> List<&2, U32>:
31	  match bytes:
32	    case b0 <> b1 <> b2 <> b3 <> rest: word(b0, b1, b2, b3) <> words(rest)
33	    case _: Nil{}
34	
35	def at(xs: List<&2, U32>, i: Nat) -> U32:
36	  match xs i:
37	    case Nil{} _: 0
38	    case x <> rest 0n: x
39	    case x <> rest 1n+p: at(rest, p)
40	
41	def key_words(+ws: List<&2, U32>) -> Key:
42	  K{at(ws, 0n), at(ws, 1n), at(ws, 2n), at(ws, 3n), at(ws, 4n), at(ws, 5n), at(ws, 6n), at(ws, 7n)}
43	
44	def nonce_words(+ws: List<&2, U32>) -> Nonce:
45	  N{at(ws, 0n), at(ws, 1n), at(ws, 2n)}
46	
47	def key(bytes: List<&2, U32>) -> Key:
48	  key_words(words(bytes))
49	
50	def nonce(bytes: List<&2, U32>) -> Nonce:
51	  nonce_words(words(bytes))
52	
53	# ---------------------------------------------------------------- block
54	
55	def state(k: Key, counter: U32, n: Nonce) -> C.State:
56	  match k n:
57	    case K{k0, k1, k2, k3, k4, k5, k6, k7} N{n0, n1, n2}:
58	      C.init(k0, k1, k2, k3, k4, k5, k6, k7, counter, n0, n1, n2)
59	
60	# The 64 keystream bytes of block `counter`, with r double rounds.
61	def keystream(+r: Nat, k: Key, counter: U32, n: Nonce) -> List<&2, U32>:
62	  C.bytes(C.block(r, state(k, counter, n)))
63	
64	# ---------------------------------------------------------------- stream
65	
66	# XOR the keystream onto the front of xs, as far as both go.
67	def xor_front(ks: List<&2, U32>, xs: List<&2, U32>) -> List<&2, U32>:
68	  match ks xs:
69	    case k <> kt x <> xt: U32.xor(x, k) <> xor_front(kt, xt)
70	    case _ _: Nil{}
71	
72	def skip(n: Nat, xs: List<&2, U32>) -> List<&2, U32>:
73	  match n xs:
74	    case 0n _: xs
75	    case 1n+p Nil{}: Nil{}
76	    case 1n+p x <> rest: skip(p, rest)
77	
78	# Blocks of 64 bytes, block j under counter + j; fuel bounds the number of
79	# blocks (the length of xs is enough).
80	def stream(fuel: Nat, +xs: List<&2, U32>, +r: Nat, +k: Key, +counter: U32, +n: Nonce) -> List<&2, U32>:
81	  match fuel xs:
82	    case 0n _: Nil{}
83	    case 1n+p Nil{}: Nil{}
84	    case 1n+p x <> xt:
85	      List.append(&2, U32, xor_front(keystream(r, k, counter, n), x <> xt),
86	        stream(p, skip(64n, x <> xt), r, k, U32.add(counter, 1), n))
87	
88	def encrypt_rounds(+r: Nat, key_bytes: List<&2, U32>, counter: U32, nonce_bytes: List<&2, U32>, +plaintext: List<&2, U32>) -> List<&2, U32>:
89	  stream(List.length(&2, U32, plaintext), plaintext, r, key(key_bytes), counter, nonce(nonce_bytes))
90	
91	def block_body(+r: Nat, key_bytes: List<&2, U32>, counter: U32, nonce_bytes: List<&2, U32>) -> List<&2, U32>:
92	  keystream(r, key(key_bytes), counter, nonce(nonce_bytes))
93	
94	# The block function entered through a match on the key's first cell: both
95	# arms are block_body; with a symbolic key the proof checker leaves the call
96	# unevaluated instead of unrolling the rounds when it compares types.
97	def block_rounds(+r: Nat, key_bytes: List<&2, U32>, counter: U32, nonce_bytes: List<&2, U32>) -> List<&2, U32>:
98	  match key_bytes:
99	    case Nil{}: block_body(r, Nil{}, counter, nonce_bytes)
100	    case b <> rest: block_body(r, b <> rest, counter, nonce_bytes)
101	
102	# ---------------------------------------------------------------- HChaCha20 / XChaCha20
103	
104	def hstate(k: Key, +ws: List<&2, U32>) -> C.State:
105	  match k:
106	    case K{k0, k1, k2, k3, k4, k5, k6, k7}:
107	      C.hinit(k0, k1, k2, k3, k4, k5, k6, k7, at(ws, 0n), at(ws, 1n), at(ws, 2n), at(ws, 3n))
108	
109	def hchacha_body(+r: Nat, key_bytes: List<&2, U32>, nonce_bytes: List<&2, U32>) -> List<&2, U32>:
110	  C.hbytes(C.rounds(r, hstate(key(key_bytes), words(nonce_bytes))))
111	
112	def hchacha_rounds(+r: Nat, key_bytes: List<&2, U32>, nonce_bytes: List<&2, U32>) -> List<&2, U32>:
113	  match key_bytes:
114	    case Nil{}: hchacha_body(r, Nil{}, nonce_bytes)
115	    case b <> rest: hchacha_body(r, b <> rest, nonce_bytes)
116	
117	def hchacha20_unchecked(key_bytes: List<&2, U32>, nonce_bytes: List<&2, U32>) -> List<&2, U32>:
118	  hchacha_rounds(10n, key_bytes, nonce_bytes)
119	
120	def take(n: Nat, xs: List<&2, U32>) -> List<&2, U32>:
121	  match n xs:
122	    case 0n _: Nil{}
123	    case 1n+p Nil{}: Nil{}
124	    case 1n+p x <> rest: x <> take(p, rest)
125	
126	# XChaCha20: ChaCha20 under the HChaCha20 subkey of nonce[0..16], with the
127	# nonce 0x00000000 || nonce[16..24].
128	def xencrypt_unchecked(key_bytes: List<&2, U32>, counter: U32, +nonce_bytes: List<&2, U32>, +plaintext: List<&2, U32>) -> List<&2, U32>:
129	  +sub = hchacha20_unchecked(key_bytes, take(16n, nonce_bytes))
130	  encrypt_rounds(10n, sub, counter, 0 <> 0 <> 0 <> 0 <> skip(16n, nonce_bytes), plaintext)
131	
132	# ---------------------------------------------------------------- checked API
133	
134	def has_length(+xs: List<&2, U32>, n: Nat) -> Bool:
135	  Nat.is_eq(List.length(&2, U32, xs), n)
136	
137	def valid(+key_bytes: List<&2, U32>, +nonce_bytes: List<&2, U32>, n: Nat) -> Bool:
138	  Bool.and(has_length(key_bytes, 32n), has_length(nonce_bytes, n))
139	
140	def when(ok: Bool, x: List<&2, U32>) -> Maybe<&2, List<&2, U32>>:
141	  match ok:
142	    case True{}: Some{x}
143	    case False{}: None{}
144	
145	# chacha20_block(key, counter, nonce): 64 bytes, or None unless the key has 32
146	# bytes and the nonce 12.
147	def chacha20_block(+key_bytes: List<&2, U32>, counter: U32, +nonce_bytes: List<&2, U32>) -> Maybe<&2, List<&2, U32>>:
148	  when(valid(key_bytes, nonce_bytes, 12n), block_rounds(10n, key_bytes, counter, nonce_bytes))
149	
150	# chacha20_encrypt(key, counter, nonce, plaintext) (encryption and decryption
151	# are the same operation).
152	def chacha20(+key_bytes: List<&2, U32>, counter: U32, +nonce_bytes: List<&2, U32>, +plaintext: List<&2, U32>) -> Maybe<&2, List<&2, U32>>:
153	  when(valid(key_bytes, nonce_bytes, 12n), encrypt_rounds(10n, key_bytes, counter, nonce_bytes, plaintext))
154	
155	# HChaCha20(key, nonce): the 32-byte subkey, or None unless 32 and 16 bytes.
156	def hchacha20(+key_bytes: List<&2, U32>, +nonce_bytes: List<&2, U32>) -> Maybe<&2, List<&2, U32>>:
157	  when(valid(key_bytes, nonce_bytes, 16n), hchacha20_unchecked(key_bytes, nonce_bytes))
158	
159	# XChaCha20 with a 24-byte nonce.
160	def xchacha20(+key_bytes: List<&2, U32>, counter: U32, +nonce_bytes: List<&2, U32>, +plaintext: List<&2, U32>) -> Maybe<&2, List<&2, U32>>:
161	  when(valid(key_bytes, nonce_bytes, 24n), xencrypt_unchecked(key_bytes, counter, nonce_bytes, plaintext))