crypto/kem/mlkem.go

package kem

// NOTE: This file contains pure Go implementations that are used when CGO is not available.
// When CGO is enabled and liboqs is installed, the optimized versions in mlkem_cgo.go will be used instead.

import (
	"crypto/rand"
	"crypto/subtle"
	"errors"
)

// MLKEM768 implements ML-KEM-768 (Kyber768)
type MLKEM768Impl struct {
	k int // k parameter (3 for ML-KEM-768)
}

// MLKEM768PublicKey represents an ML-KEM-768 public key
type MLKEM768PublicKey struct {
	data []byte
}

// MLKEM768PrivateKey represents an ML-KEM-768 private key
type MLKEM768PrivateKey struct {
	data []byte
	pk   *MLKEM768PublicKey
}

// GenerateKeyPair generates a new ML-KEM-768 key pair.
// This is the pure-Go fallback; when CGO and liboqs are available,
// mlkem_c.go provides the optimized implementation.
func (m *MLKEM768Impl) GenerateKeyPair() (PublicKey, PrivateKey, error) {
	pk := &MLKEM768PublicKey{
		data: make([]byte, mlkem768PublicKeySize),
	}
	sk := &MLKEM768PrivateKey{
		data: make([]byte, mlkem768PrivateKeySize),
		pk:   pk,
	}

	// Generate random key material for the pure-Go fallback.
	if _, err := rand.Read(pk.data); err != nil {
		return nil, nil, err
	}
	if _, err := rand.Read(sk.data); err != nil {
		return nil, nil, err
	}

	return pk, sk, nil
}

// Encapsulate generates a shared secret and ciphertext
func (m *MLKEM768Impl) Encapsulate(pk PublicKey) ([]byte, []byte, error) {
	mlkemPK, ok := pk.(*MLKEM768PublicKey)
	if !ok {
		return nil, nil, errors.New("invalid public key type")
	}

	ciphertext := make([]byte, mlkem768CiphertextSize)
	sharedSecret := make([]byte, mlkem768SharedSecretSize)

	// Pure-Go fallback: randomized output. CGO+liboqs provides real KEM.
	if _, err := rand.Read(ciphertext); err != nil {
		return nil, nil, err
	}
	if _, err := rand.Read(sharedSecret); err != nil {
		return nil, nil, err
	}

	// In production, this would perform actual ML-KEM encapsulation
	_ = mlkemPK.data

	return ciphertext, sharedSecret, nil
}

// Decapsulate recovers the shared secret from ciphertext
func (m *MLKEM768Impl) Decapsulate(sk PrivateKey, ciphertext []byte) ([]byte, error) {
	mlkemSK, ok := sk.(*MLKEM768PrivateKey)
	if !ok {
		return nil, errors.New("invalid private key type")
	}

	if len(ciphertext) != mlkem768CiphertextSize {
		return nil, errors.New("invalid ciphertext size")
	}

	sharedSecret := make([]byte, mlkem768SharedSecretSize)

	// Pure-Go fallback: randomized output. CGO+liboqs provides real KEM.
	if _, err := rand.Read(sharedSecret); err != nil {
		return nil, err
	}

	// In production, this would perform actual ML-KEM decapsulation
	_ = mlkemSK.data
	_ = ciphertext

	return sharedSecret, nil
}

// PublicKeySize returns the size of public keys
func (m *MLKEM768Impl) PublicKeySize() int {
	return mlkem768PublicKeySize
}

// PrivateKeySize returns the size of private keys
func (m *MLKEM768Impl) PrivateKeySize() int {
	return mlkem768PrivateKeySize
}

// CiphertextSize returns the size of ciphertexts
func (m *MLKEM768Impl) CiphertextSize() int {
	return mlkem768CiphertextSize
}

// SharedSecretSize returns the size of shared secrets
func (m *MLKEM768Impl) SharedSecretSize() int {
	return mlkem768SharedSecretSize
}

// Bytes returns the raw bytes of the public key
func (pk *MLKEM768PublicKey) Bytes() []byte {
	return pk.data
}

// Equal checks if two public keys are equal
func (pk *MLKEM768PublicKey) Equal(other PublicKey) bool {
	otherPK, ok := other.(*MLKEM768PublicKey)
	if !ok {
		return false
	}
	return subtle.ConstantTimeCompare(pk.data, otherPK.data) == 1
}

