/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at https://mozilla.org/MPL/2.0/.
 *
 * (c) ZeroTier, Inc.
 * https://www.zerotier.com/
 */

#ifndef ZT_AES_HPP
#define ZT_AES_HPP

#include "Constants.hpp"
#include "Utils.hpp"

// Uncomment to disable all hardware acceleration (usually for testing)
// #define ZT_AES_NO_ACCEL

#if ! defined(ZT_AES_NO_ACCEL) && defined(ZT_ARCH_X64)
#define ZT_AES_AESNI 1
#endif
#if ! defined(ZT_AES_NO_ACCEL) && defined(ZT_ARCH_ARM_HAS_NEON) && defined(ZT_ARCH_ARM_HAS_CRYPTO)
#define ZT_AES_NEON 1
#endif

#ifndef ZT_INLINE
#define ZT_INLINE inline
#endif

namespace ZeroTier {

/**
 * AES-256 and pals including GMAC, CTR, etc.
 *
 * This includes hardware acceleration for certain processors. The software
 * mode is fallback and is significantly slower.
 */
class AES {
  public:
	/**
	 * @return True if this system has hardware AES acceleration
	 */
	static ZT_INLINE bool accelerated()
	{
#ifdef ZT_AES_AESNI
		return Utils::CPUID.aes;
#else
#ifdef ZT_AES_NEON
		return Utils::ARMCAP.aes;
#else
		return false;
#endif
#endif
	}

	/**
	 * Create an un-initialized AES instance (must call init() before use)
	 */
	ZT_INLINE AES() noexcept
	{
	}

	/**
	 * Create an AES instance with the given key
	 *
	 * @param key 256-bit key
	 */
	explicit ZT_INLINE AES(const void* const key) noexcept
	{
		this->init(key);
	}

	ZT_INLINE ~AES()
	{
		Utils::burn(&p_k, sizeof(p_k));
	}

	/**
	 * Set (or re-set) this AES256 cipher's key
	 *
	 * @param key 256-bit / 32-byte key
	 */
	ZT_INLINE void init(const void* const key) noexcept
	{
#ifdef ZT_AES_AESNI
		if (likely(Utils::CPUID.aes)) {
			p_init_aesni(reinterpret_cast<const uint8_t*>(key));
			return;
		}
#endif
#ifdef ZT_AES_NEON
		if (Utils::ARMCAP.aes) {
			p_init_armneon_crypto(reinterpret_cast<const uint8_t*>(key));
			return;
		}
#endif
		p_initSW(reinterpret_cast<const uint8_t*>(key));
	}

	/**
	 * Encrypt a single AES block
	 *
	 * @param in Input block
	 * @param out Output block (can be same as input)
	 */
	ZT_INLINE void encrypt(const void* const in, void* const out) const noexcept
	{
#ifdef ZT_AES_AESNI
		if (likely(Utils::CPUID.aes)) {
			p_encrypt_aesni(in, out);
			return;
		}
#endif
#ifdef ZT_AES_NEON
		if (Utils::ARMCAP.aes) {
			p_encrypt_armneon_crypto(in, out);
			return;
		}
#endif
		p_encryptSW(reinterpret_cast<const uint8_t*>(in), reinterpret_cast<uint8_t*>(out));
	}

	/**
	 * Decrypt a single AES block
	 *
	 * @param in Input block
	 * @param out Output block (can be same as input)
	 */
	ZT_INLINE void decrypt(const void* const in, void* const out) const noexcept
	{
#ifdef ZT_AES_AESNI
		if (likely(Utils::CPUID.aes)) {
			p_decrypt_aesni(in, out);
			return;
		}
#endif
#ifdef ZT_AES_NEON
		if (Utils::ARMCAP.aes) {
			p_decrypt_armneon_crypto(in, out);
			return;
		}
#endif
		p_decryptSW(reinterpret_cast<const uint8_t*>(in), reinterpret_cast<uint8_t*>(out));
	}

	class GMACSIVEncryptor;
	class GMACSIVDecryptor;

	/**
	 * Streaming GMAC calculator
	 */
	class GMAC {
		friend class GMACSIVEncryptor;
		friend class GMACSIVDecryptor;

	  public:
		/**
		 * @return True if this system has hardware GMAC acceleration
		 */
		static ZT_INLINE bool accelerated()
		{
#ifdef ZT_AES_AESNI
			return Utils::CPUID.aes;
#else
#ifdef ZT_AES_NEON
			return Utils::ARMCAP.pmull;
#else
			return false;
#endif
#endif
		}

		/**
		 * Create a new instance of GMAC (must be initialized with init() before use)
		 *
		 * @param aes Keyed AES instance to use
		 */
		ZT_INLINE GMAC(const AES& aes) : _aes(aes)
		{
		}

