An implementation of balloon hashing in C#

BalloonHash.cs

/*

An implementation of Balloon Hashing in C#

This hasn't been thoroughly tested and it SHOULD NOT BE USED FOR PRODUCTION.

Balloon hashing is a fairly easy way to turn any cryptographic hash function
primitive into a memory-hard and parallelisable password storage funciton or
key derivation function (KDF).

This implementation follows the algorithm described in the paper, but does not
attempt to provide direct parity with the reference implementation provided by
the paper authors. It simply implements the same general algorithm.

Step 2b of the algorithm deals with the random neighbour hashing. The paper 
uses an `ints_to_block` function that is not described, along with the hash 
function, to pick a random neighbour block to include in the hash input. The 
reference source, however, does not use this approach. It instead uses AES-CTR
with a key derived from the salt to create a pseudorandom bitstream generator
which is then invoked at each step to produce a 64-bit random number from
which the neighbour block index can be derived. This is the approach I have 
taken in this code too. Ultimately both approaches have the same security 
properties, although the AES-CTR approach does have the small added bonus that
a GPU or FPGA accelerated implementation would be even higher effort.

ComputeHash is the regular single-threaded API. It takes in a hash name, the 
password, the salt, a space cost in bytes, and a time cost.

ComputeHashParallel is the parallelised variant. It takes the same parameters 
as ComputeHash, along with a thread count. The thread count is integrated into
the PRF seed, and the final combination of the hashes uses the same basic xor
approach as the paper.

ComputeHashParallelExtended is an additional variant which provides arbitrary 
output sizes. It achieves this by replacing the xor combination step with an 
invocation of PBKDF2 over all the results from the threads, using the same 
hash algorithm and an iteration count equal to the time cost. This isn't part
of the paper but ultimately has the same security properties.

Based on "Balloon Hashing: A Memory-Hard Function Providing Provable Protection Against Sequential Attacks" by Dan Boneh, Henry Corrigan-Gibbs, and Stuart Schechter.
Site: https://crypto.stanford.edu/balloon/
Paper: https://eprint.iacr.org/2016/027.pdf
Reference source: https://github.com/henrycg/balloon/

*/

static class BalloonHash
{
	/// <summary>Computes a parallel balloon function hash for the given password and salt, returning a byte sequence of the specified length.</summary>
	public static byte[] ComputeHashParallelExtended(HashAlgorithmName hashName, byte[] password, byte[] salt, int spaceCost, int timeCost, int threadCount, int outputLength)
	{
		int hashSize;
		using (var tempHash = IncrementalHash.CreateHash(hashName))
		{
			hashSize = tempHash.HashLengthInBytes;
		}
		byte[] hashes = new byte[threadCount * hashSize];
		var po = new ParallelOptions { MaxDegreeOfParallelism = threadCount };
		Parallel.For(0, threadCount, po, (int threadIndex) =>
		{
			byte[] passwordCopy = new byte[password.Length];
			Array.Copy(password, passwordCopy, password.Length);
			byte[] extendedSalt = new byte[salt.Length + sizeof(int)];
			Array.Copy(salt, extendedSalt, salt.Length);
			if (!BitConverter.TryWriteBytes(extendedSalt.AsSpan().Slice(salt.Length), threadIndex))
				throw new CryptographicUnexpectedOperationException("Failed to write to buffer.");
			
			byte[] hash = ComputeHash(hashName, passwordCopy, extendedSalt, spaceCost, timeCost);
			Array.Copy(hash, 0, hashes, threadIndex * hashSize, hashSize);
		});
		return Rfc2898DeriveBytes.Pbkdf2(hashes, salt, timeCost, hashName, outputLength);
	}

