gistfile1.txt

NTSTATUS GenDHKey(_Out_ BCRYPT_KEY_HANDLE* phKey)
{
	NTSTATUS status;
	BCRYPT_ALG_HANDLE hAlgorithm;

	if (0 <= (status = BCryptOpenAlgorithmProvider(&hAlgorithm, BCRYPT_DH_ALGORITHM, 0, 0)))
	{
		BCRYPT_KEY_HANDLE hKey;

		status = BCryptGenerateKeyPair(hAlgorithm, &hKey, 0, 0);

		BCryptCloseAlgorithmProvider(hAlgorithm, 0);

		if (0 <= status)
		{
			if (0 > (status = BCryptFinalizeKeyPair(hKey, 0)))
			{
				BCryptDestroyKey(hKey);
			}
			else
			{
				*phKey = hKey;
			}
		}
	}

	return status;
}

NTSTATUS GenAesKey(_Out_ BCRYPT_KEY_HANDLE* phKey, PBYTE pbSecret, ULONG cbSecret, ULONG mbl = 1)
{
	NTSTATUS status;
	BCRYPT_ALG_HANDLE hAlgorithm;
	if (0 <= (status = BCryptOpenAlgorithmProvider(&hAlgorithm, BCRYPT_AES_ALGORITHM, 0, 0)))
	{
		BCRYPT_KEY_HANDLE hKey;

		status = BCryptGenerateSymmetricKey(hAlgorithm, &hKey, 0, 0, pbSecret, cbSecret, 0);

		BCryptCloseAlgorithmProvider(hAlgorithm, 0);

		if (0 <= status)
		{
			if (0 <= (status = BCryptSetProperty(hKey, BCRYPT_CHAINING_MODE, (PBYTE)BCRYPT_CHAIN_MODE_CFB, sizeof(BCRYPT_CHAIN_MODE_CFB), 0)) && 
				0 <= (status = BCryptSetProperty(hKey, BCRYPT_MESSAGE_BLOCK_LENGTH, (PBYTE)&mbl, sizeof(mbl), 0)))
			{
				*phKey = hKey;
				return STATUS_SUCCESS;
			}

			BCryptDestroyKey(hKey);
		}
	}

	return status;
}

struct CClientServer
{
	CClientServer* m_pLink = 0;

	BCRYPT_KEY_HANDLE m_hPriv = 0;
	BCRYPT_KEY_HANDLE m_hPub = 0;
	BCRYPT_KEY_HANDLE m_hSession = 0;

	PCSTR m_name;
	bool m_bstart;

	CClientServer(PCSTR name, bool bstart) : m_name(name), m_bstart(bstart)
	{
	}

	~CClientServer()
	{
		if (m_hSession) BCryptDestroyKey(m_hSession);
		if (m_hPub) BCryptDestroyKey(m_hPub);
		if (m_hPriv) BCryptDestroyKey(m_hPriv);
	}

	NTSTATUS ImportPubKey(PBYTE pb, ULONG cb)
	{
		NTSTATUS status;
		BCRYPT_ALG_HANDLE hAlgorithm;

		if (0 <= (status = BCryptOpenAlgorithmProvider(&hAlgorithm, BCRYPT_DH_ALGORITHM, 0, 0)))
		{
			status = BCryptImportKeyPair(hAlgorithm, 0, BCRYPT_PUBLIC_KEY_BLOB, &m_hPub, pb, cb, 0);

			BCryptCloseAlgorithmProvider(hAlgorithm, 0);
		}

		return status;
	}

	NTSTATUS SharePubKey(BCRYPT_KEY_HANDLE hKey)
	{
		NTSTATUS status;
		PBYTE pb = 0;
		ULONG cb = 0;
		while (0 <= (status = BCryptExportKey(hKey, 0, BCRYPT_PUBLIC_KEY_BLOB, pb, cb, &cb, 0)))
		{
			if (pb)
			{
				status = m_pLink->OnData(pb, cb);
				break;
			}

			pb = (PBYTE)alloca(cb);
		}

		return status;
	}

	NTSTATUS Start()
	{
		NTSTATUS status;
		0 <= (status = GenDHKey(&m_hPriv)) && 0 <= (status = SharePubKey(m_hPriv));
		return status;
	}

	NTSTATUS SendMsg()
	{
		PBYTE psz = 0;
		ULONG len = 0;
		while (0 < (int)(len = _snprintf((PSTR)psz, len, "message from [%hs]", m_name)))
		{
			if (psz)
			{
				NTSTATUS status;
				if (0 <= (status = BCryptEncrypt(m_hSession, psz, len, 0, 0, 0, psz, len, &len, 0)))
				{
					return m_pLink->OnData(psz, len);
				}
				return status;
			}

			psz = (PBYTE)alloca(len);
		}
		return STATUS_INTERNAL_ERROR;
	}

	NTSTATUS DeriveSessionKey(PCSTR str = ">><<")
	{
		NTSTATUS status;

		BCRYPT_SECRET_HANDLE hSecret = 0;
		if (0 <= (status = BCryptSecretAgreement(m_hPriv, m_hPub, &hSecret, 0)))
		{
			if (0 <= (status = BCryptSetProperty(hSecret, BCRYPT_HKDF_HASH_ALGORITHM, 
				(PBYTE)BCRYPT_SHA256_ALGORITHM, sizeof(BCRYPT_SHA256_ALGORITHM), 0)) &&
				0 <= (status = BCryptSetProperty(hSecret, BCRYPT_HKDF_SALT_AND_FINALIZE, 0, 0, 0)))
			{
				BCryptBuffer Parameter = { (ULONG)strlen(str), KDF_HKDF_INFO, (PVOID)str };
				BCryptBufferDesc ParameterList = { BCRYPTBUFFER_VERSION, 1, &Parameter } ;
				ULONG cb;
				BYTE aes[32];
				if (0 <= (status = BCryptDeriveKey(hSecret, BCRYPT_KDF_HKDF, &ParameterList, aes, sizeof(aes), &cb, 0)))
				{
					PSTR psz = 0;
					ULONG cch = 0;
					while (CryptBinaryToStringA(aes, cb, CRYPT_STRING_BASE64, psz, &cch))
					{
						if (psz)
						{
							DbgPrint("%hs: %hs\r\n", m_name, psz);
							break;
						}

						psz = (PSTR)alloca(cch);
					}

					if (0 <= (status = GenAesKey(&m_hSession, aes, cb)))
					{
						if (m_bstart)
						{
							SendMsg();
						}
					}
				}
			}

			BCryptDestroySecret(hSecret);
		}

		return status;
	}

	NTSTATUS OnData(PBYTE pb, ULONG cb)
	{
		NTSTATUS status;
		if (!m_hPub)
		{
			if (0 <= (status = ImportPubKey(pb, cb)))
			{
				if (!m_hPriv)
				{
					if (0 > (status = Start()))
					{
						return status;
					}
				}

				return DeriveSessionKey();
			}
		}

		if (m_hSession)
		{
			if (0 <= (status = BCryptDecrypt(m_hSession, pb, cb, 0, 0, 0, pb, cb, &cb, 0)))
			{
				DbgPrint("%hs: << %.*hs\r\n", m_name, cb, pb);

				if (!m_bstart)
				{
					return SendMsg();
				}
			}
		}

		return STATUS_INVALID_PARAMETER;
	}
};

		CClientServer c("client", true), s("server", false);
		c.m_pLink = &s;
		s.m_pLink = &c;
		s.Start();