		/**
		 * Reset and initialize for a new GMAC calculation
		 *
		 * @param iv 96-bit initialization vector (pad with zeroes if actual IV is shorter)
		 */
		ZT_INLINE void init(const uint8_t iv[12]) noexcept
		{
			_rp = 0;
			_len = 0;
			// We fill the least significant 32 bits in the _iv field with 1 since in GCM mode
			// this would hold the counter, but we're not doing GCM. The counter is therefore
			// always 1.
#ifdef ZT_AES_AESNI	  // also implies an x64 processor
			*reinterpret_cast<uint64_t*>(_iv) = *reinterpret_cast<const uint64_t*>(iv);
			*reinterpret_cast<uint32_t*>(_iv + 8) = *reinterpret_cast<const uint64_t*>(iv + 8);
			*reinterpret_cast<uint32_t*>(_iv + 12) = 0x01000000;   // 0x00000001 in big-endian byte order
#else
			for (int i = 0; i < 12; ++i) {
				_iv[i] = iv[i];
			}
			_iv[12] = 0;
			_iv[13] = 0;
			_iv[14] = 0;
			_iv[15] = 1;
#endif
			_y[0] = 0;
			_y[1] = 0;
		}

		/**
		 * Process data through GMAC
		 *
		 * @param data Bytes to process
		 * @param len Length of input
		 */
		void update(const void* data, unsigned int len) noexcept;

		/**
		 * Process any remaining cached bytes and generate tag
		 *
		 * Don't call finish() more than once or you'll get an invalid result.
		 *
		 * @param tag 128-bit GMAC tag (can be truncated)
		 */
		void finish(uint8_t tag[16]) noexcept;

	  private:
#ifdef ZT_AES_AESNI
		void p_aesNIUpdate(const uint8_t* in, unsigned int len) noexcept;
		void p_aesNIFinish(uint8_t tag[16]) noexcept;
#endif
#ifdef ZT_AES_NEON
		void p_armUpdate(const uint8_t* in, unsigned int len) noexcept;
		void p_armFinish(uint8_t tag[16]) noexcept;
#endif
		const AES& _aes;
		unsigned int _rp;
		unsigned int _len;
		uint8_t _r[16];	  // remainder
		uint8_t _iv[16];
		uint64_t _y[2];
	};

	/**
	 * Streaming AES-CTR encrypt/decrypt
	 *
	 * NOTE: this doesn't support overflow of the counter in the least significant 32 bits.
	 * AES-GMAC-CTR doesn't need this, so we don't support it as an optimization.
	 */
	class CTR {
		friend class GMACSIVEncryptor;
		friend class GMACSIVDecryptor;

	  public:
		ZT_INLINE CTR(const AES& aes) noexcept : _aes(aes)
		{
		}

		/**
		 * Initialize this CTR instance to encrypt a new stream
		 *
		 * @param iv Unique initialization vector and initial 32-bit counter (least significant 32 bits, big-endian)
		 * @param output Buffer to which to store output (MUST be large enough for total bytes processed!)
		 */
		ZT_INLINE void init(const uint8_t iv[16], void* const output) noexcept
		{
			Utils::copy<16>(_ctr, iv);
			_out = reinterpret_cast<uint8_t*>(output);
			_len = 0;
		}

		/**
		 * Initialize this CTR instance to encrypt a new stream
		 *
		 * @param iv Unique initialization vector
		 * @param ic Initial counter (must be in big-endian byte order!)
		 * @param output Buffer to which to store output (MUST be large enough for total bytes processed!)
		 */
		ZT_INLINE void init(const uint8_t iv[12], const uint32_t ic, void* const output) noexcept
		{
			Utils::copy<12>(_ctr, iv);
			reinterpret_cast<uint32_t*>(_ctr)[3] = ic;
			_out = reinterpret_cast<uint8_t*>(output);
			_len = 0;
		}

		/**
		 * Encrypt or decrypt data, writing result to the output provided to init()
		 *
		 * @param input Input data
		 * @param len Length of input
		 */
		void crypt(const void* input, unsigned int len) noexcept;

		/**
		 * Finish any remaining bytes if total bytes processed wasn't a multiple of 16
		 *
		 * Don't call more than once for a given stream or data may be corrupted.
		 */
		void finish() noexcept;

	  private:
#ifdef ZT_AES_AESNI
		void p_aesNICrypt(const uint8_t* in, uint8_t* out, unsigned int len) noexcept;
#endif
#ifdef ZT_AES_NEON
		void p_armCrypt(const uint8_t* in, uint8_t* out, unsigned int len) noexcept;
#endif
		const AES& _aes;
		uint64_t _ctr[2];
		uint8_t* _out;
		unsigned int _len;
	};

