Chacha20 encryption

ChaCha20.h

/*
Base32.hpp -- Chacha20 encryption
Copyright 2018 yni3
License : Public Domain
*/

#pragma once

#ifndef __CHACHA20
#define __CHACHA20

class ChaCha20 {
	
	typedef char s8;
	typedef int s32;
	typedef unsigned long long u64;
	typedef unsigned int u32;
	typedef unsigned char u8;

	void memcpy(u8* dist, const u8* src, size_t n)
	{
		unsigned char* ss1 = dist;
		const unsigned char* ss2 = src;
		const unsigned char* end = src + n;
		while (src != end)
			*dist++ = *src++;
	}

	void zeromemcpy(u8* dist, size_t n)
	{
		unsigned char* ss1 = dist;
		const unsigned char* end = dist + n;
		while (dist != end)
			*dist++ = 0;
	}

	static constexpr u32 rotl32(const u32 &x, s32 n) {
		return (x << n) | (x >> (-n & 31));
	}

	static constexpr u32 ReadLE32(const u8 *a) {
		return
			u32(a[0] << 0 * 8) |
			u32(a[1] << 1 * 8) |
			u32(a[2] << 2 * 8) |
			u32(a[3] << 3 * 8);
	}

	static void WriteLE32(u8* dst, const u32 &src) {
		dst[0] = (src >> 0 * 8) & 0xff;
		dst[1] = (src >> 1 * 8) & 0xff;
		dst[2] = (src >> 2 * 8) & 0xff;
		dst[3] = (src >> 3 * 8) & 0xff;
	}

	void Init(const u8 key[32], const u8 nonce[8]) {
		const u8 *magic_constant = (u8*)"expand 32-byte k";
		state[0] = ReadLE32(magic_constant + 0 * 4);
		state[1] = ReadLE32(magic_constant + 1 * 4);
		state[2] = ReadLE32(magic_constant + 2 * 4);
		state[3] = ReadLE32(magic_constant + 3 * 4);
		state[4] = ReadLE32(key + 0 * 4);
		state[5] = ReadLE32(key + 1 * 4);
		state[6] = ReadLE32(key + 2 * 4);
		state[7] = ReadLE32(key + 3 * 4);
		state[8] = ReadLE32(key + 4 * 4);
		state[9] = ReadLE32(key + 5 * 4);
		state[10] = ReadLE32(key + 6 * 4);
		state[11] = ReadLE32(key + 7 * 4);
		state[12] = 0;
		state[13] = 0;
		state[14] = ReadLE32(nonce + 0 * 4);
		state[15] = ReadLE32(nonce + 1 * 4);
	}

	void seek(u64 counter) {
		state[12] = static_cast<u32>(counter);
		state[13] = static_cast<u32>(counter >> 32);
	}

	static inline void QUARTERROUND(unsigned int *x, const int a, const int b, const int c, const int d)
	{
		x[a] += x[b]; x[d] = rotl32(x[d] ^ x[a], 16);
		x[c] += x[d]; x[b] = rotl32(x[b] ^ x[c], 12);
		x[a] += x[b]; x[d] = rotl32(x[d] ^ x[a], 8);
		x[c] += x[d]; x[b] = rotl32(x[b] ^ x[c], 7);
	}

	void next(u32 out[16]) {
		for (int i = 0; i < 16; i++) out[i] = state[i];

		for (int i = 0; i < 10; i++) {
			QUARTERROUND(out, 0, 4, 8, 12);
			QUARTERROUND(out, 1, 5, 9, 13);
			QUARTERROUND(out, 2, 6, 10, 14);
			QUARTERROUND(out, 3, 7, 11, 15);
			QUARTERROUND(out, 0, 5, 10, 15);
			QUARTERROUND(out, 1, 6, 11, 12);
			QUARTERROUND(out, 2, 7, 8, 13);
			QUARTERROUND(out, 3, 4, 9, 14);
		}
		for (int i = 0; i < 16; i++) out[i] += state[i];

		//update state
		state[12] = state[12] + 1;
		if (state[12] == 0) { state[13] = state[13] + 1; }
	}

	u32 state[16];
	u32 keystream32[16];
	size_t position;

public:

	ChaCha20(
		const s8* key,
		s32 size,
		u64 counter = 0
	) : position(64) {
		u8 keybuff[32];
		u8 noncebuff[8];
		zeromemcpy(keybuff, 32);
		zeromemcpy(noncebuff, 8);
		for (s32 i = 0; i < size && i < 32; i++) {
			keybuff[i] = key[i];
		}
		Init(keybuff, noncebuff);
		seek(counter);
	}

	ChaCha20(
		const s8* key,
		s32 skey,
		const s8* nonce,
		s32 snonce,
		u64 counter = 0
	) : position(64) {
		u8 keybuff[32];
		u8 noncebuff[8];
		zeromemcpy(keybuff, 32);
		zeromemcpy(noncebuff, 8);
		for (s32 i = 0; i < 32 && i < skey; i++) {
			keybuff[i] = key[i];
		}
		for (s32 i = 0; i < 8 && i < snonce; i++) {
			noncebuff[i] = nonce[i];
		}
		Init(keybuff, noncebuff);
		seek(counter);
	}

	ChaCha20() : position(64) {
		zeromemcpy(reinterpret_cast<u8*>(&state), sizeof(u32) * 16);
		zeromemcpy(reinterpret_cast<u8*>(&keystream32), sizeof(u32) * 16);
	}

	void setPassWrod(const s8* key,s32 size, u64 counter = 0) {
		u8 keybuff[32];
		u8 noncebuff[8];
		zeromemcpy(keybuff, 32);
		zeromemcpy(noncebuff, 8);
		for (s32 i = 0; i < size; i++) {
			keybuff[i] = key[i];
		}
		Init(keybuff, noncebuff);
		seek(counter);
	}

	//encryption
	void crypt(u8 *bytes, size_t n_bytes) {
		u8 blk[64];
		u8 *keystream8 = (u8*)keystream32;
		for (size_t i = 0; i < n_bytes; i++) {
			if (position >= 64) {
				next(keystream32);
				position = 0;
			}
			for (s32 j = 0; j < 16; ++j) { WriteLE32(blk + (4 * i), keystream32[i]); }
			bytes[i] ^= blk[position];
			position++;
		}
	}

	//encryption
	void update(u8 *dst, const u8 *src, s32 srclen) {
		memcpy(dst, src, srclen);
		crypt(dst, srclen);
	}

};

#endif