	/// <summary>Computes a parallel balloon function hash for the given password and salt.</summary>
	public static byte[] ComputeHashParallel(HashAlgorithmName hashName, byte[] password, byte[] salt, int spaceCost, int timeCost, int threadCount)
	{
		byte[][] hashes = new byte[threadCount][];
		var po = new ParallelOptions { MaxDegreeOfParallelism = threadCount };
		Parallel.For(0, threadCount, po, (int threadIndex) => {
			byte[] passwordCopy = new byte[password.Length];
			Array.Copy(password, passwordCopy, password.Length);
			byte[] extendedSalt = new byte[salt.Length + sizeof(int)];
			Array.Copy(salt, extendedSalt, salt.Length);
			if (!BitConverter.TryWriteBytes(extendedSalt.AsSpan().Slice(salt.Length), threadIndex))
				throw new CryptographicUnexpectedOperationException("Failed to write to buffer.");
			hashes[threadIndex] = ComputeHash(hashName, passwordCopy, extendedSalt, spaceCost, timeCost);
		});
		for (int n = 1; n < hashes.Length; n++)
		{
			for (int i = 0; i < hashes[0].Length; i++)
			{
				hashes[0][i] ^= hashes[n][i];
			}
		}
		return hashes[0];
	}
	
	/// <summary>Computes a balloon function hash for the given password and salt.</summary>
	public static byte[] ComputeHash(HashAlgorithmName hashName, ReadOnlySpan<byte> password, ReadOnlySpan<byte> salt, int spaceCost, int timeCost)
	{
		if (hashName == null)
		{
			throw new ArgumentNullException(nameof(hashName));
		}
		if (password == null)
		{
			throw new ArgumentNullException(nameof(password));
		}
		if (salt == null)
		{
			throw new ArgumentNullException(nameof(salt));
		}
		if (timeCost < 1)
		{
			throw new ArgumentOutOfRangeException(nameof(timeCost), "Time cost must be greater than or equal to 1.");
		}
		using var hashFunction = IncrementalHash.CreateHash(hashName);
		int hashSizeBytes = hashFunction.HashLengthInBytes;
		if (spaceCost < hashSizeBytes * 2)
		{
			throw new ArgumentOutOfRangeException(nameof(spaceCost), "Space cost must be greater than or equal to the output size of the hash function multiplied by two.");
		}
		if ((spaceCost % hashSizeBytes) != 0)
		{
			// round up the space cost so it matches the output size of the hash function
			spaceCost += hashSizeBytes;
			spaceCost -= spaceCost % hashSizeBytes;
		}
		int spaceCostBytes = spaceCost;
		spaceCost /= hashSizeBytes;

		const int delta = 3;

		byte[] bitstreamSeed = new byte[salt.Length + sizeof(int) + sizeof(int)];
		var bitstreamSpan = bitstreamSeed.AsSpan();
		salt.CopyTo(bitstreamSpan);
		if (!BitConverter.TryWriteBytes(bitstreamSpan.Slice(salt.Length, sizeof(int)), spaceCost))
			throw new CryptographicUnexpectedOperationException("Failed to write to buffer.");
		if (!BitConverter.TryWriteBytes(bitstreamSpan.Slice(salt.Length + sizeof(int), sizeof(int)), timeCost))
			throw new CryptographicUnexpectedOperationException("Failed to write to buffer.");
		using var bitStream = new BalloonHashBitstream(hashName, bitstreamSpan);

		Span<byte> buffer = new byte[spaceCostBytes];

		UInt64 counter = 0;
		Span<byte> counterBuffer = stackalloc byte[sizeof(UInt64)];
		Span<byte> indexBuffer = stackalloc byte[hashSizeBytes];
		
		if (!BitConverter.TryWriteBytes(counterBuffer, counter))
			throw new CryptographicUnexpectedOperationException("Failed to write to buffer.");

		// 1. expand input into buffer
		hashFunction.AppendData(counterBuffer);
		hashFunction.AppendData(password);
		hashFunction.AppendData(salt);
		hashFunction.GetHashAndReset(buffer);
		for (int m = 1; m < spaceCost; m++)
		{
			unchecked { counter++; }
			if (!BitConverter.TryWriteBytes(counterBuffer, counter))
				throw new CryptographicUnexpectedOperationException("Failed to write to buffer.");
			hashFunction.AppendData(counterBuffer); // ctr
			hashFunction.AppendData(buffer.Slice((m - 1) * hashSizeBytes, hashSizeBytes)); // buf[m-1]
			hashFunction.GetHashAndReset(buffer.Slice(m * hashSizeBytes));
		}

		// 2. mix buffer contents
		for (int t = 0; t < timeCost; t++)
		{
			for (int m = 0; m < spaceCost; m++)
			{
				int prevIndex = (m + spaceCost - 1) % spaceCost;

				// 2a. hash last and current blocks
				unchecked { counter++; }
				if (!BitConverter.TryWriteBytes(counterBuffer, counter))
					throw new CryptographicUnexpectedOperationException("Failed to write to buffer.");
				
				hashFunction.AppendData(counterBuffer); // ctr
				hashFunction.AppendData(buffer.Slice(prevIndex * hashSizeBytes, hashSizeBytes)); // prev
				hashFunction.AppendData(buffer.Slice(m * hashSizeBytes, hashSizeBytes)); // buf[m]
				hashFunction.GetHashAndReset(buffer.Slice(m * hashSizeBytes));
        
				// 2b. hash in pseudorandomly chosen blocks
				for (int i = 0; i < delta; i++)
				{
					// get a block index. this comes from an AES-CTR based PRNG seeded with the salt
					UInt64 idx_block = bitStream.GetRandomUInt64();
					int otherIdx = (int)(idx_block % (UInt64)spaceCost);

					// hash in the neighbour
					unchecked { counter++; }
					if (!BitConverter.TryWriteBytes(counterBuffer, counter))
						throw new CryptographicUnexpectedOperationException("Failed to write to buffer.");
					hashFunction.AppendData(counterBuffer); // ctr
					hashFunction.AppendData(buffer.Slice(m * hashSizeBytes, hashSizeBytes)); // buf[m]
					hashFunction.AppendData(buffer.Slice(otherIdx * hashSizeBytes, hashSizeBytes)); // buf[other]
					hashFunction.GetHashAndReset(buffer.Slice(m * hashSizeBytes));
				}
			}
		}

		return buffer.Slice((spaceCost - 1) * hashSizeBytes, hashSizeBytes).ToArray();
	}