	/**
	 * Encryptor for AES-GMAC-SIV.
	 *
	 * Encryption requires two passes. The first pass starts after init
	 * with aad (if any) followed by update1() and finish1(). Then the
	 * update2() and finish2() methods must be used over the same data
	 * (but NOT AAD) again.
	 *
	 * This supports encryption of a maximum of 2^31 bytes of data per
	 * call to init().
	 */
	class GMACSIVEncryptor {
	  public:
		/**
		 * Create a new AES-GMAC-SIV encryptor keyed with the provided AES instances
		 *
		 * @param k0 First of two AES instances keyed with K0
		 * @param k1 Second of two AES instances keyed with K1
		 */
		ZT_INLINE GMACSIVEncryptor(const AES& k0, const AES& k1) noexcept
			: _gmac(k0)
			, _ctr(k1)
		{
		}

		/**
		 * Initialize AES-GMAC-SIV
		 *
		 * @param iv IV in network byte order (byte order in which it will appear on the wire)
		 * @param output Pointer to buffer to receive ciphertext, must be large enough for all to-be-processed data!
		 */
		ZT_INLINE void init(const uint64_t iv, void* const output) noexcept
		{
			// Output buffer to receive the result of AES-CTR encryption.
			_output = output;

			// Initialize GMAC with 64-bit IV (and remaining 32 bits padded to zero).
			_tag[0] = iv;
			_tag[1] = 0;
			_gmac.init(reinterpret_cast<const uint8_t*>(_tag));
		}

		/**
		 * Process AAD (additional authenticated data) that is not being encrypted.
		 *
		 * If such data exists this must be called before update1() and finish1().
		 *
		 * Note: current code only supports one single chunk of AAD. Don't call this
		 * multiple times per message.
		 *
		 * @param aad Additional authenticated data
		 * @param len Length of AAD in bytes
		 */
		ZT_INLINE void aad(const void* const aad, unsigned int len) noexcept
		{
			// Feed ADD into GMAC first
			_gmac.update(aad, len);

			// End of AAD is padded to a multiple of 16 bytes to ensure unique encoding.
			len &= 0xfU;
			if (len != 0) {
				_gmac.update(Utils::ZERO256, 16 - len);
			}
		}

		/**
		 * First pass plaintext input function
		 *
		 * @param input Plaintext chunk
		 * @param len Length of plaintext chunk
		 */
		ZT_INLINE void update1(const void* const input, const unsigned int len) noexcept
		{
			_gmac.update(input, len);
		}

		/**
		 * Finish first pass, compute CTR IV, initialize second pass.
		 */
		ZT_INLINE void finish1() noexcept
		{
			// Compute 128-bit GMAC tag.
			uint64_t tmp[2];
			_gmac.finish(reinterpret_cast<uint8_t*>(tmp));

			// Shorten to 64 bits, concatenate with message IV, and encrypt with AES to
			// yield the CTR IV and opaque IV/MAC blob. In ZeroTier's use of GMAC-SIV
			// this get split into the packet ID (64 bits) and the MAC (64 bits) in each
			// packet and then recombined on receipt for legacy reasons (but with no
			// cryptographic or performance impact).
			_tag[1] = tmp[0] ^ tmp[1];
			_ctr._aes.encrypt(_tag, _tag);

			// Initialize CTR with 96-bit CTR nonce and 32-bit counter. The counter
			// incorporates 31 more bits of entropy which should raise our security margin
			// a bit, but this is not included in the worst case analysis of GMAC-SIV.
			// The most significant bit of the counter is masked to zero to allow up to
			// 2^31 bytes to be encrypted before the counter loops. Some CTR implementations
			// increment the whole big-endian 128-bit integer in which case this could be
			// used for more than 2^31 bytes, but ours does not for performance reasons
			// and so 2^31 should be considered the input limit.
			tmp[0] = _tag[0];
			tmp[1] = _tag[1] & ZT_CONST_TO_BE_UINT64(0xffffffff7fffffffULL);
			_ctr.init(reinterpret_cast<const uint8_t*>(tmp), _output);
		}

		/**
		 * Second pass plaintext input function
		 *
		 * The same plaintext must be fed in the second time in the same order,
		 * though chunk boundaries do not have to be the same.
		 *
		 * @param input Plaintext chunk
		 * @param len Length of plaintext chunk
		 */
		ZT_INLINE void update2(const void* const input, const unsigned int len) noexcept
		{
			_ctr.crypt(input, len);
		}

		/**
		 * Finish second pass and return a pointer to the opaque 128-bit IV+MAC block
		 *
		 * The returned pointer remains valid as long as this object exists and init()
		 * is not called again.
		 *
		 * @return Pointer to 128-bit opaque IV+MAC (packed into two 64-bit integers)
		 */
		ZT_INLINE const uint64_t* finish2()
		{
			_ctr.finish();
			return _tag;
		}

