// Copyright (c) 2024 blackshirt.
// Use of this source code is governed by an MIT license
// that can be found in the LICENSE file.
//
// AEAD_CHACHA20_POLY1305 is an authenticated encryption with additional data algorithm.
// The inputs to AEAD_CHACHA20_POLY1305 are:
//   A 256-bit 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

import crypto.cipher
import encoding.binary
import crypto.internal.subtle
import x.crypto.chacha20
import x.crypto.poly1305

// 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, 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, 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 {
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 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) !&cipher.AEAD {
	if key.len != key_size {
		return error('chacha20poly1305: bad key size')
	}
	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')
	}
	return &Chacha20Poly1305{
		key:   key
		nsize: nsize
		opt:   opt
	}
}

// nonce_size returns the size of underlying nonce (in bytes) of AEAD algorithm.
pub fn (c Chacha20Poly1305) nonce_size() int {
	return c.nsize
}

// overhead returns maximum difference between the lengths of a plaintext to be encrypted and
// ciphertext's output. In the context of Chacha20Poly1305, `.overhead() == .tag_size`.
pub fn (c Chacha20Poly1305) overhead() int {
	return tag_size
}

// encrypt encrypts plaintext, along with nonce and additional data and generates
// authenticated tag appended into ciphertext's output.
pub fn (c Chacha20Poly1305) encrypt(plaintext []u8, nonce []u8, ad []u8) ![]u8 {
	// makes sure if the nonce length is matching with internal nonce size
	if nonce.len != c.nonce_size() {
		return error('chacha20poly1305: unmatching nonce size')
	}
	// check if the plaintext length doesn't exceed the amount of limit.
	// its comes from the internal of chacha20 mechanism, where the counter are u32
	// with the facts of chacha20 operates on 64 bytes block, we can measure the amount
	// of encrypted data possible in a single invocation, ie.,
	// amount = (2^32-1)*64 = 274,877,906,880 bytes, or nearly 256 GB
	if u64(plaintext.len) > (u64(1) << 38) - 64 {
		panic('chacha20poly1305: plaintext too large')
	}
	if ad.len > max_u64 {
		return error('chacha20poly1305: something bad in your additional data')
	}
	return c.generic_crypt(plaintext, nonce, ad, .encrypt)!
}

// decrypt decrypts ciphertext along with provided nonce and additional data.
// Decryption is similar with the encryption process with slight differences in:
// The roles of ciphertext and plaintext are reversed, so the ChaCha20 encryption
// function is applied to the ciphertext, producing the plaintext.
// The Poly1305 function is still run on the AAD and the ciphertext, not the plaintext.
// 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')
	}
	// ciphertext max = plaintext max length  + tag length
	// ie, (2^32-1)*64 + overhead = (u64(1) << 38) - 64 + 16 = 274,877,906,896 octets.
	if u64(ciphertext.len) > (u64(1) << 38) - 48 {
		return error('chacha20poly1305: ciphertext too large')
	}
	return c.generic_crypt(ciphertext, nonce, ad, .decrypt)!
}

// Helpers

// generic_crypt direction
enum Mode {
	encrypt = 0
	decrypt = 1
}

// 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()}

	// 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(otkey)!
	po.update(constructed_msg)
	po.finish(mut tag)

	// 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
		}
	}
	// In the encryption mode, appends the tag into end of destination buffer
	dst << tag
	return dst
}

// pad x to 16 bytes block
@[direct_array_access; inline]
fn pad_to_16(x []u8) []u8 {
	if x.len % 16 == 0 {
		return x
	}
	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
}

// 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).
// 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))
}