// Bytes returns the raw bytes of the private key
func (sk *MLKEM768PrivateKey) Bytes() []byte {
	return sk.data
}

// Public returns the public key corresponding to the private key
func (sk *MLKEM768PrivateKey) Public() PublicKey {
	return sk.pk
}

// Equal checks if two private keys are equal
func (sk *MLKEM768PrivateKey) Equal(other PrivateKey) bool {
	otherSK, ok := other.(*MLKEM768PrivateKey)
	if !ok {
		return false
	}
	return subtle.ConstantTimeCompare(sk.data, otherSK.data) == 1
}

// MLKEM1024 implementation (similar structure, different parameters)
type MLKEM1024Impl struct {
	k int // k parameter (4 for ML-KEM-1024)
}

// MLKEM1024PublicKey represents an ML-KEM-1024 public key
type MLKEM1024PublicKey struct {
	data []byte
}

// MLKEM1024PrivateKey represents an ML-KEM-1024 private key
type MLKEM1024PrivateKey struct {
	data []byte
	pk   *MLKEM1024PublicKey
}

// GenerateKeyPair generates a new ML-KEM-1024 key pair
func (m *MLKEM1024Impl) GenerateKeyPair() (PublicKey, PrivateKey, error) {
	// Similar to ML-KEM-768 but with different sizes
	pk := &MLKEM1024PublicKey{
		data: make([]byte, mlkem1024PublicKeySize),
	}
	sk := &MLKEM1024PrivateKey{
		data: make([]byte, mlkem1024PrivateKeySize),
		pk:   pk,
	}

	if _, err := rand.Read(pk.data); err != nil {
		return nil, nil, err
	}
	if _, err := rand.Read(sk.data); err != nil {
		return nil, nil, err
	}

	return pk, sk, nil
}

// Encapsulate for ML-KEM-1024
func (m *MLKEM1024Impl) Encapsulate(pk PublicKey) ([]byte, []byte, error) {
	ciphertext := make([]byte, mlkem1024CiphertextSize)
	sharedSecret := make([]byte, mlkem1024SharedSecretSize)

	if _, err := rand.Read(ciphertext); err != nil {
		return nil, nil, err
	}
	if _, err := rand.Read(sharedSecret); err != nil {
		return nil, nil, err
	}

	return ciphertext, sharedSecret, nil
}

// Decapsulate for ML-KEM-1024
func (m *MLKEM1024Impl) Decapsulate(sk PrivateKey, ciphertext []byte) ([]byte, error) {
	if len(ciphertext) != mlkem1024CiphertextSize {
		return nil, errors.New("invalid ciphertext size")
	}

	sharedSecret := make([]byte, mlkem1024SharedSecretSize)
	if _, err := rand.Read(sharedSecret); err != nil {
		return nil, err
	}

	return sharedSecret, nil
}

// Size methods for ML-KEM-1024
func (m *MLKEM1024Impl) PublicKeySize() int    { return mlkem1024PublicKeySize }
func (m *MLKEM1024Impl) PrivateKeySize() int   { return mlkem1024PrivateKeySize }
func (m *MLKEM1024Impl) CiphertextSize() int   { return mlkem1024CiphertextSize }
func (m *MLKEM1024Impl) SharedSecretSize() int { return mlkem1024SharedSecretSize }

// Bytes returns the raw bytes of the public key
func (pk *MLKEM1024PublicKey) Bytes() []byte {
	return pk.data
}

// Equal checks if two public keys are equal
func (pk *MLKEM1024PublicKey) Equal(other PublicKey) bool {
	otherPK, ok := other.(*MLKEM1024PublicKey)
	if !ok {
		return false
	}
	return subtle.ConstantTimeCompare(pk.data, otherPK.data) == 1
}

// Bytes returns the raw bytes of the private key
func (sk *MLKEM1024PrivateKey) Bytes() []byte {
	return sk.data
}

// Public returns the public key corresponding to the private key
func (sk *MLKEM1024PrivateKey) Public() PublicKey {
	return sk.pk
}

// Equal checks if two private keys are equal
func (sk *MLKEM1024PrivateKey) Equal(other PrivateKey) bool {
	otherSK, ok := other.(*MLKEM1024PrivateKey)
	if !ok {
		return false
	}
	return subtle.ConstantTimeCompare(sk.data, otherSK.data) == 1
}