// Copyright 2023 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.


// Package mlkem implements the quantum-resistant key encapsulation method
// ML-KEM (formerly known as Kyber), as specified in [NIST FIPS 203].
//
// [NIST FIPS 203]: https://doi.org/10.6028/NIST.FIPS.203
package mlkem

// This package targets security, correctness, simplicity, readability, and
// reviewability as its primary goals. All critical operations are performed in
// constant time.
//
// Variable and function names, as well as code layout, are selected to
// facilitate reviewing the implementation against the NIST FIPS 203 document.
//
// Reviewers unfamiliar with polynomials or linear algebra might find the
// background at https://words.filippo.io/kyber-math/ useful.
//
// This file implements the recommended parameter set ML-KEM-768. The ML-KEM-1024
// parameter set implementation is auto-generated from this file.
//
//go:generate go run generate1024.go -input mlkem768.go -output mlkem1024.go

import (
	
	
	
	
	
)

const (
	// ML-KEM global constants.
	n = 256
	q = 3329

	// encodingSizeX is the byte size of a ringElement or nttElement encoded
	// by ByteEncode_X (FIPS 203, Algorithm 5).
	encodingSize12 = n * 12 / 8
	encodingSize11 = n * 11 / 8
	encodingSize10 = n * 10 / 8
	encodingSize5  = n * 5 / 8
	encodingSize4  = n * 4 / 8
	encodingSize1  = n * 1 / 8

	messageSize = encodingSize1

	SharedKeySize = 32
	SeedSize      = 32 + 32
)

// ML-KEM-768 parameters.
const (
	k = 3

	CiphertextSize768       = k*encodingSize10 + encodingSize4
	EncapsulationKeySize768 = k*encodingSize12 + 32
)

// ML-KEM-1024 parameters.
const (
	k1024 = 4

	CiphertextSize1024       = k1024*encodingSize11 + encodingSize5
	EncapsulationKeySize1024 = k1024*encodingSize12 + 32
)

// A DecapsulationKey768 is the secret key used to decapsulate a shared key from a
// ciphertext. It includes various precomputed values.
type DecapsulationKey768 struct {
	d [32]byte // decapsulation key seed
	z [32]byte // implicit rejection sampling seed

	ρ [32]byte // sampleNTT seed for A, stored for the encapsulation key
	h [32]byte // H(ek), stored for ML-KEM.Decaps_internal

	encryptionKey
	decryptionKey
}

// Bytes returns the decapsulation key as a 64-byte seed in the "d || z" form.
//
// The decapsulation key must be kept secret.
func ( *DecapsulationKey768) () []byte {
	var  [SeedSize]byte
	copy([:], .d[:])
	copy([32:], .z[:])
	return [:]
}

// EncapsulationKey returns the public encapsulation key necessary to produce
// ciphertexts.
func ( *DecapsulationKey768) () *EncapsulationKey768 {
	return &EncapsulationKey768{
		ρ:             .ρ,
		h:             .h,
		encryptionKey: .encryptionKey,
	}
}

// An EncapsulationKey768 is the public key used to produce ciphertexts to be
// decapsulated by the corresponding [DecapsulationKey768].
type EncapsulationKey768 struct {
	ρ [32]byte // sampleNTT seed for A
	h [32]byte // H(ek)
	encryptionKey
}

// Bytes returns the encapsulation key as a byte slice.
func ( *EncapsulationKey768) () []byte {
	// The actual logic is in a separate function to outline this allocation.
	 := make([]byte, 0, EncapsulationKeySize768)
	return .bytes()
}

func ( *EncapsulationKey768) ( []byte) []byte {
	for  := range .t {
		 = polyByteEncode(, .t[])
	}
	 = append(, .ρ[:]...)
	return 
}

// encryptionKey is the parsed and expanded form of a PKE encryption key.
type encryptionKey struct {
	t [k]nttElement     // ByteDecode₁₂(ek[:384k])
	a [k * k]nttElement // A[i*k+j] = sampleNTT(ρ, j, i)
}

// decryptionKey is the parsed and expanded form of a PKE decryption key.
type decryptionKey struct {
	s [k]nttElement // ByteDecode₁₂(dk[:decryptionKeySize])
}

// GenerateKey768 generates a new decapsulation key, drawing random bytes from
// a DRBG. The decapsulation key must be kept secret.
func () (*DecapsulationKey768, error) {
	// The actual logic is in a separate function to outline this allocation.
	 := &DecapsulationKey768{}
	return generateKey()
}

func generateKey( *DecapsulationKey768) (*DecapsulationKey768, error) {
	var  [32]byte
	drbg.Read([:])
	var  [32]byte
	drbg.Read([:])
	kemKeyGen(, &, &)
	if  := fips140.PCT("ML-KEM PCT", func() error { return kemPCT() });  != nil {
		// This clearly can't happen, but FIPS 140-3 requires us to check.
		panic()
	}
	fips140.RecordApproved()
	return , nil
}

