1	import Base
2	import ./chacha.bend as CH
3	import ./poly1305.bend as PS
4	import ./subtle.bend as EQ
5	
6	# Specification of the ChaCha20-Poly1305 AEAD (RFC 8439 sections 2.6 and
7	# 2.8) and of XChaCha20-Poly1305 (draft-irtf-cfrg-xchacha-03 section 2),
8	# composed from the ChaCha20 and Poly1305 specifications (as HACL*'s
9	# Spec.Chacha20Poly1305). The sealed message is the ciphertext followed by
10	# the 16-byte tag (the layout of RFC 5116 and of Go's and Python's AEAD
11	# APIs). Opening returns None when the input is shorter than a tag or the
12	# tag is not the one computed over the ciphertext (compared as lists).
13	
14	# 2.6 poly1305_key_gen: the first 32 bytes of the block with counter 0.
15	def poly_key(+key: List<&2, U32>, +nonce: List<&2, U32>) -> List<&2, U32>:
16	  CH.prefix(32n, CH.block(key, 0, nonce))
17	
18	# pad16(x): zero bytes up to a multiple of 16, i.e. (16 - len(x) % 16) % 16
19	# of them ("if (len(x) % 16) == 0 then NULL else copies(0, 16 - len % 16)").
20	def pad16(+xs: List<&2, U32>) -> List<&2, U32>:
21	  List.replicate(U32, Nat.mod(Nat.sub(16n, Nat.mod(List.length(&2, U32, xs), 16n)), 16n), 0)
22	
23	# num_to_8_le_bytes
24	def le8(n: Nat, +x: Nat) -> List<&2, U32>:
25	  match n:
26	    case 0n: Nil{}
27	    case 1n+k: U32.from_nat(Nat.mod(x, 256n)) <> le8(k, Nat.div(x, 256n))
28	
29	# mac_data = aad | pad16(aad) | ciphertext | pad16(ciphertext)
30	#            | num_to_8_le_bytes(aad.length) | num_to_8_le_bytes(ciphertext.length)
31	def mac_data(+aad: List<&2, U32>, +ct: List<&2, U32>) -> List<&2, U32>:
32	  List.append(&2, U32, aad, List.append(&2, U32, pad16(aad), List.append(&2, U32, ct,
33	    List.append(&2, U32, pad16(ct), List.append(&2, U32, le8(8n, List.length(&2, U32, aad)), le8(8n, List.length(&2, U32, ct)))))))
34	
35	def tag(+key: List<&2, U32>, +nonce: List<&2, U32>, +aad: List<&2, U32>, +ct: List<&2, U32>) -> List<&2, U32>:
36	  PS.mac(poly_key(key, nonce), mac_data(aad, ct))
37	
38	# 2.8 chacha20_aead_encrypt: ciphertext = chacha20_encrypt(key, 1, nonce,
39	# plaintext), then the tag over aad and ciphertext.
40	def seal(+key: List<&2, U32>, +nonce: List<&2, U32>, +aad: List<&2, U32>, +pt: List<&2, U32>) -> List<&2, U32>:
41	  +ct = CH.encrypt(key, 1, nonce, pt)
42	  List.append(&2, U32, ct, tag(key, nonce, aad, ct))
43	
44	# Decryption: the tag is recomputed over the received ciphertext and the
45	# plaintext released only when it matches.
46	def open_tag(ok: Bool, +key: List<&2, U32>, +nonce: List<&2, U32>, +ct: List<&2, U32>) -> Maybe<&2, List<&2, U32>>:
47	  match ok:
48	    case True{}: Some{CH.encrypt(key, 1, nonce, ct)}
49	    case False{}: None{}
50	
51	def open_split(+key: List<&2, U32>, +nonce: List<&2, U32>, +aad: List<&2, U32>, +ct: List<&2, U32>, +t: List<&2, U32>) -> Maybe<&2, List<&2, U32>>:
52	  open_tag(EQ.equal(t, tag(key, nonce, aad, ct)), key, nonce, ct)
53	
54	def open_len(short: Bool, +key: List<&2, U32>, +nonce: List<&2, U32>, +aad: List<&2, U32>, +data: List<&2, U32>) -> Maybe<&2, List<&2, U32>>:
55	  match short:
56	    case True{}: None{}
57	    case False{}:
58	      +n = Nat.sub(List.length(&2, U32, data), 16n)
59	      open_split(key, nonce, aad, CH.prefix(n, data), CH.suffix(n, data))
60	
61	def open(+key: List<&2, U32>, +nonce: List<&2, U32>, +aad: List<&2, U32>, +data: List<&2, U32>) -> Maybe<&2, List<&2, U32>>:
62	  open_len(Nat.is_lt(List.length(&2, U32, data), 16n), key, nonce, aad, data)
63	
64	# XChaCha20-Poly1305: ChaCha20-Poly1305 under the HChaCha20 subkey of the
65	# first 16 nonce bytes, with the nonce 0x00000000 || nonce[16..24].
66	def xseal(+key: List<&2, U32>, +nonce: List<&2, U32>, +aad: List<&2, U32>, +pt: List<&2, U32>) -> List<&2, U32>:
67	  seal(CH.xsubkey(key, nonce), CH.xnonce(nonce), aad, pt)
68	
69	def xopen(+key: List<&2, U32>, +nonce: List<&2, U32>, +aad: List<&2, U32>, +data: List<&2, U32>) -> Maybe<&2, List<&2, U32>>:
70	  open(CH.xsubkey(key, nonce), CH.xnonce(nonce), aad, data)