crypto/sign/mldsa.go

// Package sign provides post-quantum signature algorithms
package sign

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

// SigID identifies a signature algorithm
type SigID string

const (
	MLDSA2 SigID = "mldsa2"
	MLDSA3 SigID = "mldsa3"
	SLHDSA SigID = "slhdsa"
)

// Signer interface for signature algorithms
type Signer interface {
	// GenerateKeyPair generates a new signing key pair
	GenerateKeyPair() (PublicKey, PrivateKey, error)

	// Sign creates a signature
	Sign(sk PrivateKey, message []byte) ([]byte, error)

	// Verify verifies a signature
	Verify(pk PublicKey, message, signature []byte) bool

	// PublicKeySize returns the size of public keys
	PublicKeySize() int

	// PrivateKeySize returns the size of private keys
	PrivateKeySize() int

	// SignatureSize returns the size of signatures
	SignatureSize() int
}

// PublicKey represents a signature public key
type PublicKey interface {
	Bytes() []byte
	Equal(PublicKey) bool
}

// PrivateKey represents a signature private key
type PrivateKey interface {
	Bytes() []byte
	Public() PublicKey
	Equal(PrivateKey) bool
}

// MLDSA2Impl implements ML-DSA-44 (Dilithium2)
type MLDSA2Impl struct{}

// NewMLDSA2 creates a new ML-DSA-44 instance
func NewMLDSA2() Signer {
	return &MLDSA2Impl{}
}

// MLDSA2PublicKey represents an ML-DSA-44 public key
type MLDSA2PublicKey struct {
	data []byte
}

// MLDSA2PrivateKey represents an ML-DSA-44 private key
type MLDSA2PrivateKey struct {
	data []byte
	pk   *MLDSA2PublicKey
}

// Constants for ML-DSA-44 (Dilithium2)
const (
	mldsa2PublicKeySize  = 1312
	mldsa2PrivateKeySize = 2528
	mldsa2SignatureSize  = 2420
)

// GenerateKeyPair generates a new ML-DSA-44 key pair
func (m *MLDSA2Impl) GenerateKeyPair() (PublicKey, PrivateKey, error) {
	// Placeholder for actual ML-DSA key generation
	// In production, this would use liboqs or native implementation

	pk := &MLDSA2PublicKey{
		data: make([]byte, mldsa2PublicKeySize),
	}
	sk := &MLDSA2PrivateKey{
		data: make([]byte, mldsa2PrivateKeySize),
		pk:   pk,
	}

	// Generate random key material (placeholder)
	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
}

// Sign creates a signature
func (m *MLDSA2Impl) Sign(sk PrivateKey, message []byte) ([]byte, error) {
	mldsaSK, ok := sk.(*MLDSA2PrivateKey)
	if !ok {
		return nil, errors.New("invalid private key type")
	}

	signature := make([]byte, mldsa2SignatureSize)

	// Placeholder for actual ML-DSA signing
	if _, err := rand.Read(signature); err != nil {
		return nil, err
	}

	// In production, this would perform actual ML-DSA signing
	_ = mldsaSK.data
	_ = message

	return signature, nil
}

// Verify verifies a signature
func (m *MLDSA2Impl) Verify(pk PublicKey, message, signature []byte) bool {
	mldsaPK, ok := pk.(*MLDSA2PublicKey)
	if !ok {
		return false
	}

	if len(signature) != mldsa2SignatureSize {
		return false
	}

	// Placeholder for actual ML-DSA verification
	// In production, this would perform actual verification
	_ = mldsaPK.data
	_ = message
	_ = signature

	// Placeholder: always return true for now
	return true
}

// Size methods for ML-DSA-44
func (m *MLDSA2Impl) PublicKeySize() int  { return mldsa2PublicKeySize }
func (m *MLDSA2Impl) PrivateKeySize() int { return mldsa2PrivateKeySize }
func (m *MLDSA2Impl) SignatureSize() int  { return mldsa2SignatureSize }