// GenerateKeyInternal768 is a derandomized version of GenerateKey768,
// exclusively for use in tests.
func (,  *[32]byte) *DecapsulationKey768 {
	 := &DecapsulationKey768{}
	kemKeyGen(, , )
	return 
}

// NewDecapsulationKey768 parses a decapsulation key from a 64-byte
// seed in the "d || z" form. The seed must be uniformly random.
func ( []byte) (*DecapsulationKey768, error) {
	// The actual logic is in a separate function to outline this allocation.
	 := &DecapsulationKey768{}
	return newKeyFromSeed(, )
}

func newKeyFromSeed( *DecapsulationKey768,  []byte) (*DecapsulationKey768, error) {
	if len() != SeedSize {
		return nil, errors.New("mlkem: invalid seed length")
	}
	 := (*[32]byte)([:32])
	 := (*[32]byte)([32:])
	kemKeyGen(, , )
	if  := fips140.PCT("ML-KEM PCT", func() error { return kemPCT() });  != nil {
		// This clearly can't happen, but FIPS 140-3 requires us to check.
		panic()
	}
	fips140.RecordApproved()
	return , nil
}

// kemKeyGen generates a decapsulation key.
//
// It implements ML-KEM.KeyGen_internal according to FIPS 203, Algorithm 16, and
// K-PKE.KeyGen according to FIPS 203, Algorithm 13. The two are merged to save
// copies and allocations.
func kemKeyGen( *DecapsulationKey768, ,  *[32]byte) {
	.d = *
	.z = *

	 := sha3.New512()
	.Write([:])
	.Write([]byte{k}) // Module dimension as a domain separator.
	 := .Sum(make([]byte, 0, 64))
	,  := [:32], [32:]
	.ρ = [32]byte()

	 := &.a
	for  := byte(0);  < k; ++ {
		for  := byte(0);  < k; ++ {
			[*k+] = sampleNTT(, , )
		}
	}

	var  byte
	 := &.s
	for  := range  {
		[] = ntt(samplePolyCBD(, ))
		++
	}
	 := make([]nttElement, k)
	for  := range  {
		[] = ntt(samplePolyCBD(, ))
		++
	}

	 := &.t
	for  := range  { // t = A ◦ s + e
		[] = []
		for  := range  {
			[] = polyAdd([], nttMul([*k+], []))
		}
	}

	 := sha3.New256()
	 := .EncapsulationKey().Bytes()
	.Write()
	.Sum(.h[:0])
}

// kemPCT performs a Pairwise Consistency Test per FIPS 140-3 IG 10.3.A
// Additional Comment 1: "For key pairs generated for use with approved KEMs in
// FIPS 203, the PCT shall consist of applying the encapsulation key ek to
// encapsulate a shared secret K leading to ciphertext c, and then applying
// decapsulation key dk to retrieve the same shared secret K. The PCT passes if
// the two shared secret K values are equal. The PCT shall be performed either
// when keys are generated/imported, prior to the first exportation, or prior to
// the first operational use (if not exported before the first use)."
func kemPCT( *DecapsulationKey768) error {
	 := .EncapsulationKey()
	,  := .Encapsulate()
	,  := .Decapsulate()
	if  != nil {
		return 
	}
	if subtle.ConstantTimeCompare(, ) != 1 {
		return errors.New("mlkem: PCT failed")
	}
	return nil
}

// Encapsulate generates a shared key and an associated ciphertext from an
// encapsulation key, drawing random bytes from a DRBG.
//
// The shared key must be kept secret.
func ( *EncapsulationKey768) () (,  []byte) {
	// The actual logic is in a separate function to outline this allocation.
	var  [CiphertextSize768]byte
	return .encapsulate(&)
}

func ( *EncapsulationKey768) ( *[CiphertextSize768]byte) (,  []byte) {
	var  [messageSize]byte
	drbg.Read([:])
	// Note that the modulus check (step 2 of the encapsulation key check from
	// FIPS 203, Section 7.2) is performed by polyByteDecode in parseEK.
	fips140.RecordApproved()
	return kemEncaps(, , &)
}

// EncapsulateInternal is a derandomized version of Encapsulate, exclusively for
// use in tests.
func ( *EncapsulationKey768) ( *[32]byte) (,  []byte) {
	 := &[CiphertextSize768]byte{}
	return kemEncaps(, , )
}

// kemEncaps generates a shared key and an associated ciphertext.
//
// It implements ML-KEM.Encaps_internal according to FIPS 203, Algorithm 17.
func kemEncaps( *[CiphertextSize768]byte,  *EncapsulationKey768,  *[messageSize]byte) (,  []byte) {
	 := sha3.New512()
	.Write([:])
	.Write(.h[:])
	 := .Sum(nil)
	,  := [:SharedKeySize], [SharedKeySize:]
	 = pkeEncrypt(, &.encryptionKey, , )
	return , 
}

