// 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{})
}