module hkdf

import crypto.sha1
import crypto.sha256
import encoding.hex
import hash

struct HkdfTest {
	new_hash fn () hash.Hash @[required]
	master   string
	salt     string
	prk      string
	info     string
	out      string
}

fn sha1_hash() hash.Hash {
	return sha1.new()
}

fn sha256_hash() hash.Hash {
	return sha256.new()
}

fn decode(s string) []u8 {
	return hex.decode(s) or { panic(err) }
}

fn hkdf_tests() []HkdfTest {
	// Tests from RFC 5869
	return [
		HkdfTest{
			new_hash: sha256_hash
			master:   '0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b'
			salt:     '000102030405060708090a0b0c'
			prk:      '077709362c2e32df0ddc3f0dc47bba6390b6c73bb50f9c3122ec844ad7c2b3e5'
			info:     'f0f1f2f3f4f5f6f7f8f9'
			out:      '3cb25f25faacd57a90434f64d0362f2a2d2d0a90cf1a5a4c5db02d56ecc4c5bf34007208d5b887185865'
		},
		HkdfTest{
			new_hash: sha256_hash
			master:   '000102030405060708090a0b0c0d0e0f' + '101112131415161718191a1b1c1d1e1f' +
				'202122232425262728292a2b2c2d2e2f' + '303132333435363738393a3b3c3d3e3f' +
				'404142434445464748494a4b4c4d4e4f'
			salt:     '606162636465666768696a6b6c6d6e6f' + '707172737475767778797a7b7c7d7e7f' +
				'808182838485868788898a8b8c8d8e8f' + '909192939495969798999a9b9c9d9e9f' +
				'a0a1a2a3a4a5a6a7a8a9aaabacadaeaf'
			prk:      '06a6b88c5853361a06104c9ceb35b45cef760014904671014a193f40c15fc244'
			info:     'b0b1b2b3b4b5b6b7b8b9babbbcbdbebf' + 'c0c1c2c3c4c5c6c7c8c9cacbcccdcecf' +
				'd0d1d2d3d4d5d6d7d8d9dadbdcdddedf' + 'e0e1e2e3e4e5e6e7e8e9eaebecedeeef' +
				'f0f1f2f3f4f5f6f7f8f9fafbfcfdfeff'
			out:      'b11e398dc80327a1c8e7f78c596a4934' + '4f012eda2d4efad8a050cc4c19afa97c' +
				'59045a99cac7827271cb41c65e590e09' + 'da3275600c2f09b8367793a9aca3db71' +
				'cc30c58179ec3e87c14c01d5c1f3434f1d87'
		},
		HkdfTest{
			new_hash: sha256_hash
			master:   '0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b'
			salt:     ''
			prk:      '19ef24a32c717b167f33a91d6f648bdf96596776afdb6377ac434c1c293ccb04'
			info:     ''
			out:      '8da4e775a563c18f715f802a063c5a31b8a11f5c5ee1879ec3454e5f3c738d2d9d201395faa4b61a96c8'
		},
		HkdfTest{
			new_hash: sha256_hash
			master:   '0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b'
			salt:     ''
			prk:      '19ef24a32c717b167f33a91d6f648bdf96596776afdb6377ac434c1c293ccb04'
			info:     ''
			out:      '8da4e775a563c18f715f802a063c5a31b8a11f5c5ee1879ec3454e5f3c738d2d9d201395faa4b61a96c8'
		},
		HkdfTest{
			new_hash: sha1_hash
			master:   '0b0b0b0b0b0b0b0b0b0b0b'
			salt:     '000102030405060708090a0b0c'
			prk:      '9b6c18c432a7bf8f0e71c8eb88f4b30baa2ba243'
			info:     'f0f1f2f3f4f5f6f7f8f9'
			out:      '085a01ea1b10f36933068b56efa5ad81a4f14b822f5b091568a9cdd4f155fda2c22e422478d305f3f896'
		},
		HkdfTest{
			new_hash: sha1_hash
			master:   '000102030405060708090a0b0c0d0e0f' + '101112131415161718191a1b1c1d1e1f' +
				'202122232425262728292a2b2c2d2e2f' + '303132333435363738393a3b3c3d3e3f' +
				'404142434445464748494a4b4c4d4e4f'
			salt:     '606162636465666768696a6b6c6d6e6f' + '707172737475767778797a7b7c7d7e7f' +
				'808182838485868788898a8b8c8d8e8f' + '909192939495969798999a9b9c9d9e9f' +
				'a0a1a2a3a4a5a6a7a8a9aaabacadaeaf'
			prk:      '8adae09a2a307059478d309b26c4115a224cfaf6'
			info:     'b0b1b2b3b4b5b6b7b8b9babbbcbdbebf' + 'c0c1c2c3c4c5c6c7c8c9cacbcccdcecf' +
				'd0d1d2d3d4d5d6d7d8d9dadbdcdddedf' + 'e0e1e2e3e4e5e6e7e8e9eaebecedeeef' +
				'f0f1f2f3f4f5f6f7f8f9fafbfcfdfeff'
			out:      '0bd770a74d1160f7c9f12cd5912a06eb' + 'ff6adcae899d92191fe4305673ba2ffe' +
				'8fa3f1a4e5ad79f3f334b3b202b2173c' + '486ea37ce3d397ed034c7f9dfeb15c5e' +
				'927336d0441f4c4300e2cff0d0900b52d3b4'
		},
		HkdfTest{
			new_hash: sha1_hash
			master:   '0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b'
			salt:     ''
			prk:      'da8c8a73c7fa77288ec6f5e7c297786aa0d32d01'
			info:     ''
			out:      '0ac1af7002b3d761d1e55298da9d0506b9ae52057220a306e07b6b87e8df21d0ea00033de03984d34918'
		},
		HkdfTest{
			new_hash: sha1_hash
			master:   '0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c'
			salt:     ''
			prk:      '2adccada18779e7c2077ad2eb19d3f3e731385dd'
			info:     ''
			out:      '2c91117204d745f3500d636a62f64f0ab3bae548aa53d423b0d1f27ebba6f5e5673a081d70cce7acfc48'
		},
	]
}