// NewEncapsulationKey768 parses an encapsulation key from its encoded form.
// If the encapsulation key is not valid, NewEncapsulationKey768 returns an error.
func ( []byte) (*EncapsulationKey768, error) {
	// The actual logic is in a separate function to outline this allocation.
	 := &EncapsulationKey768{}
	return parseEK(, )
}

// parseEK parses an encryption key from its encoded form.
//
// It implements the initial stages of K-PKE.Encrypt according to FIPS 203,
// Algorithm 14.
func parseEK( *EncapsulationKey768,  []byte) (*EncapsulationKey768, error) {
	if len() != EncapsulationKeySize768 {
		return nil, errors.New("mlkem: invalid encapsulation key length")
	}

	 := sha3.New256()
	.Write()
	.Sum(.h[:0])

	for  := range .t {
		var  error
		.t[],  = polyByteDecode[nttElement]([:encodingSize12])
		if  != nil {
			return nil, 
		}
		 = [encodingSize12:]
	}
	copy(.ρ[:], )

	for  := byte(0);  < k; ++ {
		for  := byte(0);  < k; ++ {
			.a[*k+] = sampleNTT(.ρ[:], , )
		}
	}

	return , nil
}

// pkeEncrypt encrypt a plaintext message.
//
// It implements K-PKE.Encrypt according to FIPS 203, Algorithm 14, although the
// computation of t and AT is done in parseEK.
func pkeEncrypt( *[CiphertextSize768]byte,  *encryptionKey,  *[messageSize]byte,  []byte) []byte {
	var  byte
	,  := make([]nttElement, k), make([]ringElement, k)
	for  := range  {
		[] = ntt(samplePolyCBD(, ))
		++
	}
	for  := range  {
		[] = samplePolyCBD(, )
		++
	}
	 := samplePolyCBD(, )

	 := make([]ringElement, k) // NTT⁻¹(AT ◦ r) + e1
	for  := range  {
		[] = []
		for  := range  {
			// Note that i and j are inverted, as we need the transposed of A.
			[] = polyAdd([], inverseNTT(nttMul(.a[*k+], [])))
		}
	}

	 := ringDecodeAndDecompress1()

	var  nttElement // t⊺ ◦ r
	for  := range .t {
		 = polyAdd(, nttMul(.t[], []))
	}
	 := polyAdd(polyAdd(inverseNTT(), ), )

	 := [:0]
	for ,  := range  {
		 = ringCompressAndEncode10(, )
	}
	 = ringCompressAndEncode4(, )

	return 
}

// Decapsulate generates a shared key from a ciphertext and a decapsulation key.
// If the ciphertext is not valid, Decapsulate returns an error.
//
// The shared key must be kept secret.
func ( *DecapsulationKey768) ( []byte) ( []byte,  error) {
	if len() != CiphertextSize768 {
		return nil, errors.New("mlkem: invalid ciphertext length")
	}
	 := (*[CiphertextSize768]byte)()
	// Note that the hash check (step 3 of the decapsulation input check from
	// FIPS 203, Section 7.3) is foregone as a DecapsulationKey is always
	// validly generated by ML-KEM.KeyGen_internal.
	return kemDecaps(, ), nil
}

// kemDecaps produces a shared key from a ciphertext.
//
// It implements ML-KEM.Decaps_internal according to FIPS 203, Algorithm 18.
func kemDecaps( *DecapsulationKey768,  *[CiphertextSize768]byte) ( []byte) {
	fips140.RecordApproved()
	 := pkeDecrypt(&.decryptionKey, )
	 := sha3.New512()
	.Write([:])
	.Write(.h[:])
	 := .Sum(make([]byte, 0, 64))
	,  := [:SharedKeySize], [SharedKeySize:]
	 := sha3.NewShake256()
	.Write(.z[:])
	.Write([:])
	 := make([]byte, SharedKeySize)
	.Read()
	var  [CiphertextSize768]byte
	 := pkeEncrypt(&, &.encryptionKey, (*[32]byte)(), )

	subtle.ConstantTimeCopy(subtle.ConstantTimeCompare([:], ), , )
	return 
}

// pkeDecrypt decrypts a ciphertext.
//
// It implements K-PKE.Decrypt according to FIPS 203, Algorithm 15,
// although s is retained from kemKeyGen.
func pkeDecrypt( *decryptionKey,  *[CiphertextSize768]byte) []byte {
	 := make([]ringElement, k)
	for  := range  {
		 := (*[encodingSize10]byte)([encodingSize10* : encodingSize10*(+1)])
		[] = ringDecodeAndDecompress10()
	}

	 := (*[encodingSize4]byte)([encodingSize10*k:])
	 := ringDecodeAndDecompress4()

	var  nttElement // s⊺ ◦ NTT(u)
	for  := range .s {
		 = polyAdd(, nttMul(.s[], ntt([])))
	}
	 := polySub(, inverseNTT())

	return ringCompressAndEncode1(nil, )
}