Add performance benchmarks

all_test.go

@@ -263,7 +263,7 @@ func testPrecompiles(t *testing.T) {
 func testCGOPerformance(t *testing.T) {
 	// This test is for comparing performance when CGO optimizations are available
 	// CGO implementations are opt-in only with CGO=1
-	
+
 	t.Run("ML-KEM Performance", func(t *testing.T) {
 		message := make([]byte, 32)
 		rand.Read(message)
@@ -276,7 +276,7 @@ func testCGOPerformance(t *testing.T) {
 			priv.PublicKey.Encapsulate(rand.Reader)
 		}
 		duration := time.Since(start)
-		
+
 		t.Logf("ML-KEM-768 Encapsulate (100 ops): %v", duration)
 	})
 
@@ -292,7 +292,7 @@ func testCGOPerformance(t *testing.T) {
 			priv.Sign(rand.Reader, message, nil)
 		}
 		duration := time.Since(start)
-		
+
 		t.Logf("ML-DSA-65 Sign (100 ops): %v", duration)
 	})
 
@@ -308,7 +308,7 @@ func testCGOPerformance(t *testing.T) {
 			priv.Sign(rand.Reader, message, nil)
 		}
 		duration := time.Since(start)
-		
+
 		t.Logf("SLH-DSA-128f Sign (10 ops): %v", duration)
 	})
 }

audit_test.go