fn test_hkdf() {
	for i, tt in hkdf_tests() {
		master := decode(tt.master)
		salt := decode(tt.salt)
		prk_expected := decode(tt.prk)
		info := decode(tt.info)
		out_expected := decode(tt.out)

		prk := extract(tt.new_hash, master, salt)!
		assert prk == prk_expected, 'test ${i}: incorrect PRK: have ${prk}, need ${prk_expected}'

		derived_key := key(tt.new_hash, master, salt, info.bytestr(), out_expected.len)!
		assert derived_key == out_expected, 'test ${i}: incorrect output: have ${derived_key}, need ${out_expected}'

		expanded := expand(tt.new_hash, prk, info.bytestr(), out_expected.len)!
		assert expanded == out_expected, 'test ${i}: incorrect output from expand: have ${expanded}, need ${out_expected}'
	}
}

fn test_hkdf_limit() {
	master := []u8{len: 4, init: u8(index)}
	info := ''
	limit := sha1.new().size() * 255

	// The maximum output bytes should be extractable
	out := key(sha1_hash, master, []u8{}, info, limit)!
	assert out.len == limit

	// Reading one more should return an error
	if _ := key(sha1_hash, master, []u8{}, info, limit + 1) {
		assert false, 'expected key derivation to fail, but it succeeded'
	} else {
		assert err.msg() == 'hkdf: requested key length too large'
	}
}

fn test_direct_hash_constructor() {
	out := key(sha256.new, [u8(0), 1, 2, 3], []u8{}, 'context', sha256.size)!
	assert out.len == sha256.size
	assert out != []u8{len: sha256.size}
}