package hkdf
import (
  "crypto/hmac"
  "hash"
  "io"
)
type HKDF struct {
  hashFunc hash.Hash
  size     int
  info    []byte
  counter byte
  prev []byte
  buf  []byte
}
// Extract HKDF-Extract(salt, IKM) -> PRK
func Extract(hash func() hash.Hash, secret, salt []byte) []byte {
  if salt == nil {
    salt = make([]byte, hash().Size())
  }
  hashFunc := hmac.New(hash, salt)
  hashFunc.Write(secret)
  return hashFunc.Sum(nil)
}
func (f *HKDF) Read(p []byte) (int, error) {
  need := len(p)
  n := copy(p, f.buf)
  p = p[n:]
  for len(p) > 0 {
    f.hashFunc.Reset()
    f.hashFunc.Write(f.prev)
    f.hashFunc.Write(f.info)
    f.hashFunc.Write([]byte{f.counter})
    f.prev = f.hashFunc.Sum(f.prev[:0])
    f.counter++
    f.buf = f.prev
    n = copy(p, f.buf)
    p = p[n:]
  }
  f.buf = f.buf[n:]
  return need, nil
}
// Expand HKDF-Expand(PRK, info, L) -> OKM
func Expand(hash func() hash.Hash, prk, info []byte) io.Reader {
  hashFunc := hmac.New(hash, prk)
  return &HKDF{hashFunc, hashFunc.Size(), info, 1, nil, nil}
}
func New(hash func() hash.Hash, secret, salt, info []byte) io.Reader {
  prk := Extract(hash, secret, salt)
  return Expand(hash, prk, info)
}