	private class BalloonHashBitstream : IDisposable
	{
		private Aes _aes;
		private ICryptoTransform _encrypt;
		private ulong _ctr = 0; // will wrap at 2^64 steps but that's an unreasonable expectation here
		private readonly object _syncRoot = new object();
		bool disposedValue;

		public BalloonHashBitstream(HashAlgorithmName hashName, ReadOnlySpan<byte> seed)
		{
			if (seed == null)
			{
				throw new ArgumentNullException(nameof(seed));
			}
			using var hash = IncrementalHash.CreateHash(hashName);
			const int AesKeySize = 128;
			if (hash.HashLengthInBytes < AesKeySize / 8)
			{
				throw new CryptographicException("The provided hash function must have a block size of at least 128 bits.");
			}
			_aes = Aes.Create();
			_aes.KeySize = AesKeySize;
			_aes.Mode = CipherMode.ECB; // we're rolling our own CTR here
			hash.AppendData(seed);
			Span<byte> key = hash.GetHashAndReset();
			_aes.Key = key.Slice(0, AesKeySize / 8).ToArray();
			_aes.IV = new byte[_aes.BlockSize / 8];
			_encrypt = _aes.CreateEncryptor();
		}

		public UInt64 GetRandomUInt64()
		{
			byte[] bufferIn = new byte[_aes.BlockSize / 8];
			byte[] bufferOut = new byte[_aes.BlockSize / 8];
			lock (_syncRoot)
			{
				unchecked { _ctr++; }
				BitConverter.TryWriteBytes(bufferIn.AsSpan(), _ctr);
				_encrypt.TransformBlock(bufferIn, 0, bufferIn.Length, bufferOut, 0);
			}
			return BitConverter.ToUInt64(bufferOut, 0);
		}

		protected virtual void Dispose(bool disposing)
		{
			if (!disposedValue)
			{
				if (disposing)
				{
					_encrypt.Dispose();
					_aes.Dispose();
				}
				disposedValue = true;
			}
		}

		public void Dispose()
		{
			// Do not change this code. Put cleanup code in 'Dispose(bool disposing)' method
			Dispose(disposing: true);
			GC.SuppressFinalize(this);
		}
	}
}