// PublicKey methods
func (pk *MLDSA2PublicKey) Bytes() []byte {
	return pk.data
}

func (pk *MLDSA2PublicKey) Equal(other PublicKey) bool {
	otherPK, ok := other.(*MLDSA2PublicKey)
	if !ok {
		return false
	}
	return subtle.ConstantTimeCompare(pk.data, otherPK.data) == 1
}

// PrivateKey methods
func (sk *MLDSA2PrivateKey) Bytes() []byte {
	return sk.data
}

func (sk *MLDSA2PrivateKey) Public() PublicKey {
	return sk.pk
}

func (sk *MLDSA2PrivateKey) Equal(other PrivateKey) bool {
	otherSK, ok := other.(*MLDSA2PrivateKey)
	if !ok {
		return false
	}
	return subtle.ConstantTimeCompare(sk.data, otherSK.data) == 1
}

// MLDSA3Impl implements ML-DSA-65 (Dilithium3)
type MLDSA3Impl struct{}

// NewMLDSA3 creates a new ML-DSA-65 instance
func NewMLDSA3() Signer {
	return &MLDSA3Impl{}
}

// Constants for ML-DSA-65 (Dilithium3)
const (
	mldsa3PublicKeySize  = 1952
	mldsa3PrivateKeySize = 4000
	mldsa3SignatureSize  = 3293
)

// GenerateKeyPair generates a new ML-DSA-65 key pair
func (m *MLDSA3Impl) GenerateKeyPair() (PublicKey, PrivateKey, error) {
	pk := &MLDSA2PublicKey{
		data: make([]byte, mldsa3PublicKeySize),
	}
	sk := &MLDSA2PrivateKey{
		data: make([]byte, mldsa3PrivateKeySize),
		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
}

// Sign creates a signature
func (m *MLDSA3Impl) Sign(sk PrivateKey, message []byte) ([]byte, error) {
	signature := make([]byte, mldsa3SignatureSize)

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

	return signature, nil
}

// Verify verifies a signature
func (m *MLDSA3Impl) Verify(pk PublicKey, message, signature []byte) bool {
	if len(signature) != mldsa3SignatureSize {
		return false
	}

	// Placeholder
	return true
}

// Size methods for ML-DSA-65
func (m *MLDSA3Impl) PublicKeySize() int  { return mldsa3PublicKeySize }
func (m *MLDSA3Impl) PrivateKeySize() int { return mldsa3PrivateKeySize }
func (m *MLDSA3Impl) SignatureSize() int  { return mldsa3SignatureSize }

// GetSigner returns a Signer implementation for the given ID
func GetSigner(id SigID) (Signer, error) {
	switch id {
	case MLDSA2:
		return NewMLDSA2(), nil
	case MLDSA3:
		return NewMLDSA3(), nil
	case SLHDSA:
		// SLH-DSA would require additional implementation
		return nil, errors.New("SLH-DSA not yet implemented")
	default:
		return nil, fmt.Errorf("unsupported signature algorithm: %s", id)
	}
}

// TranscriptSigner signs protocol transcripts
type TranscriptSigner struct {
	signer Signer
	sk     PrivateKey
}

// NewTranscriptSigner creates a new transcript signer
func NewTranscriptSigner(signer Signer, sk PrivateKey) *TranscriptSigner {
	return &TranscriptSigner{
		signer: signer,
		sk:     sk,
	}
}

// SignTranscript signs a protocol transcript
func (ts *TranscriptSigner) SignTranscript(transcript []byte) ([]byte, error) {
	// Add context string to prevent cross-protocol attacks
	context := []byte("QZMQ-Transcript-v1")
	message := append(context, transcript...)

	return ts.signer.Sign(ts.sk, message)
}

// VerifyTranscript verifies a transcript signature
func VerifyTranscript(signer Signer, pk PublicKey, transcript, signature []byte) bool {
	context := []byte("QZMQ-Transcript-v1")
	message := append(context, transcript...)

	return signer.Verify(pk, message, signature)
}