x.crypto.chacha20poly1305: enable support for (x)chacha20 with 64-bit counter construct (#25421)

Author: blackshirt

vlib/x/crypto/chacha20poly1305/chacha20poly1305.v

@@ -5,7 +5,7 @@
 // AEAD_CHACHA20_POLY1305 is an authenticated encryption with additional data algorithm.
 // The inputs to AEAD_CHACHA20_POLY1305 are:
 //   A 256-bit key
-//   A 96-bit nonce (or bigger 192 bit nonce) -- different for each invocation with the same key
+//   A 64-bit nonce, 96-bit (or bigger 192 bit nonce) -- different for each invocation with the same key
 //   An arbitrary length plaintext
 //   Arbitrary length additional authenticated data (AAD)
 module chacha20poly1305
@@ -41,55 +41,61 @@ pub interface AEAD {
 
 // key_size is the size of key (in bytes) which the Chacha20Poly1305 AEAD accepts.
 pub const key_size = 32
+
+// orig_nonce_size is the size (in bytes) of nonce of the original (DJ Bernstein) variant
+// which the Chacha20Poly1305 AEAD accepts.
+pub const orig_nonce_size = 8
 // nonce_size is the size of the standard nonce (in bytes) which the Chacha20Poly1305 AEAD accepts.
 pub const nonce_size = 12
 // nonce_size is the size of the extended nonce (in bytes) which the Chacha20Poly1305 AEAD accepts.
 pub const x_nonce_size = 24
+
 // tag_size is the size of the message authenticated code (in bytes) produced by Chacha20Poly1305 AEAD.
 pub const tag_size = 16
 
 // encrypt does one-shot encryption of given plaintext with associated key, nonce and additional data.
 // It return ciphertext output and authenticated tag appended into it.
-pub fn encrypt(plaintext []u8, key []u8, nonce []u8, ad []u8) ![]u8 {
-	mut c := new(key, nonce.len)!
-	out := c.encrypt(plaintext, nonce, ad)!
-	return out
+pub fn encrypt(plaintext []u8, key []u8, nonce []u8, ad []u8, opt chacha20.Options) ![]u8 {
+	mut c := new(key, nonce.len, opt)!
+	return c.encrypt(plaintext, nonce, ad)!
 }
 
 // decrypt does one-shot decryption of given ciphertext with associated key, nonce and additional data.
 // It return plaintext output and verify if resulting tag is a valid message authenticated code (mac)
 // for given message, key and additional data.
-pub fn decrypt(ciphertext []u8, key []u8, nonce []u8, ad []u8) ![]u8 {
-	mut c := new(key, nonce.len)!
-	out := c.decrypt(ciphertext, nonce, ad)!
-	return out
+pub fn decrypt(ciphertext []u8, key []u8, nonce []u8, ad []u8, opt chacha20.Options) ![]u8 {
+	mut c := new(key, nonce.len, opt)!
+	return c.decrypt(ciphertext, nonce, ad)!
 }
 
 // Chacha20Poly1305 represents AEAD algorithm backed by `x.crypto.chacha20` and `x.crypto.poly1305`.
+@[noinit]
 struct Chacha20Poly1305 {
-	key    []u8 = []u8{len: key_size}
-	ncsize int  = nonce_size
+mut:
+	key   []u8 = []u8{len: key_size}
+	nsize int
+	opt   chacha20.Options // for XChaCha20 construct
 }
 
 // new creates a new Chacha20Poly1305 AEAD instance with given 32 bytes of key
-// and the nonce size in ncsize. The ncsize should be 12 or 24 length, otherwise it would return error.
-pub fn new(key []u8, ncsize int) !&AEAD {
+// and the nonce size in nsize. The nsize should be 8, 12 or 24 length, otherwise it would return error.
+pub fn new(key []u8, nsize int, opt chacha20.Options) !&AEAD {
 	if key.len != key_size {
 		return error('chacha20poly1305: bad key size')
 	}
-	if ncsize != nonce_size && ncsize != x_nonce_size {
-		return error('chacha20poly1305: bad nonce size supplied, its should 12 or 24')
+	if nsize != orig_nonce_size && nsize != nonce_size && nsize != x_nonce_size {
+		return error('chacha20poly1305: bad nonce size supplied, its should 8, 12 or 24')
 	}
-	c := &Chacha20Poly1305{
-		key:    key
-		ncsize: ncsize
+	return &Chacha20Poly1305{
+		key:   key
+		nsize: nsize
+		opt:   opt
 	}
-	return c
 }
 
 // nonce_size returns the size of underlying nonce (in bytes) of AEAD algorithm.
 pub fn (c Chacha20Poly1305) nonce_size() int {
-	return c.ncsize
+	return c.nsize
 }
 
 // overhead returns maximum difference between the lengths of a plaintext to be encrypted and
@@ -116,42 +122,7 @@ pub fn (c Chacha20Poly1305) encrypt(plaintext []u8, nonce []u8, ad []u8) ![]u8 {
 	if ad.len > max_u64 {
 		return error('chacha20poly1305: something bad in your additional data')
 	}
-	return c.encrypt_generic(plaintext, nonce, ad)
-}
-
-// encrypt_generic encrypts plaintext along with nonce and additional data
-fn (c Chacha20Poly1305) encrypt_generic(plaintext []u8, nonce []u8, ad []u8) ![]u8 {
-	// First, generates a Poly1305 one-time key from the 256-bit key
-	// and given nonce. Actually its generates by performing ChaCha20 key stream function,
-	// and take the first 32 bytes as a one-time key for Poly1305 from 64 bytes results.
-	// see https://datatracker.ietf.org/doc/html/rfc8439#section-2.6
-	mut polykey := []u8{len: key_size}
-	mut s := chacha20.new_cipher(c.key, nonce)!
-	s.encrypt(mut polykey, polykey)!
-
-	// Next, the ChaCha20 encryption function is called to encrypt the plaintext,
-	// using the same key and nonce, and with the initial ChaCha20 counter set to 1.
-	mut ciphertext := []u8{len: plaintext.len}
-	s.set_counter(1)
-	s.encrypt(mut ciphertext, plaintext)!
-
-	// Finally, the Poly1305 function is called with the generated Poly1305 one-time key
-	// calculated above, and a message constructed as described in
-	// https://datatracker.ietf.org/doc/html/rfc8439#section-2.8
-	mut constructed_msg := []u8{}
-	poly1305_construct_msg(mut constructed_msg, ad, ciphertext)
-
-	// Lets creates Poly1305 instance with one-time key generates in above step,
-	// updates Poly1305 state with this constructed_msg and finally generates tag.
-	mut tag := []u8{len: tag_size}
-	mut po := poly1305.new(polykey)!
-	po.update(constructed_msg)
-	po.finish(mut tag)
-
-	// add this tag to ciphertext output
-	ciphertext << tag
-
-	return ciphertext
+	return c.generic_crypt(plaintext, nonce, ad, .encrypt)!
 }
 
 // decrypt decrypts ciphertext along with provided nonce and additional data.
@@ -162,6 +133,10 @@ fn (c Chacha20Poly1305) encrypt_generic(plaintext []u8, nonce []u8, ad []u8) ![]
 // The calculated mac is bitwise compared to the received mac.
 // The message is authenticated if and only if the tags match, return error if failed to verify.
 pub fn (c Chacha20Poly1305) decrypt(ciphertext []u8, nonce []u8, ad []u8) ![]u8 {
+	// Preliminary check
+	if ciphertext.len < tag_size {
+		return error('chacha20poly1305: ciphertext length does not meet minimum required length')
+	}
 	if nonce.len != c.nonce_size() {
 		return error('chacha20poly1305: unmatching nonce size')
 	}
@@ -170,74 +145,128 @@ pub fn (c Chacha20Poly1305) decrypt(ciphertext []u8, nonce []u8, ad []u8) ![]u8
 	if u64(ciphertext.len) > (u64(1) << 38) - 48 {
 		return error('chacha20poly1305: ciphertext too large')
 	}
-	return c.decrypt_generic(ciphertext, nonce, ad)
+	return c.generic_crypt(ciphertext, nonce, ad, .decrypt)!
 }
 
-fn (c Chacha20Poly1305) decrypt_generic(ciphertext []u8, nonce []u8, ad []u8) ![]u8 {
-	// generates poly1305 one-time key for later calculation
-	mut polykey := []u8{len: key_size}
-	mut s := chacha20.new_cipher(c.key, nonce)!
-	s.encrypt(mut polykey, polykey)!
+// Helpers
 
-	// Remember, ciphertext is concatenation of associated cipher output plus tag (mac) bytes
-	encrypted := ciphertext[0..ciphertext.len - c.overhead()]
-	mac := ciphertext[ciphertext.len - c.overhead()..]
+// generic_crypt direction
+enum Mode {
+	encrypt = 0
+	decrypt = 1
+}
 
-	mut plaintext := []u8{len: encrypted.len}
-	s.set_counter(1)
-	// doing reverse encrypt on cipher output part produces plaintext
-	s.encrypt(mut plaintext, encrypted)!
+// generic_crypt does generic encryption or decryption based on the mode flag was passed.
+// See AEAD Construction at https://datatracker.ietf.org/doc/html/rfc8439#section-2.8
+@[direct_array_access; inline]
+fn (c Chacha20Poly1305) generic_crypt(msg []u8, nonce []u8, ad []u8, mode Mode) ![]u8 {
+	// Setup some values
+	mut src := unsafe { msg }
+	mut mac := []u8{} // used in decryption
+	if mode == .decrypt {
+		src = unsafe { msg[0..msg.len - c.overhead()] }
+		mac = unsafe { msg[msg.len - c.overhead()..] }
+	}
+
+	// generates 32-bytes of one-time key for later poly1305 operation
+	mut otkey := []u8{len: key_size}
+	mut s := chacha20.new_cipher(c.key, nonce, c.opt)!
+	s.encrypt(mut otkey, otkey)!
+
+	// destination buffer, with overhead spaces for generated tag without reallocating
+	mut dst := []u8{len: src.len, cap: src.len + c.overhead()}
 
-	// authenticated messages part
-	mut constructed_msg := []u8{}
-	poly1305_construct_msg(mut constructed_msg, ad, encrypted)
+	// Next, the ChaCha20 encryption function is called to encrypt (decrypt) message input,
+	// using the same key and nonce, and with the initial counter set to 1.
+	s.set_counter(1)
+	s.encrypt(mut dst, src)!
 
+	// Finally, the Poly1305 function is called with the Poly1305 key calculated above
+	// to build message authentication code (tag).
+	// length of constructed message
+	cm_length := if mode == .encrypt {
+		length_constructed_msg(ad, dst)
+	} else {
+		length_constructed_msg(ad, src)
+	}
+	mut constructed_msg := []u8{len: cm_length}
+	if mode == .encrypt {
+		construct_msg(mut constructed_msg, ad, dst)
+	} else {
+		construct_msg(mut constructed_msg, ad, src)
+	}
 	mut tag := []u8{len: tag_size}
-	mut po := poly1305.new(polykey)!
+	mut po := poly1305.new(otkey)!
 	po.update(constructed_msg)
 	po.finish(mut tag)
 
-	// lets verify if received mac is matching with calculated tag,
-	// return error on fail and free allocated resource.
-	if subtle.constant_time_compare(mac, tag) != 1 {
-		// free allocated resource
-		unsafe {
-			s.free()
-			tag.free()
-			polykey.free()
-			plaintext.free()
+	// If this a decryption mode, lets verify whether this calculated tag was matching
+	// with the supplied mac, otherwise return error on fails and free allocated resources.
+	if mode == .decrypt {
+		if subtle.constant_time_compare(mac, tag) != 1 {
+			// free allocated resource
+			unsafe {
+				s.free()
+				tag.free()
+				otkey.free()
+				dst.free()
+				constructed_msg.free()
+			}
+			return error('chacha20poly1305: unmatching tag')
+		} else {
+			// return the decrypted message (plaintext) when the tag was matching
+			return dst
 		}
-		return error('chacha20poly1305: unmatching tag')
 	}
-
-	return plaintext
+	// In the encryption mode, appends the tag into end of destination buffer
+	dst << tag
+	return dst
 }
 
-// Helper function
-
 // pad x to 16 bytes block
+@[direct_array_access; inline]
 fn pad_to_16(x []u8) []u8 {
-	mut buf := x.clone()
-	if buf.len % 16 == 0 {
-		return buf
+	if x.len % 16 == 0 {
+		return x
 	}
-	pad_bytes := []u8{len: 16 - buf.len % 16, init: 0}
-	buf << pad_bytes
-	return buf
+	mut out := []u8{len: x.len + (16 - x.len % 16)}
+	_ := copy(mut out, x)
+	return out
+}
+
+// The length of padded x
+@[inline]
+fn length_pad_to_16(x []u8) int {
+	if x.len % 16 == 0 {
+		return x.len
+	}
+	//
+	return x.len + (16 - x.len % 16)
+}
+
+// The length of constructed message
+@[inline]
+fn length_constructed_msg(ad []u8, bytes []u8) int {
+	mut n := 0
+	n += length_pad_to_16(ad)
+	n += length_pad_to_16(bytes)
+	n += 16 // 2 * 8 bytes
+	return n
 }
 
-// poly1305_construct_msg constructs poly1305 message for later usage and stores to out.
+// construct_msg builds a message for later usage and stored into out.
+// The last step on the AEAD Construction on the how the message was constructed.
+// See the details on the [2.8](https://datatracker.ietf.org/doc/html/rfc8439#section-2.8)
 // The message constructed as a concatenation of the following:
 // 	*  padded to multiple of 16 bytes block of the additional data bytes
 // 	*  padded to multiple of 16 bytes block of the ciphertext (or plaintext) bytes
 // 	*  The length of the additional data in octets (as a 64-bit little-endian integer).
 // 	*  The length of the ciphertext (or plaintext) in octets (as a 64-bit little-endian integer).
-fn poly1305_construct_msg(mut out []u8, ad []u8, bytes []u8) {
-	mut b8 := []u8{len: 8}
-	out << pad_to_16(ad)
-	out << pad_to_16(bytes)
-	binary.little_endian_put_u64(mut b8, u64(ad.len))
-	out << b8
-	binary.little_endian_put_u64(mut b8, u64(bytes.len))
-	out << b8
+// Assumed the output buffer length was correctly initialized.
+@[direct_array_access; inline]
+fn construct_msg(mut out []u8, ad []u8, bytes []u8) {
+	n0 := copy(mut out, pad_to_16(ad))
+	n1 := copy(mut out[n0..], pad_to_16(bytes))
+	binary.little_endian_put_u64(mut out[n0 + n1..], u64(ad.len))
+	binary.little_endian_put_u64(mut out[n0 + n1 + 8..], u64(bytes.len))
 }