package lafsession

import (
	"crypto/sha256"
	"io"

	"golang.org/x/crypto/chacha20poly1305"
	"golang.org/x/crypto/hkdf"
)

const (
	keyLen   = chacha20poly1305.KeySize
	nonceLen = chacha20poly1305.NonceSize
	tagLen   = 16
)

type trafficKeys struct {
	c2s [keyLen]byte
	s2c [keyLen]byte
}

func deriveTrafficKeys(sharedSecret, transcriptHash []byte, version uint16, suite uint16) (*trafficKeys, error) {
	return deriveTrafficKeysImpl(sharedSecret, transcriptHash, version, suite)
}

func deriveTrafficKeysImpl(sharedSecret, transcriptHash []byte, version uint16, suite uint16) (*trafficKeys, error) {
	info := buildHKDFInfo(version, suite, transcriptHash)
	r := hkdf.New(sha256.New, sharedSecret, nil, info)
	var tk trafficKeys
	if _, err := io.ReadFull(r, tk.c2s[:]); err != nil {
		return nil, err
	}
	if _, err := io.ReadFull(r, tk.s2c[:]); err != nil {
		return nil, err
	}
	return &tk, nil
}

func buildHKDFInfo(version, suite uint16, transcriptHash []byte) []byte {
	const label = "lafsession ecdh-aead v1\x00"
	info := make([]byte, 0, len(label)+2+2+len(transcriptHash))
	info = append(info, label...)
	info = append(info, byte(version>>8), byte(version))
	info = append(info, byte(suite>>8), byte(suite))
	info = append(info, transcriptHash...)
	return info
}

func zero(b []byte) {
	for i := range b {
		b[i] = 0
	}
}

func (tk *trafficKeys) wipe() {
	if tk == nil {
		return
	}
	zero(tk.c2s[:])
	zero(tk.s2c[:])
}