// Copyright (c) 2019-2024 Alexander Medvednikov. All rights reserved.
// Use of this source code is governed by an MIT license
// that can be found in the LICENSE file.
// Package hkdf implements the HMAC-based Extract-and-Expand Key Derivation
// Function (HKDF) as defined in RFC 5869.
//
// HKDF is a cryptographic key derivation function (KDF) with the goal of
// expanding limited input keying material into one or more cryptographically
// strong secret keys.
// Based off:   https://github.com/golang/go/tree/master/src/crypto/hkdf
module hkdf

// extract generates a pseudorandom key for use with `expand` from an input
// secret and an optional independent salt.
//
// Only use this function if you need to reuse the extracted key with multiple
// `expand` invocations and different context values. Most common scenarios,
// including the generation of multiple keys, should use `key` instead.
pub fn extract[H](h fn () H, secret []u8, salt []u8) ![]u8 {
	check_fips140_only(h, secret)!
	mut hash := h()
	mut xkey := salt.clone()
	if xkey.len == 0 {
		xkey = []u8{len: hash.size()}
	}
	return hmac_sum(h, xkey, secret)
}

// expand derives a key from the given hash, key, and optional context info,
// returning a []u8 of length key_length that can be used as cryptographic key.
// The extraction step is skipped.
//
// The key should have been generated by `extract`, or be a uniformly
// random or pseudorandom cryptographically strong key. See RFC 5869, Section
// 3.3. Most common scenarios will want to use `key` instead.
pub fn expand[H](h fn () H, pseudorandom_key []u8, info string, key_length int) ![]u8 {
	check_fips140_only(h, pseudorandom_key)!
	mut hash := h()
	if key_length < 0 {
		return error('hkdf: requested key length must be non-negative')
	}
	limit := hash.size() * 255
	if key_length > limit {
		return error('hkdf: requested key length too large')
	}
	mut out := []u8{cap: key_length}
	mut counter := u8(0)
	mut buf := []u8{}
	info_bytes := info.bytes()
	for out.len < key_length {
		counter++
		if counter == 0 {
			panic('hkdf: counter overflow')
		}
		mut data := []u8{cap: buf.len + info_bytes.len + 1}
		data << buf
		data << info_bytes
		data << counter
		buf = hmac_sum(h, pseudorandom_key, data)
		remaining := key_length - out.len
		if remaining < buf.len {
			out << buf[..remaining]
		} else {
			out << buf
		}
	}
	return out
}

// key derives a key from the given hash, secret, salt and context info,
// returning a []u8 of length key_length that can be used as cryptographic key.
// Salt and info can be empty.
pub fn key[H](h fn () H, secret []u8, salt []u8, info string, key_length int) ![]u8 {
	check_fips140_only(h, secret)!
	if key_length < 0 {
		return error('hkdf: requested key length must be non-negative')
	}
	mut hash := h()
	limit := hash.size() * 255
	if key_length > limit {
		return error('hkdf: requested key length too large')
	}
	prk := extract(h, secret, salt)!
	return expand(h, prk, info, key_length)
}

fn check_fips140_only[H](_h fn () H, _key []u8) ! {
	// V does not currently have a FIPS 140-only mode.
	return
}

fn hmac_sum[H](h fn () H, key []u8, data []u8) []u8 {
	mut hash := h()
	block_size := hash.block_size()
	mut b_key := if key.len <= block_size { key.clone() } else { hash_sum(h, key) }
	if b_key.len > block_size {
		b_key = b_key[..block_size].clone()
	}
	mut inner := []u8{len: block_size, init: 0x36}
	mut outer := []u8{len: block_size, init: 0x5c}
	for i, b in b_key {
		inner[i] = b ^ 0x36
		outer[i] = b ^ 0x5c
	}
	inner << data
	inner_hash := hash_sum(h, inner)
	outer << inner_hash
	return hash_sum(h, outer)
}

fn hash_sum[H](h fn () H, data []u8) []u8 {
	mut hash := h()
	hash.write(data) or { panic(err) }
	return hash.sum([]u8{})
}