@@ -0,0 +1,253 @@
+package crypto
+
+import (
+	"crypto/rand"
+	"fmt"
+	"testing"
+
+	"github.com/luxfi/crypto/mldsa"
+	"github.com/luxfi/crypto/mlkem"
+	"github.com/luxfi/crypto/slhdsa"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+// TestMLKEMEdgeCases tests edge cases and potential bugs
+func TestMLKEMEdgeCases(t *testing.T) {
+	t.Run("Invalid Mode", func(t *testing.T) {
+		_, err := mlkem.GenerateKeyPair(rand.Reader, mlkem.Mode(99))
+		assert.Error(t, err)
+	})
+
+	t.Run("Nil Random Source", func(t *testing.T) {
+		_, err := mlkem.GenerateKeyPair(nil, mlkem.MLKEM768)
+		assert.Error(t, err)
+	})
+
+	t.Run("Empty Ciphertext", func(t *testing.T) {
+		priv, _ := mlkem.GenerateKeyPair(rand.Reader, mlkem.MLKEM768)
+		_, err := priv.Decapsulate([]byte{})
+		assert.Error(t, err)
+	})
+
+	t.Run("Wrong Size Ciphertext", func(t *testing.T) {
+		priv, _ := mlkem.GenerateKeyPair(rand.Reader, mlkem.MLKEM768)
+		wrongCT := make([]byte, 100) // Wrong size
+		_, err := priv.Decapsulate(wrongCT)
+		assert.Error(t, err)
+	})
+
+	t.Run("Serialization Round Trip", func(t *testing.T) {
+		modes := []mlkem.Mode{mlkem.MLKEM512, mlkem.MLKEM768, mlkem.MLKEM1024}
+		for _, mode := range modes {
+			priv1, _ := mlkem.GenerateKeyPair(rand.Reader, mode)
+			
+			// Serialize
+			privBytes := priv1.Bytes()
+			pubBytes := priv1.PublicKey.Bytes()
+			
+			// Deserialize
+			priv2, err := mlkem.PrivateKeyFromBytes(privBytes, mode)
+			require.NoError(t, err)
+			pub2, err := mlkem.PublicKeyFromBytes(pubBytes, mode)
+			require.NoError(t, err)
+			
+			// Verify they work the same
+			result1, _ := priv1.PublicKey.Encapsulate(rand.Reader)
+			secret1, _ := priv1.Decapsulate(result1.Ciphertext)
+			
+			result2, _ := pub2.Encapsulate(rand.Reader)
+			secret2, _ := priv2.Decapsulate(result2.Ciphertext)
+			
+			// Both should produce valid shared secrets
+			assert.Len(t, secret1, 32)
+			assert.Len(t, secret2, 32)
+		}
+	})
+
+	t.Run("Deterministic Public Key", func(t *testing.T) {
+		// Same private key seed should generate same public key
+		privBytes := make([]byte, mlkem.MLKEM768PrivateKeySize)
+		copy(privBytes, []byte("deterministic seed for testing"))
+		
+		priv1, _ := mlkem.PrivateKeyFromBytes(privBytes, mlkem.MLKEM768)
+		priv2, _ := mlkem.PrivateKeyFromBytes(privBytes, mlkem.MLKEM768)
+		
+		assert.Equal(t, priv1.PublicKey.Bytes(), priv2.PublicKey.Bytes())
+	})
+}
+
+// TestMLDSAEdgeCases tests ML-DSA edge cases
+func TestMLDSAEdgeCases(t *testing.T) {
+	t.Run("Invalid Mode", func(t *testing.T) {
+		_, err := mldsa.GenerateKey(rand.Reader, mldsa.Mode(99))
+		assert.Error(t, err)
+	})
+
+	t.Run("Empty Message", func(t *testing.T) {
+		priv, _ := mldsa.GenerateKey(rand.Reader, mldsa.MLDSA65)
+		sig, err := priv.Sign(rand.Reader, []byte{}, nil)
+		require.NoError(t, err)
+		assert.True(t, priv.PublicKey.Verify([]byte{}, sig))
+	})
+
+	t.Run("Large Message", func(t *testing.T) {
+		priv, _ := mldsa.GenerateKey(rand.Reader, mldsa.MLDSA65)
+		largeMsg := make([]byte, 10000)
+		rand.Read(largeMsg)
+		
+		sig, err := priv.Sign(rand.Reader, largeMsg, nil)
+		require.NoError(t, err)
+		assert.True(t, priv.PublicKey.Verify(largeMsg, sig))
+	})
+
+	t.Run("Signature Malleability", func(t *testing.T) {
+		priv, _ := mldsa.GenerateKey(rand.Reader, mldsa.MLDSA65)
+		msg := []byte("test message")
+		
+		sig1, _ := priv.Sign(rand.Reader, msg, nil)
+		sig2, _ := priv.Sign(rand.Reader, msg, nil)
+		
+		// Signatures should be deterministic in our implementation
+		assert.Equal(t, sig1, sig2)
+	})
+
+	t.Run("Wrong Signature Size", func(t *testing.T) {
+		priv, _ := mldsa.GenerateKey(rand.Reader, mldsa.MLDSA65)
+		msg := []byte("test")
+		
+		wrongSig := make([]byte, 100) // Wrong size
+		assert.False(t, priv.PublicKey.Verify(msg, wrongSig))
+	})
+
+	t.Run("Cross Mode Verification", func(t *testing.T) {
+		priv44, _ := mldsa.GenerateKey(rand.Reader, mldsa.MLDSA44)
+		priv65, _ := mldsa.GenerateKey(rand.Reader, mldsa.MLDSA65)
+		msg := []byte("test")
+		
+		sig44, _ := priv44.Sign(rand.Reader, msg, nil)
+		
+		// ML-DSA65 key shouldn't verify ML-DSA44 signature
+		assert.False(t, priv65.PublicKey.Verify(msg, sig44))
+	})
+}
+
+// TestSLHDSAEdgeCases tests SLH-DSA edge cases
+func TestSLHDSAEdgeCases(t *testing.T) {
+	t.Run("Deterministic Signatures", func(t *testing.T) {
+		priv, _ := slhdsa.GenerateKey(rand.Reader, slhdsa.SLHDSA128f)
+		msg := []byte("deterministic test")
+		
+		sig1, _ := priv.Sign(rand.Reader, msg, nil)
+		sig2, _ := priv.Sign(rand.Reader, msg, nil)
+		
+		// SLH-DSA is deterministic - same message should produce same signature
+		assert.Equal(t, sig1, sig2)
+	})
+
+	t.Run("Large Signature Sizes", func(t *testing.T) {
+		modes := []struct {
+			mode slhdsa.Mode
+			name string
+			size int
+		}{
+			{slhdsa.SLHDSA128f, "128f", slhdsa.SLHDSA128fSignatureSize},
+			{slhdsa.SLHDSA192f, "192f", slhdsa.SLHDSA192fSignatureSize},
+		}
+		
+		for _, m := range modes {
+			t.Run(m.name, func(t *testing.T) {
+				priv, _ := slhdsa.GenerateKey(rand.Reader, m.mode)
+				msg := []byte("test")
+				sig, _ := priv.Sign(rand.Reader, msg, nil)
+				
+				assert.Len(t, sig, m.size)
+			})
+		}
+	})
+}
+
+// TestConcurrency tests thread safety
+func TestConcurrency(t *testing.T) {
+	t.Run("ML-KEM Concurrent Operations", func(t *testing.T) {
+		priv, _ := mlkem.GenerateKeyPair(rand.Reader, mlkem.MLKEM768)
+		
+		// Run concurrent encapsulations
+		done := make(chan bool, 10)
+		for i := 0; i < 10; i++ {
+			go func() {
+				result, err := priv.PublicKey.Encapsulate(rand.Reader)
+				assert.NoError(t, err)
+				secret, err := priv.Decapsulate(result.Ciphertext)
+				assert.NoError(t, err)
+				assert.Equal(t, result.SharedSecret, secret)
+				done <- true
+			}()
+		}
+		
+		for i := 0; i < 10; i++ {
+			<-done
+		}
+	})
+
+	t.Run("ML-DSA Concurrent Signing", func(t *testing.T) {
+		priv, _ := mldsa.GenerateKey(rand.Reader, mldsa.MLDSA65)
+		
+		done := make(chan bool, 10)
+		for i := 0; i < 10; i++ {
+			go func(id int) {
+				msg := []byte(fmt.Sprintf("message %d", id))
+				sig, err := priv.Sign(rand.Reader, msg, nil)
+				assert.NoError(t, err)
+				assert.True(t, priv.PublicKey.Verify(msg, sig))
+				done <- true
+			}(i)
+		}
+		
+		for i := 0; i < 10; i++ {
+			<-done
+		}
+	})
+}
+
+// TestMemoryLeaks checks for potential memory issues
+func TestMemoryLeaks(t *testing.T) {
+	t.Run("ML-KEM No Leak", func(t *testing.T) {
+		// This would need proper memory profiling
+		// For now, just ensure no panics on repeated operations
+		for i := 0; i < 100; i++ {
+			priv, _ := mlkem.GenerateKeyPair(rand.Reader, mlkem.MLKEM768)
+			result, _ := priv.PublicKey.Encapsulate(rand.Reader)
+			priv.Decapsulate(result.Ciphertext)
+		}
+	})
+}
+
+// TestParameterValidation ensures all parameters match NIST specs
+func TestParameterValidation(t *testing.T) {
+	// ML-KEM parameters from FIPS 203
+	assert.Equal(t, 800, mlkem.MLKEM512PublicKeySize)
+	assert.Equal(t, 1632, mlkem.MLKEM512PrivateKeySize)
+	assert.Equal(t, 768, mlkem.MLKEM512CiphertextSize)
+	
+	assert.Equal(t, 1184, mlkem.MLKEM768PublicKeySize)
+	assert.Equal(t, 2400, mlkem.MLKEM768PrivateKeySize)
+	assert.Equal(t, 1088, mlkem.MLKEM768CiphertextSize)
+	
+	assert.Equal(t, 1568, mlkem.MLKEM1024PublicKeySize)
+	assert.Equal(t, 3168, mlkem.MLKEM1024PrivateKeySize)
+	assert.Equal(t, 1568, mlkem.MLKEM1024CiphertextSize)
+
+	// ML-DSA parameters from FIPS 204
+	assert.Equal(t, 1312, mldsa.MLDSA44PublicKeySize)
+	assert.Equal(t, 2528, mldsa.MLDSA44PrivateKeySize)
+	assert.Equal(t, 2420, mldsa.MLDSA44SignatureSize)
+	
+	assert.Equal(t, 1952, mldsa.MLDSA65PublicKeySize)
+	assert.Equal(t, 4000, mldsa.MLDSA65PrivateKeySize)
+	assert.Equal(t, 3293, mldsa.MLDSA65SignatureSize)
+	
+	assert.Equal(t, 2592, mldsa.MLDSA87PublicKeySize)
+	assert.Equal(t, 4864, mldsa.MLDSA87PrivateKeySize)
+	assert.Equal(t, 4595, mldsa.MLDSA87SignatureSize)
+}

common/hash.go

@@ -0,0 +1,96 @@
+// Package common provides shared utilities for post-quantum crypto implementations
+package common
+
+import (
+	"crypto/sha256"
+	"crypto/sha512"
+	"hash"
+)
+
+// HashFunc represents available hash functions
+type HashFunc int
+
+const (
+	SHA256Hash HashFunc = iota
+	SHA512Hash
+	SHA3_256Hash
+	SHA3_512Hash
+)
+
+// GetHasher returns a hash function instance
+func GetHasher(h HashFunc) hash.Hash {
+	switch h {
+	case SHA512Hash:
+		return sha512.New()
+	default:
+		return sha256.New()
+	}
+}
+
+// DeriveKey derives deterministic key material from seed
+func DeriveKey(seed []byte, label string, outputLen int) []byte {
+	h := sha256.New()
+	h.Write(seed)
+	h.Write([]byte(label))
+	baseHash := h.Sum(nil)
+	
+	output := make([]byte, outputLen)
+	for i := 0; i < outputLen; i += len(baseHash) {
+		h.Reset()
+		h.Write(baseHash)
+		h.Write([]byte{byte(i / len(baseHash))})
+		chunk := h.Sum(nil)
+		
+		end := i + len(baseHash)
+		if end > outputLen {
+			end = outputLen
+		}
+		copy(output[i:end], chunk)
+		baseHash = chunk
+	}
+	
+	return output
+}
+
+// XOF extends output using a hash function as XOF
+func XOF(seed []byte, outputLen int) []byte {
+	output := make([]byte, outputLen)
+	h := sha256.New()
+	h.Write(seed)
+	hash := h.Sum(nil)
+	
+	for i := 0; i < outputLen; i += len(hash) {
+		end := i + len(hash)
+		if end > outputLen {
+			end = outputLen
+		}
+		copy(output[i:end], hash)
+		if end < outputLen {
+			h.Reset()
+			h.Write(hash)
+			hash = h.Sum(nil)
+		}
+	}
+	
+	return output
+}
+
+// SecureCompare performs constant-time comparison
+func SecureCompare(a, b []byte) bool {
+	if len(a) != len(b) {
+		return false
+	}
+	
+	var result byte
+	for i := range a {
+		result |= a[i] ^ b[i]
+	}
+	return result == 0
+}
+
+// ClearBytes securely clears sensitive data
+func ClearBytes(b []byte) {
+	for i := range b {
+		b[i] = 0
+	}
+}

common/utils.go

@@ -0,0 +1,139 @@
+// Package common provides shared utilities for post-quantum crypto implementations
+package common
+
+import (
+	"crypto/rand"
+	"errors"
+	"io"
+)
+
+// ValidateMode checks if a mode value is within valid range
+func ValidateMode(mode int, min, max int) error {
+	if mode < min || mode > max {
+		return errors.New("invalid mode")
+	}
+	return nil
+}
+
+// ValidateRandomSource checks if random source is not nil
+func ValidateRandomSource(rand io.Reader) error {
+	if rand == nil {
+		return errors.New("random source is nil")
+	}
+	return nil
+}
+
+// GenerateRandomBytes generates random bytes with validation
+func GenerateRandomBytes(rand io.Reader, size int) ([]byte, error) {
+	if err := ValidateRandomSource(rand); err != nil {
+		return nil, err
+	}
+	
+	bytes := make([]byte, size)
+	if _, err := io.ReadFull(rand, bytes); err != nil {
+		return nil, err
+	}
+	
+	return bytes, nil
+}
+
+// AllocateCombined allocates a single buffer for multiple data segments
+func AllocateCombined(sizes ...int) []byte {
+	total := 0
+	for _, size := range sizes {
+		total += size
+	}
+	return make([]byte, total)
+}
+
+// SplitBuffer splits a buffer into segments of specified sizes
+func SplitBuffer(buffer []byte, sizes ...int) [][]byte {
+	segments := make([][]byte, len(sizes))
+	offset := 0
+	
+	for i, size := range sizes {
+		if offset+size > len(buffer) {
+			panic("buffer too small for requested segments")
+		}
+		segments[i] = buffer[offset : offset+size]
+		offset += size
+	}
+	
+	return segments
+}
+
+// CopyWithPadding copies source to destination with padding if needed
+func CopyWithPadding(dst, src []byte, padValue byte) {
+	copy(dst, src)
+	if len(src) < len(dst) {
+		for i := len(src); i < len(dst); i++ {
+			dst[i] = padValue
+		}
+	}
+}
+
+// ConstantTimeSelect returns a if v == 1, b if v == 0
+func ConstantTimeSelect(v int, a, b []byte) []byte {
+	if len(a) != len(b) {
+		panic("slices must have equal length")
+	}
+	
+	result := make([]byte, len(a))
+	for i := range result {
+		result[i] = byte(v)*a[i] + byte(1-v)*b[i]
+	}
+	return result
+}
+
+// ValidateBufferSize checks if buffer has expected size
+func ValidateBufferSize(buf []byte, expectedSize int, name string) error {
+	if len(buf) != expectedSize {
+		return errors.New(name + " has invalid size")
+	}
+	return nil
+}
+
+// SafeCopy performs bounds-checked copy
+func SafeCopy(dst, src []byte) error {
+	if len(dst) < len(src) {
+		return errors.New("destination buffer too small")
+	}
+	copy(dst, src)
+	return nil
+}
+
+// GenerateDeterministicBytes generates deterministic bytes from seed
+func GenerateDeterministicBytes(seed []byte, length int) []byte {
+	return DeriveKey(seed, "deterministic", length)
+}
+
+// CreateSignatureBuffer creates a signature buffer with proper initialization
+func CreateSignatureBuffer(size int) []byte {
+	return make([]byte, size)
+}
+
+// FillRandomBytes fills existing buffer with random bytes
+func FillRandomBytes(rand io.Reader, buf []byte) error {
+	if err := ValidateRandomSource(rand); err != nil {
+		return err
+	}
+	
+	_, err := io.ReadFull(rand, buf)
+	return err
+}
+
+// Min returns the minimum of two integers
+func Min(a, b int) int {
+	if a < b {
+		return a
+	}
+	return b
+}
+
+// Max returns the maximum of two integers
+func Max(a, b int) int {
+	if a > b {
+		return a
+	}
+	return b
+}

mldsa/mldsa.go

@@ -68,19 +68,24 @@ func GenerateKey(rand io.Reader, mode Mode) (*PrivateKey, error) {
 		return nil, errors.New("invalid ML-DSA mode")
 	}
 
+	// Check for nil random source
+	if rand == nil {
+		return nil, errors.New("random source is nil")
+	}
+
 	// Placeholder implementation - generate random private key
 	privBytes := make([]byte, privKeySize)
 	if _, err := io.ReadFull(rand, privBytes); err != nil {
 		return nil, err
 	}
-	
+
 	// Derive public key from private key for consistency
 	// In real ML-DSA, public key is derived from private key seed
 	h := sha256.New()
 	h.Write(privBytes[:32]) // Use first 32 bytes as seed
 	h.Write([]byte("public"))
 	pubSeed := h.Sum(nil)
-	
+
 	pubBytes := make([]byte, pubKeySize)
 	// Fill public key with deterministic data
 	for i := 0; i < pubKeySize; i += 32 {
@@ -129,13 +134,13 @@ func (priv *PrivateKey) Sign(rand io.Reader, message []byte, opts crypto.SignerO
 	signature := make([]byte, sigSize)
 	// Copy the hash to the beginning of signature
 	copy(signature[:32], hash)
-	
+
 	// Fill rest with deterministic data based on private key
 	h.Reset()
 	h.Write(priv.data)
 	h.Write(message)
 	privHash := h.Sum(nil)
-	
+
 	for i := 32; i < sigSize; i += len(privHash) {
 		end := i + len(privHash)
 		if end > sigSize {
@@ -175,20 +180,20 @@ func (pub *PublicKey) Verify(message, signature []byte) bool {
 	h.Write(pub.data)
 	h.Write(message)
 	expectedSigStart := h.Sum(nil)
-	
+
 	// Check if first 32 bytes of signature match expected
 	// This is a simplified check for our placeholder
 	if len(signature) < 32 {
 		return false
 	}
-	
+
 	// Compare first 32 bytes
 	for i := 0; i < 32; i++ {
 		if signature[i] != expectedSigStart[i] {
 			return false
 		}
 	}
-	
+
 	return true
 }
 
@@ -262,7 +267,7 @@ func PrivateKeyFromBytes(data []byte, mode Mode) (*PrivateKey, error) {
 	h.Write(privData[:32]) // Use first 32 bytes as seed
 	h.Write([]byte("public"))
 	pubSeed := h.Sum(nil)
-	
+
 	pubData := make([]byte, expectedPubSize)
 	// Fill public key with deterministic data
 	for i := 0; i < expectedPubSize; i += 32 {

mldsa/mldsa_bench_test.go

@@ -0,0 +1,149 @@
+package mldsa
+
+import (
+	"crypto/rand"
+	"fmt"
+	"testing"
+)
+
+func BenchmarkMLDSA44(b *testing.B) {
+	benchmarkMLDSA(b, MLDSA44)
+}
+
+func BenchmarkMLDSA65(b *testing.B) {
+	benchmarkMLDSA(b, MLDSA65)
+}
+
+func BenchmarkMLDSA87(b *testing.B) {
+	benchmarkMLDSA(b, MLDSA87)
+}
+
+func benchmarkMLDSA(b *testing.B, mode Mode) {
+	message := make([]byte, 32)
+	rand.Read(message)
+
+	b.Run("GenerateKey", func(b *testing.B) {
+		b.ReportAllocs()
+		for i := 0; i < b.N; i++ {
+			_, err := GenerateKey(rand.Reader, mode)
+			if err != nil {
+				b.Fatal(err)
+			}
+		}
+	})
+
+	priv, err := GenerateKey(rand.Reader, mode)
+	if err != nil {
+		b.Fatal(err)
+	}
+
+	b.Run("Sign", func(b *testing.B) {
+		b.ReportAllocs()
+		for i := 0; i < b.N; i++ {
+			_, err := priv.Sign(rand.Reader, message, nil)
+			if err != nil {
+				b.Fatal(err)
+			}
+		}
+	})
+
+	sig, err := priv.Sign(rand.Reader, message, nil)
+	if err != nil {
+		b.Fatal(err)
+	}
+
+	b.Run("Verify", func(b *testing.B) {
+		b.ReportAllocs()
+		for i := 0; i < b.N; i++ {
+			valid := priv.PublicKey.Verify(message, sig)
+			if !valid {
+				b.Fatal("verification failed")
+			}
+		}
+	})
+
+	b.Run("Serialize", func(b *testing.B) {
+		b.ReportAllocs()
+		for i := 0; i < b.N; i++ {
+			_ = priv.Bytes()
+			_ = priv.PublicKey.Bytes()
+		}
+	})
+
+	privBytes := priv.Bytes()
+	pubBytes := priv.PublicKey.Bytes()
+
+	b.Run("Deserialize", func(b *testing.B) {
+		b.ReportAllocs()
+		for i := 0; i < b.N; i++ {
+			_, err := PrivateKeyFromBytes(privBytes, mode)
+			if err != nil {
+				b.Fatal(err)
+			}
+			_, err = PublicKeyFromBytes(pubBytes, mode)
+			if err != nil {
+				b.Fatal(err)
+			}
+		}
+	})
+}
+
+// Batch verification benchmark
+func BenchmarkMLDSABatchVerify(b *testing.B) {
+	modes := []struct {
+		name string
+		mode Mode
+	}{
+		{"MLDSA44", MLDSA44},
+		{"MLDSA65", MLDSA65},
+		{"MLDSA87", MLDSA87},
+	}
+
+	for _, m := range modes {
+		b.Run(m.name, func(b *testing.B) {
+			// Generate multiple signatures
+			numSigs := 10
+			messages := make([][]byte, numSigs)
+			signatures := make([][]byte, numSigs)
+			
+			priv, _ := GenerateKey(rand.Reader, m.mode)
+			for i := 0; i < numSigs; i++ {
+				messages[i] = make([]byte, 32)
+				rand.Read(messages[i])
+				signatures[i], _ = priv.Sign(rand.Reader, messages[i], nil)
+			}
+
+			b.ReportAllocs()
+			b.ResetTimer()
+			for i := 0; i < b.N; i++ {
+				// Verify all signatures
+				for j := 0; j < numSigs; j++ {
+					if !priv.PublicKey.Verify(messages[j], signatures[j]) {
+						b.Fatal("verification failed")
+					}
+				}
+			}
+		})
+	}
+}
+
+// Different message sizes
+func BenchmarkMLDSAMessageSizes(b *testing.B) {
+	sizes := []int{32, 64, 128, 256, 512, 1024}
+	priv, _ := GenerateKey(rand.Reader, MLDSA65)
+
+	for _, size := range sizes {
+		b.Run(fmt.Sprintf("Size%d", size), func(b *testing.B) {
+			message := make([]byte, size)
+			rand.Read(message)
+			
+			b.ReportAllocs()
+			for i := 0; i < b.N; i++ {
+				sig, _ := priv.Sign(rand.Reader, message, nil)
+				if !priv.PublicKey.Verify(message, sig) {
+					b.Fatal("verification failed")
+				}
+			}
+		})
+	}
+}

mldsa/mldsa_cgo.go

@@ -8,7 +8,7 @@ package mldsa
 //
 // TODO: Implement actual CGO optimizations using C libraries for:
 // - ML-DSA-44/65/87 from NIST reference implementation
-// - CRYSTALS-Dilithium optimized implementations  
+// - CRYSTALS-Dilithium optimized implementations
 // - AVX2/AVX512 optimizations for x86_64
 // - NEON optimizations for ARM64
-// - Batch verification optimizations
+// - Batch verification optimizations

mldsa/mldsa_optimized.go

@@ -0,0 +1,298 @@
+package mldsa
+
+import (
+	"crypto"
+	"crypto/rand"
+	"crypto/sha256"
+	"errors"
+	"io"
+	"sync"
+)
+
+// Pool for reusing signature buffers
+var sigBufferPool = sync.Pool{
+	New: func() interface{} {
+		return make([]byte, MLDSA87SignatureSize) // Max size
+	},
+}
+
+// Pool for reusing hash buffers
+var hashBufferPool = sync.Pool{
+	New: func() interface{} {
+		return make([]byte, 64) // SHA-512 output size
+	},
+}
+
+// getSignatureBuffer gets a buffer from the pool
+func getSignatureBuffer(size int) []byte {
+	buf := sigBufferPool.Get().([]byte)
+	if cap(buf) < size {
+		return make([]byte, size)
+	}
+	return buf[:size]
+}
+
+// putSignatureBuffer returns a buffer to the pool
+func putSignatureBuffer(buf []byte) {
+	if cap(buf) >= MLDSA44SignatureSize { // Only pool larger buffers
+		sigBufferPool.Put(buf)
+	}
+}
+
+// Optimized key generation with pre-allocated buffers
+func GenerateKeyOptimized(rand io.Reader, mode Mode) (*PrivateKey, error) {
+	var pubKeySize, privKeySize int
+
+	switch mode {
+	case MLDSA44:
+		pubKeySize = MLDSA44PublicKeySize
+		privKeySize = MLDSA44PrivateKeySize
+	case MLDSA65:
+		pubKeySize = MLDSA65PublicKeySize
+		privKeySize = MLDSA65PrivateKeySize
+	case MLDSA87:
+		pubKeySize = MLDSA87PublicKeySize
+		privKeySize = MLDSA87PrivateKeySize
+	default:
+		return nil, errors.New("invalid ML-DSA mode")
+	}
+
+	// Check for nil random source
+	if rand == nil {
+		return nil, errors.New("random source is nil")
+	}
+
+	// Use single allocation for both keys
+	totalSize := privKeySize + pubKeySize
+	allBytes := make([]byte, totalSize)
+	
+	privBytes := allBytes[:privKeySize]
+	if _, err := io.ReadFull(rand, privBytes); err != nil {
+		return nil, err
+	}
+	
+	// Derive public key in-place
+	pubBytes := allBytes[privKeySize:]
+	derivePublicKeyOptimized(privBytes[:32], pubBytes)
+
+	return &PrivateKey{
+		PublicKey: &PublicKey{
+			mode: mode,
+			data: pubBytes,
+		},
+		data: privBytes,
+	}, nil
+}
+
+// Optimized public key derivation
+func derivePublicKeyOptimized(seed []byte, output []byte) {
+	h := sha256.New()
+	h.Write(seed)
+	h.Write([]byte("public"))
+	hash := h.Sum(nil)
+	
+	// Unroll loop for better performance
+	outputLen := len(output)
+	for i := 0; i < outputLen; {
+		remaining := outputLen - i
+		if remaining >= 32 {
+			copy(output[i:i+32], hash)
+			i += 32
+		} else {
+			copy(output[i:], hash[:remaining])
+			break
+		}
+		
+		if i < outputLen {
+			h.Reset()
+			h.Write(hash)
+			h.Write([]byte{byte(i / 32)})
+			hash = h.Sum(hash[:0]) // Reuse slice
+		}
+	}
+}
+
+// OptimizedSign performs signing with buffer pooling
+func (priv *PrivateKey) OptimizedSign(rand io.Reader, message []byte, opts crypto.SignerOpts) ([]byte, error) {
+	var sigSize int
+	
+	switch priv.PublicKey.mode {
+	case MLDSA44:
+		sigSize = MLDSA44SignatureSize
+	case MLDSA65:
+		sigSize = MLDSA65SignatureSize
+	case MLDSA87:
+		sigSize = MLDSA87SignatureSize
+	default:
+		return nil, errors.New("invalid ML-DSA mode")
+	}
+
+	// Get buffer from pool
+	signature := getSignatureBuffer(sigSize)
+	defer putSignatureBuffer(signature)
+	
+	// Create hash with pooled buffer
+	hashBuf := hashBufferPool.Get().([]byte)
+	defer hashBufferPool.Put(hashBuf)
+	
+	h := sha256.New()
+	h.Write(priv.PublicKey.data)
+	h.Write(message)
+	copy(hashBuf, h.Sum(nil))
+	
+	// Copy hash to signature
+	copy(signature[:32], hashBuf)
+	
+	// Fill rest with deterministic data
+	h.Reset()
+	h.Write(priv.data[:32])
+	h.Write(message)
+	deterministicBytes := h.Sum(nil)
+	
+	for i := 32; i < sigSize; i += 32 {
+		end := i + 32
+		if end > sigSize {
+			end = sigSize
+		}
+		copy(signature[i:end], deterministicBytes)
+		if end < sigSize {
+			h.Reset()
+			h.Write(deterministicBytes)
+			h.Write([]byte{byte(i / 32)})
+			deterministicBytes = h.Sum(deterministicBytes[:0])
+		}
+	}
+	
+	// Return a copy since we're returning the pooled buffer
+	result := make([]byte, sigSize)
+	copy(result, signature)
+	return result, nil
+}
+
+// BatchDSA provides batch signing operations
+type BatchDSA struct {
+	mode Mode
+	keys []*PrivateKey
+	mu   sync.Mutex
+}
+
+// NewBatchDSA creates a batch DSA processor
+func NewBatchDSA(mode Mode, numKeys int) (*BatchDSA, error) {
+	keys := make([]*PrivateKey, numKeys)
+	for i := range keys {
+		key, err := GenerateKey(rand.Reader, mode)
+		if err != nil {
+			return nil, err
+		}
+		keys[i] = key
+	}
+	
+	return &BatchDSA{
+		mode: mode,
+		keys: keys,
+	}, nil
+}
+
+// SignBatch performs batch signing in parallel
+func (b *BatchDSA) SignBatch(messages [][]byte) ([][]byte, error) {
+	if len(messages) != len(b.keys) {
+		return nil, errors.New("message count mismatch")
+	}
+	
+	signatures := make([][]byte, len(messages))
+	errors := make([]error, len(messages))
+	
+	// Process in parallel
+	var wg sync.WaitGroup
+	for i := range messages {
+		wg.Add(1)
+		go func(idx int) {
+			defer wg.Done()
+			sig, err := b.keys[idx].Sign(rand.Reader, messages[idx], nil)
+			if err != nil {
+				errors[idx] = err
+				return
+			}
+			signatures[idx] = sig
+		}(i)
+	}
+	
+	wg.Wait()
+	
+	// Check for errors
+	for _, err := range errors {
+		if err != nil {
+			return nil, err
+		}
+	}
+	
+	return signatures, nil
+}
+
+// VerifyBatch performs batch verification in parallel
+func (b *BatchDSA) VerifyBatch(messages [][]byte, signatures [][]byte) ([]bool, error) {
+	if len(messages) != len(signatures) || len(messages) != len(b.keys) {
+		return nil, errors.New("input count mismatch")
+	}
+	
+	results := make([]bool, len(messages))
+	
+	// Process in parallel
+	var wg sync.WaitGroup
+	for i := range messages {
+		wg.Add(1)
+		go func(idx int) {
+			defer wg.Done()
+			results[idx] = b.keys[idx].PublicKey.Verify(messages[idx], signatures[idx])
+		}(i)
+	}
+	
+	wg.Wait()
+	
+	return results, nil
+}
+
+// PrecomputedMLDSA stores precomputed values for faster operations
+type PrecomputedMLDSA struct {
+	mode        Mode
+	privKey     *PrivateKey
+	hashCache   map[string][]byte
+	mu          sync.RWMutex
+}
+
+// NewPrecomputedMLDSA creates a new precomputed ML-DSA instance
+func NewPrecomputedMLDSA(privKey *PrivateKey) *PrecomputedMLDSA {
+	return &PrecomputedMLDSA{
+		mode:      privKey.PublicKey.mode,
+		privKey:   privKey,
+		hashCache: make(map[string][]byte),
+	}
+}
+
+// SignCached signs with caching for repeated messages
+func (p *PrecomputedMLDSA) SignCached(message []byte) ([]byte, error) {
+	// Create cache key
+	h := sha256.New()
+	h.Write(message)
+	cacheKey := string(h.Sum(nil))
+	
+	// Check cache
+	p.mu.RLock()
+	if sig, ok := p.hashCache[cacheKey]; ok {
+		p.mu.RUnlock()
+		return sig, nil
+	}
+	p.mu.RUnlock()
+	
+	// Sign and cache
+	sig, err := p.privKey.Sign(rand.Reader, message, nil)
+	if err != nil {
+		return nil, err
+	}
+	
+	p.mu.Lock()
+	p.hashCache[cacheKey] = sig
+	p.mu.Unlock()
+	
+	return sig, nil
+}

mldsa/mldsa_test.go

@@ -26,10 +26,3 @@ func TestMLDSA(t *testing.T) {
 		}
 	})
 }
-
-func BenchmarkMLDSA65(b *testing.B) {
-	for i := 0; i < b.N; i++ {
-		// Placeholder benchmark
-		_ = MLDSA65
-	}
-}

mlkem/mlkem.go

@@ -76,6 +76,11 @@ func GenerateKeyPair(rand io.Reader, mode Mode) (*PrivateKey, error) {
 		return nil, errors.New("invalid ML-KEM mode")
 	}
 
+	// Check for nil random source
+	if rand == nil {
+		return nil, errors.New("random source is nil")
+	}
+	
 	// Placeholder implementation - generate random private key
 	privBytes := make([]byte, privKeySize)
 	if _, err := io.ReadFull(rand, privBytes); err != nil {
@@ -264,4 +269,3 @@ func PrivateKeyFromBytes(data []byte, mode Mode) (*PrivateKey, error) {
 		data: data,
 	}, nil
 }
-

mlkem/mlkem_bench_test.go

@@ -0,0 +1,112 @@
+package mlkem
+
+import (
+	"crypto/rand"
+	"testing"
+)
+
+func BenchmarkMLKEM512(b *testing.B) {
+	benchmarkMLKEM(b, MLKEM512)
+}
+
+func BenchmarkMLKEM768(b *testing.B) {
+	benchmarkMLKEM(b, MLKEM768)
+}
+
+func BenchmarkMLKEM1024(b *testing.B) {
+	benchmarkMLKEM(b, MLKEM1024)
+}
+
+func benchmarkMLKEM(b *testing.B, mode Mode) {
+	b.Run("GenerateKeyPair", func(b *testing.B) {
+		b.ReportAllocs()
+		for i := 0; i < b.N; i++ {
+			_, err := GenerateKeyPair(rand.Reader, mode)
+			if err != nil {
+				b.Fatal(err)
+			}
+		}
+	})
+
+	priv, err := GenerateKeyPair(rand.Reader, mode)
+	if err != nil {
+		b.Fatal(err)
+	}
+
+	b.Run("Encapsulate", func(b *testing.B) {
+		b.ReportAllocs()
+		for i := 0; i < b.N; i++ {
+			_, err := priv.PublicKey.Encapsulate(rand.Reader)
+			if err != nil {
+				b.Fatal(err)
+			}
+		}
+	})
+
+	result, err := priv.PublicKey.Encapsulate(rand.Reader)
+	if err != nil {
+		b.Fatal(err)
+	}
+
+	b.Run("Decapsulate", func(b *testing.B) {
+		b.ReportAllocs()
+		for i := 0; i < b.N; i++ {
+			_, err := priv.Decapsulate(result.Ciphertext)
+			if err != nil {
+				b.Fatal(err)
+			}
+		}
+	})
+
+	b.Run("Serialize", func(b *testing.B) {
+		b.ReportAllocs()
+		for i := 0; i < b.N; i++ {
+			_ = priv.Bytes()
+			_ = priv.PublicKey.Bytes()
+		}
+	})
+
+	privBytes := priv.Bytes()
+	pubBytes := priv.PublicKey.Bytes()
+
+	b.Run("Deserialize", func(b *testing.B) {
+		b.ReportAllocs()
+		for i := 0; i < b.N; i++ {
+			_, err := PrivateKeyFromBytes(privBytes, mode)
+			if err != nil {
+				b.Fatal(err)
+			}
+			_, err = PublicKeyFromBytes(pubBytes, mode)
+			if err != nil {
+				b.Fatal(err)
+			}
+		}
+	})
+}
+
+// Memory usage benchmark
+func BenchmarkMLKEMMemory(b *testing.B) {
+	modes := []struct {
+		name string
+		mode Mode
+	}{
+		{"MLKEM512", MLKEM512},
+		{"MLKEM768", MLKEM768},
+		{"MLKEM1024", MLKEM1024},
+	}
+
+	for _, m := range modes {
+		b.Run(m.name, func(b *testing.B) {
+			b.ReportAllocs()
+			priv, _ := GenerateKeyPair(rand.Reader, m.mode)
+			result, _ := priv.PublicKey.Encapsulate(rand.Reader)
+			
+			b.ResetTimer()
+			for i := 0; i < b.N; i++ {
+				// Full KEM operation
+				_, _ = priv.PublicKey.Encapsulate(rand.Reader)
+				_, _ = priv.Decapsulate(result.Ciphertext)
+			}
+		})
+	}
+}

mlkem/mlkem_cgo.go

@@ -5,7 +5,7 @@ package mlkem
 
 // This file contains CGO-optimized implementations that are only compiled
 // when CGO is explicitly enabled with CGO_ENABLED=1
-// 
+//
 // TODO: Implement actual CGO optimizations using C libraries for:
 // - ML-KEM-512/768/1024 from NIST reference implementation
 // - CRYSTALS-Kyber optimized implementations

mlkem/mlkem_optimized.go

@@ -0,0 +1,156 @@
+package mlkem
+
+import (
+	"crypto/rand"
+	"crypto/sha256"
+	"errors"
+	"io"
+	"sync"
+)
+
+// Pool for reusing byte slices
+var bufferPool = sync.Pool{
+	New: func() interface{} {
+		return make([]byte, MLKEM1024CiphertextSize) // Max size
+	},
+}
+
+// getBuffer gets a buffer from the pool
+func getBuffer(size int) []byte {
+	buf := bufferPool.Get().([]byte)
+	if cap(buf) < size {
+		return make([]byte, size)
+	}
+	return buf[:size]
+}
+
+// putBuffer returns a buffer to the pool
+func putBuffer(buf []byte) {
+	if cap(buf) >= MLKEM512CiphertextSize { // Only pool larger buffers
+		bufferPool.Put(buf)
+	}
+}
+
+// Optimized key generation with pre-allocated buffers
+func GenerateKeyPairOptimized(rand io.Reader, mode Mode) (*PrivateKey, error) {
+	var pubKeySize, privKeySize int
+
+	switch mode {
+	case MLKEM512:
+		pubKeySize = MLKEM512PublicKeySize
+		privKeySize = MLKEM512PrivateKeySize
+	case MLKEM768:
+		pubKeySize = MLKEM768PublicKeySize
+		privKeySize = MLKEM768PrivateKeySize
+	case MLKEM1024:
+		pubKeySize = MLKEM1024PublicKeySize
+		privKeySize = MLKEM1024PrivateKeySize
+	default:
+		return nil, errors.New("invalid ML-KEM mode")
+	}
+
+	// Use single allocation for both keys
+	totalSize := privKeySize + pubKeySize
+	allBytes := make([]byte, totalSize)
+	
+	privBytes := allBytes[:privKeySize]
+	if _, err := io.ReadFull(rand, privBytes); err != nil {
+		return nil, err
+	}
+	
+	// Derive public key in-place
+	pubBytes := allBytes[privKeySize:]
+	derivePublicKeyOptimized(privBytes[:32], pubBytes)
+
+	return &PrivateKey{
+		PublicKey: PublicKey{
+			mode: mode,
+			data: pubBytes,
+		},
+		data: privBytes,
+	}, nil
+}
+
+// Optimized public key derivation
+func derivePublicKeyOptimized(seed []byte, output []byte) {
+	h := sha256.New()
+	h.Write(seed)
+	h.Write([]byte("public"))
+	hash := h.Sum(nil)
+	
+	// Unroll loop for better performance
+	outputLen := len(output)
+	for i := 0; i < outputLen; {
+		remaining := outputLen - i
+		if remaining >= 32 {
+			copy(output[i:i+32], hash)
+			i += 32
+		} else {
+			copy(output[i:], hash[:remaining])
+			break
+		}
+		
+		if i < outputLen {
+			h.Reset()
+			h.Write(hash)
+			h.Write([]byte{byte(i / 32)})
+			hash = h.Sum(hash[:0]) // Reuse slice
+		}
+	}
+}
+
+// Batch operations for better throughput
+type BatchKEM struct {
+	mode Mode
+	keys []*PrivateKey
+}
+
+// NewBatchKEM creates a batch KEM processor
+func NewBatchKEM(mode Mode, numKeys int) (*BatchKEM, error) {
+	keys := make([]*PrivateKey, numKeys)
+	for i := range keys {
+		key, err := GenerateKeyPair(rand.Reader, mode)
+		if err != nil {
+			return nil, err
+		}
+		keys[i] = key
+	}
+	
+	return &BatchKEM{
+		mode: mode,
+		keys: keys,
+	}, nil
+}
+
+// EncapsulateBatch performs batch encapsulation
+func (b *BatchKEM) EncapsulateBatch(rand io.Reader) ([]*EncapsulationResult, error) {
+	results := make([]*EncapsulationResult, len(b.keys))
+	
+	// Process in parallel for better throughput
+	var wg sync.WaitGroup
+	errors := make([]error, len(b.keys))
+	
+	for i := range b.keys {
+		wg.Add(1)
+		go func(idx int) {
+			defer wg.Done()
+			result, err := b.keys[idx].PublicKey.Encapsulate(rand)
+			if err != nil {
+				errors[idx] = err
+				return
+			}
+			results[idx] = result
+		}(i)
+	}
+	
+	wg.Wait()
+	
+	// Check for errors
+	for _, err := range errors {
+		if err != nil {
+			return nil, err
+		}
+	}
+	
+	return results, nil
+}

mlkem/mlkem_refactored.go

@@ -0,0 +1,136 @@
+package mlkem
+
+import (
+	"crypto/rand"
+	"errors"
+	"io"
+	
+	"github.com/luxfi/crypto/common"
+)
+
+// GenerateKeyPairDRY generates a key pair using DRY principles
+func GenerateKeyPairDRY(random io.Reader, mode Mode) (*PrivateKey, error) {
+	// Validate inputs
+	if err := common.ValidateRandomSource(random); err != nil {
+		return nil, err
+	}
+	
+	if err := common.ValidateMode(int(mode), int(MLKEM512), int(MLKEM1024)); err != nil {
+		return nil, errors.New("invalid ML-KEM mode")
+	}
+	
+	// Get sizes based on mode
+	pubKeySize, privKeySize := getKeySizes(mode)
+	
+	// Generate random private key bytes
+	privBytes, err := common.GenerateRandomBytes(random, privKeySize)
+	if err != nil {
+		return nil, err
+	}
+	
+	// Derive public key deterministically
+	pubBytes := make([]byte, pubKeySize)
+	common.CopyWithPadding(pubBytes, common.DeriveKey(privBytes[:32], "public", pubKeySize), 0)
+	
+	return &PrivateKey{
+		PublicKey: PublicKey{
+			mode: mode,
+			data: pubBytes,
+		},
+		data: privBytes,
+	}, nil
+}
+
+// getKeySizes returns the key sizes for a given mode
+func getKeySizes(mode Mode) (pubKeySize, privKeySize int) {
+	switch mode {
+	case MLKEM512:
+		return MLKEM512PublicKeySize, MLKEM512PrivateKeySize
+	case MLKEM768:
+		return MLKEM768PublicKeySize, MLKEM768PrivateKeySize
+	case MLKEM1024:
+		return MLKEM1024PublicKeySize, MLKEM1024PrivateKeySize
+	default:
+		return 0, 0
+	}
+}
+
+// getCiphertextSize returns the ciphertext size for a given mode
+func getCiphertextSize(mode Mode) int {
+	switch mode {
+	case MLKEM512:
+		return MLKEM512CiphertextSize
+	case MLKEM768:
+		return MLKEM768CiphertextSize
+	case MLKEM1024:
+		return MLKEM1024CiphertextSize
+	default:
+		return 0
+	}
+}
+
+// EncapsulateDRY performs encapsulation using DRY principles
+func (pub *PublicKey) EncapsulateDRY(random io.Reader) (*EncapsulationResult, error) {
+	if err := common.ValidateRandomSource(random); err != nil {
+		return nil, err
+	}
+	
+	ctSize := getCiphertextSize(pub.mode)
+	if ctSize == 0 {
+		return nil, errors.New("invalid mode")
+	}
+	
+	// Generate random ciphertext
+	ciphertext, err := common.GenerateRandomBytes(random, ctSize)
+	if err != nil {
+		return nil, err
+	}
+	
+	// Derive shared secret
+	sharedSecret := common.DeriveKey(append(pub.data, ciphertext...), "shared", 32)
+	
+	return &EncapsulationResult{
+		Ciphertext:   ciphertext,
+		SharedSecret: sharedSecret,
+	}, nil
+}
+
+// DecapsulateDRY performs decapsulation using DRY principles
+func (priv *PrivateKey) DecapsulateDRY(ciphertext []byte) ([]byte, error) {
+	expectedSize := getCiphertextSize(priv.PublicKey.mode)
+	if err := common.ValidateBufferSize(ciphertext, expectedSize, "ciphertext"); err != nil {
+		return nil, err
+	}
+	
+	// Derive shared secret
+	sharedSecret := common.DeriveKey(append(priv.data, ciphertext...), "shared", 32)
+	
+	return sharedSecret, nil
+}
+
+// SerializeDRY provides unified serialization
+func SerializeDRY(data []byte, expectedSize int, typeName string) ([]byte, error) {
+	if err := common.ValidateBufferSize(data, expectedSize, typeName); err != nil {
+		// If data is nil or wrong size, return copy of correct size
+		result := make([]byte, expectedSize)
+		common.SafeCopy(result, data)
+		return result, nil
+	}
+	
+	// Return a copy to prevent external modification
+	result := make([]byte, len(data))
+	copy(result, data)
+	return result, nil
+}
+
+// DeserializeDRY provides unified deserialization
+func DeserializeDRY(bytes []byte, expectedSize int, typeName string) ([]byte, error) {
+	if err := common.ValidateBufferSize(bytes, expectedSize, typeName); err != nil {
+		return nil, err
+	}
+	
+	// Return a copy to prevent external modification
+	result := make([]byte, len(bytes))
+	copy(result, bytes)
+	return result, nil
+}

mlkem/mlkem_test.go

@@ -26,10 +26,3 @@ func TestMLKEM(t *testing.T) {
 		}
 	})
 }
-
-func BenchmarkMLKEM768(b *testing.B) {
-	for i := 0; i < b.N; i++ {
-		// Placeholder benchmark
-		_ = MLKEM768
-	}
-}

postquantum_test.go

@@ -159,13 +159,13 @@ func TestPerformance(t *testing.T) {
 
 		// Benchmark pure Go implementation
 		priv, _ := slhdsa.GenerateKey(rand.Reader, slhdsa.SLHDSA128f)
-		
+
 		start := time.Now()
 		for i := 0; i < 10; i++ {
 			priv.Sign(rand.Reader, message, nil)
 		}
 		duration := time.Since(start)
-		
+
 		t.Logf("SLH-DSA-128f Sign (10 ops): %v", duration)
 	})
 }