	  private:
		void* _output;
		uint64_t _tag[2];
		AES::GMAC _gmac;
		AES::CTR _ctr;
	};

	/**
	 * Decryptor for AES-GMAC-SIV.
	 *
	 * GMAC-SIV decryption is single-pass. AAD (if any) must be processed first.
	 */
	class GMACSIVDecryptor {
	  public:
		ZT_INLINE GMACSIVDecryptor(const AES& k0, const AES& k1) noexcept
			: _ctr(k1)
			, _gmac(k0)
		{
		}

		/**
		 * Initialize decryptor for a new message
		 *
		 * @param tag 128-bit combined IV/MAC originally created by GMAC-SIV encryption
		 * @param output Buffer in which to write output plaintext (must be large enough!)
		 */
		ZT_INLINE void init(const uint64_t tag[2], void* const output) noexcept
		{
			uint64_t tmp[2];
			tmp[0] = tag[0];
			tmp[1] = tag[1] & ZT_CONST_TO_BE_UINT64(0xffffffff7fffffffULL);
			_ctr.init(reinterpret_cast<const uint8_t*>(tmp), output);

			_ctr._aes.decrypt(tag, _ivMac);

			tmp[0] = _ivMac[0];
			tmp[1] = 0;
			_gmac.init(reinterpret_cast<const uint8_t*>(tmp));

			_output = output;
			_decryptedLen = 0;
		}

		/**
		 * Process AAD (additional authenticated data) that wasn't encrypted
		 *
		 * @param aad Additional authenticated data
		 * @param len Length of AAD in bytes
		 */
		ZT_INLINE void aad(const void* const aad, unsigned int len) noexcept
		{
			_gmac.update(aad, len);
			len &= 0xfU;
			if (len != 0) {
				_gmac.update(Utils::ZERO256, 16 - len);
			}
		}

		/**
		 * Feed ciphertext into the decryptor
		 *
		 * Unlike encryption, GMAC-SIV decryption requires only one pass.
		 *
		 * @param input Input ciphertext
		 * @param len Length of ciphertext
		 */
		ZT_INLINE void update(const void* const input, const unsigned int len) noexcept
		{
			_ctr.crypt(input, len);
			_decryptedLen += len;
		}

		/**
		 * Flush decryption, compute MAC, and verify
		 *
		 * @return True if resulting plaintext (and AAD) pass message authentication check
		 */
		ZT_INLINE bool finish() noexcept
		{
			_ctr.finish();

			uint64_t gmacTag[2];
			_gmac.update(_output, _decryptedLen);
			_gmac.finish(reinterpret_cast<uint8_t*>(gmacTag));
			return (gmacTag[0] ^ gmacTag[1]) == _ivMac[1];
		}

	  private:
		uint64_t _ivMac[2];
		AES::CTR _ctr;
		AES::GMAC _gmac;
		void* _output;
		unsigned int _decryptedLen;
	};

  private:
	static const uint32_t Te0[256];
	static const uint32_t Te4[256];
	static const uint32_t Td0[256];
	static const uint8_t Td4[256];
	static const uint32_t rcon[15];

	void p_initSW(const uint8_t* key) noexcept;
	void p_encryptSW(const uint8_t* in, uint8_t* out) const noexcept;
	void p_decryptSW(const uint8_t* in, uint8_t* out) const noexcept;

	union {
#ifdef ZT_AES_AESNI
		struct {
			__m128i k[28];
			__m128i h[4];	 // h, hh, hhh, hhhh
			__m128i h2[4];	 // _mm_xor_si128(_mm_shuffle_epi32(h, 78), h), etc.
		} ni;
#endif

#ifdef ZT_AES_NEON
		struct {
			uint64_t hsw[2];   // in case it has AES but not PMULL, not sure if that ever happens
			uint8x16_t ek[15];
			uint8x16_t dk[15];
			uint8x16_t h;
		} neon;
#endif

		struct {
			uint64_t h[2];
			uint32_t ek[60];
			uint32_t dk[60];
		} sw;
	} p_k;

#ifdef ZT_AES_AESNI
	void p_init_aesni(const uint8_t* key) noexcept;
	void p_encrypt_aesni(const void* in, void* out) const noexcept;
	void p_decrypt_aesni(const void* in, void* out) const noexcept;
#endif

#ifdef ZT_AES_NEON
	void p_init_armneon_crypto(const uint8_t* key) noexcept;
	void p_encrypt_armneon_crypto(const void* in, void* out) const noexcept;
	void p_decrypt_armneon_crypto(const void* in, void* out) const noexcept;
#endif
};

}	// namespace ZeroTier

#endif