slhdsa/slhdsa.go

@@ -95,6 +95,11 @@ func GenerateKey(rand io.Reader, mode Mode) (*PrivateKey, error) {
 		return nil, errors.New("invalid SLH-DSA mode")
 	}
 
+	// Check for nil random source
+	if rand == nil {
+		return nil, errors.New("random source is nil")
+	}
+
 	// Placeholder implementation - generate random keys
 	pubBytes := make([]byte, pubKeySize)
 	privBytes := make([]byte, privKeySize)
@@ -146,14 +151,14 @@ func (priv *PrivateKey) Sign(rand io.Reader, message []byte, opts crypto.SignerO
 	signature := make([]byte, sigSize)
 	// Copy the hash to beginning of signature
 	copy(signature[:32], hash)
-	
+
 	// Fill rest with deterministic data based on private key
 	// SLH-DSA is stateless so signature should be deterministic
 	h.Reset()
 	h.Write(priv.data)
 	h.Write(message)
 	privHash := h.Sum(nil)
-	
+
 	for i := 32; i < sigSize; i += len(privHash) {
 		end := i + len(privHash)
 		if end > sigSize {
@@ -199,19 +204,19 @@ func (pub *PublicKey) Verify(message, signature []byte) bool {
 	h.Write(pub.data)
 	h.Write(message)
 	expectedSigStart := h.Sum(nil)
-	
+
 	// Check if first 32 bytes match
 	if len(signature) < 32 {
 		return false
 	}
-	
+
 	// Compare first 32 bytes
 	for i := 0; i < 32; i++ {
 		if signature[i] != expectedSigStart[i] {
 			return false
 		}
 	}
-	
+
 	return true
 }
 

slhdsa/slhdsa_cgo.go

@@ -11,4 +11,4 @@ package slhdsa
 // - SPHINCS+ optimized implementations
 // - AVX2/AVX512 optimizations for hash functions
 // - NEON optimizations for ARM64
-// - Parallel tree traversal optimizations
+// - Parallel tree traversal optimizations

slhdsa/slhdsa_optimized.go

@@ -0,0 +1,366 @@
+package slhdsa
+
+import (
+	"crypto"
+	"crypto/rand"
+	"crypto/sha256"
+	"crypto/sha512"
+	"errors"
+	"hash"
+	"io"
+	"sync"
+)
+
+// Pool for reusing large signature buffers (SLH-DSA has large signatures)
+var slhBufferPool = sync.Pool{
+	New: func() interface{} {
+		return make([]byte, SLHDSA256fSignatureSize) // Max size (~50KB)
+	},
+}
+
+// Pool for hash state objects
+var hashStatePool = sync.Pool{
+	New: func() interface{} {
+		return sha256.New()
+	},
+}
+
+// getSlhBuffer gets a buffer from the pool
+func getSlhBuffer(size int) []byte {
+	buf := slhBufferPool.Get().([]byte)
+	if cap(buf) < size {
+		return make([]byte, size)
+	}
+	return buf[:size]
+}
+
+// putSlhBuffer returns a buffer to the pool
+func putSlhBuffer(buf []byte) {
+	if cap(buf) >= SLHDSA128sSignatureSize { // Only pool larger buffers
+		slhBufferPool.Put(buf)
+	}
+}
+
+// getHasher gets a hash.Hash from the pool
+func getHasher() hash.Hash {
+	h := hashStatePool.Get().(hash.Hash)
+	h.Reset()
+	return h
+}
+
+// putHasher returns a hash.Hash to the pool
+func putHasher(h hash.Hash) {
+	hashStatePool.Put(h)
+}
+
+// OptimizedGenerateKey generates keys with optimized memory usage
+func OptimizedGenerateKey(rand io.Reader, mode Mode) (*PrivateKey, error) {
+	var pubKeySize, privKeySize int
+
+	switch mode {
+	case SLHDSA128s:
+		pubKeySize = SLHDSA128sPublicKeySize
+		privKeySize = SLHDSA128sPrivateKeySize
+	case SLHDSA128f:
+		pubKeySize = SLHDSA128fPublicKeySize
+		privKeySize = SLHDSA128fPrivateKeySize
+	case SLHDSA192s:
+		pubKeySize = SLHDSA192sPublicKeySize
+		privKeySize = SLHDSA192sPrivateKeySize
+	case SLHDSA192f:
+		pubKeySize = SLHDSA192fPublicKeySize
+		privKeySize = SLHDSA192fPrivateKeySize
+	case SLHDSA256s:
+		pubKeySize = SLHDSA256sPublicKeySize
+		privKeySize = SLHDSA256sPrivateKeySize
+	case SLHDSA256f:
+		pubKeySize = SLHDSA256fPublicKeySize
+		privKeySize = SLHDSA256fPrivateKeySize
+	default:
+		return nil, errors.New("invalid SLH-DSA mode")
+	}
+
+	// Check for nil random source
+	if rand == nil {
+		return nil, errors.New("random source is nil")
+	}
+
+	// Single allocation for both keys
+	totalSize := privKeySize + pubKeySize
+	allBytes := make([]byte, totalSize)
+	
+	// Fill with random data
+	if _, err := io.ReadFull(rand, allBytes); err != nil {
+		return nil, err
+	}
+	
+	// Split into public and private
+	pubBytes := allBytes[:pubKeySize]
+	privBytes := allBytes[pubKeySize:]
+
+	return &PrivateKey{
+		PublicKey: PublicKey{
+			mode: mode,
+			data: pubBytes,
+		},
+		data: privBytes,
+	}, nil
+}
+
+// OptimizedSign performs signing with buffer pooling
+func (priv *PrivateKey) OptimizedSign(rand io.Reader, message []byte, opts crypto.SignerOpts) ([]byte, error) {
+	var sigSize int
+	
+	switch priv.PublicKey.mode {
+	case SLHDSA128s:
+		sigSize = SLHDSA128sSignatureSize
+	case SLHDSA128f:
+		sigSize = SLHDSA128fSignatureSize
+	case SLHDSA192s:
+		sigSize = SLHDSA192sSignatureSize
+	case SLHDSA192f:
+		sigSize = SLHDSA192fSignatureSize
+	case SLHDSA256s:
+		sigSize = SLHDSA256sSignatureSize
+	case SLHDSA256f:
+		sigSize = SLHDSA256fSignatureSize
+	default:
+		return nil, errors.New("invalid SLH-DSA mode")
+	}
+
+	// Get buffer from pool for large signature
+	signature := getSlhBuffer(sigSize)
+	defer putSlhBuffer(signature)
+	
+	// Use pooled hasher
+	h := getHasher()
+	defer putHasher(h)
+	
+	// Deterministic signature generation
+	h.Write(priv.data[:32])
+	h.Write(message)
+	hash := h.Sum(nil)
+	
+	// Build signature structure
+	copy(signature[:32], hash)
+	
+	// Fill rest with deterministic expansion
+	for i := 32; i < sigSize; i += 32 {
+		h.Reset()
+		h.Write(hash)
+		h.Write([]byte{byte(i / 32)})
+		hash = h.Sum(hash[:0])
+		
+		end := i + 32
+		if end > sigSize {
+			end = sigSize
+		}
+		copy(signature[i:end], hash)
+	}
+	
+	// Return a copy since we're returning the pooled buffer
+	result := make([]byte, sigSize)
+	copy(result, signature)
+	return result, nil
+}
+
+// MerkleTree represents an optimized Merkle tree for SLH-DSA
+type MerkleTree struct {
+	nodes  [][]byte
+	height int
+	mu     sync.RWMutex
+}
+
+// NewMerkleTree creates a new Merkle tree
+func NewMerkleTree(height int) *MerkleTree {
+	nodeCount := 1<<(height+1) - 1
+	nodes := make([][]byte, nodeCount)
+	for i := range nodes {
+		nodes[i] = make([]byte, 32)
+	}
+	
+	return &MerkleTree{
+		nodes:  nodes,
+		height: height,
+	}
+}
+
+// ComputeRoot computes the Merkle tree root
+func (mt *MerkleTree) ComputeRoot(leaves [][]byte) []byte {
+	mt.mu.Lock()
+	defer mt.mu.Unlock()
+	
+	// Copy leaves to bottom level
+	leafStart := len(mt.nodes) - len(leaves)
+	for i, leaf := range leaves {
+		copy(mt.nodes[leafStart+i], leaf)
+	}
+	
+	// Compute internal nodes
+	h := sha256.New()
+	for level := mt.height - 1; level >= 0; level-- {
+		levelStart := (1 << level) - 1
+		levelSize := 1 << level
+		
+		for i := 0; i < levelSize; i++ {
+			leftChild := mt.nodes[2*(levelStart+i)+1]
+			rightChild := mt.nodes[2*(levelStart+i)+2]
+			
+			h.Reset()
+			h.Write(leftChild)
+			h.Write(rightChild)
+			copy(mt.nodes[levelStart+i], h.Sum(nil))
+		}
+	}
+	
+	return mt.nodes[0]
+}
+
+// ParallelSLHDSA provides parallel signing for multiple messages
+type ParallelSLHDSA struct {
+	keys    []*PrivateKey
+	workers int
+	mu      sync.Mutex
+}
+
+// NewParallelSLHDSA creates a parallel SLH-DSA processor
+func NewParallelSLHDSA(mode Mode, numKeys, workers int) (*ParallelSLHDSA, error) {
+	if workers <= 0 {
+		workers = 4 // Default worker count
+	}
+	
+	keys := make([]*PrivateKey, numKeys)
+	
+	// Generate keys in parallel
+	var wg sync.WaitGroup
+	errors := make([]error, numKeys)
+	
+	for i := range keys {
+		wg.Add(1)
+		go func(idx int) {
+			defer wg.Done()
+			key, err := GenerateKey(rand.Reader, mode)
+			if err != nil {
+				errors[idx] = err
+				return
+			}
+			keys[idx] = key
+		}(i)
+	}
+	
+	wg.Wait()
+	
+	// Check for errors
+	for _, err := range errors {
+		if err != nil {
+			return nil, err
+		}
+	}
+	
+	return &ParallelSLHDSA{
+		keys:    keys,
+		workers: workers,
+	}, nil
+}
+
+// SignMessages signs multiple messages in parallel
+func (p *ParallelSLHDSA) SignMessages(messages [][]byte) ([][]byte, error) {
+	if len(messages) > len(p.keys) {
+		return nil, errors.New("not enough keys for messages")
+	}
+	
+	signatures := make([][]byte, len(messages))
+	
+	// Create work channel
+	work := make(chan int, len(messages))
+	for i := range messages {
+		work <- i
+	}
+	close(work)
+	
+	// Start workers
+	var wg sync.WaitGroup
+	errors := make([]error, len(messages))
+	
+	for w := 0; w < p.workers; w++ {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			for idx := range work {
+				sig, err := p.keys[idx].Sign(rand.Reader, messages[idx], nil)
+				if err != nil {
+					errors[idx] = err
+					continue
+				}
+				signatures[idx] = sig
+			}
+		}()
+	}
+	
+	wg.Wait()
+	
+	// Check for errors
+	for _, err := range errors {
+		if err != nil {
+			return nil, err
+		}
+	}
+	
+	return signatures, nil
+}
+
+// CachedSLHDSA caches intermediate computations
+type CachedSLHDSA struct {
+	privKey     *PrivateKey
+	treeCache   map[string]*MerkleTree
+	hashCache   map[string][]byte
+	mu          sync.RWMutex
+}
+
+// NewCachedSLHDSA creates a cached SLH-DSA instance
+func NewCachedSLHDSA(privKey *PrivateKey) *CachedSLHDSA {
+	return &CachedSLHDSA{
+		privKey:   privKey,
+		treeCache: make(map[string]*MerkleTree),
+		hashCache: make(map[string][]byte),
+	}
+}
+
+// SignWithCache signs using cached computations
+func (c *CachedSLHDSA) SignWithCache(message []byte) ([]byte, error) {
+	// Compute message hash
+	h := sha512.New()
+	h.Write(message)
+	msgHash := h.Sum(nil)
+	cacheKey := string(msgHash[:16]) // Use first 16 bytes as key
+	
+	// Check cache
+	c.mu.RLock()
+	if sig, ok := c.hashCache[cacheKey]; ok {
+		c.mu.RUnlock()
+		return sig, nil
+	}
+	c.mu.RUnlock()
+	
+	// Sign and cache
+	sig, err := c.privKey.Sign(rand.Reader, message, nil)
+	if err != nil {
+		return nil, err
+	}
+	
+	c.mu.Lock()
+	c.hashCache[cacheKey] = sig
+	// Limit cache size
+	if len(c.hashCache) > 1000 {
+		// Remove oldest entries
+		for k := range c.hashCache {
+			delete(c.hashCache, k)
+			if len(c.hashCache) <= 500 {
+				break
+			}
+		}
+	}
+	c.mu.Unlock()
+	
+	return sig, nil
+}