fix: kem cleanup

This commit is contained in:
Hanzo AI
2025-12-27 18:12:27 -08:00
parent f937139f29
commit ba4f2a2995
3 changed files with 52 additions and 0 deletions
+27
View File
@@ -3,6 +3,7 @@ package kem
import ( import (
"crypto/sha256" "crypto/sha256"
"errors" "errors"
"fmt"
"io" "io"
"golang.org/x/crypto/hkdf" "golang.org/x/crypto/hkdf"
@@ -159,6 +160,32 @@ func (h *HybridKEMImpl) SharedSecretSize() int {
return 32 // HKDF output size return 32 // HKDF output size
} }
// ParseHybridPublicKey reconstructs a HybridPublicKey from concatenated bytes
// (X25519 32 bytes || ML-KEM-768 encapsulation key).
func ParseHybridPublicKey(data []byte) (*HybridPublicKey, error) {
x25519Size := 32
mlkemSize := mlkem768PublicKeySize
expected := x25519Size + mlkemSize
if len(data) != expected {
return nil, fmt.Errorf("invalid hybrid public key size: got %d, expected %d", len(data), expected)
}
x25519PK, err := ParseX25519PublicKey(data[:x25519Size])
if err != nil {
return nil, fmt.Errorf("invalid X25519 component: %w", err)
}
mlkemPK, err := ParseMLKEM768PublicKey(data[x25519Size:])
if err != nil {
return nil, fmt.Errorf("invalid ML-KEM-768 component: %w", err)
}
return &HybridPublicKey{
X25519PK: x25519PK,
MLKEMPK: mlkemPK,
}, nil
}
// Bytes returns the concatenated public key bytes // Bytes returns the concatenated public key bytes
func (pk *HybridPublicKey) Bytes() []byte { func (pk *HybridPublicKey) Bytes() []byte {
return append(pk.X25519PK.Bytes(), pk.MLKEMPK.Bytes()...) return append(pk.X25519PK.Bytes(), pk.MLKEMPK.Bytes()...)
+15
View File
@@ -7,6 +7,7 @@ import (
"crypto/mlkem" "crypto/mlkem"
"crypto/subtle" "crypto/subtle"
"errors" "errors"
"fmt"
) )
// MLKEM768Impl implements ML-KEM-768 using Go's crypto/mlkem (FIPS 203) // MLKEM768Impl implements ML-KEM-768 using Go's crypto/mlkem (FIPS 203)
@@ -85,6 +86,20 @@ func (m *MLKEM768Impl) CiphertextSize() int { return mlkem.CiphertextSize768 }
// SharedSecretSize returns the size of ML-KEM-768 shared secrets. // SharedSecretSize returns the size of ML-KEM-768 shared secrets.
func (m *MLKEM768Impl) SharedSecretSize() int { return mlkem.SharedKeySize } func (m *MLKEM768Impl) SharedSecretSize() int { return mlkem.SharedKeySize }
// ParseMLKEM768PublicKey reconstructs an MLKEM768PublicKey from raw encapsulation key bytes.
func ParseMLKEM768PublicKey(data []byte) (*MLKEM768PublicKey, error) {
if len(data) != mlkem.EncapsulationKeySize768 {
return nil, fmt.Errorf("invalid ML-KEM-768 public key size: got %d, expected %d", len(data), mlkem.EncapsulationKeySize768)
}
// Validate by attempting to construct the Go stdlib encapsulation key
if _, err := mlkem.NewEncapsulationKey768(data); err != nil {
return nil, fmt.Errorf("invalid ML-KEM-768 public key: %w", err)
}
cp := make([]byte, len(data))
copy(cp, data)
return &MLKEM768PublicKey{data: cp}, nil
}
// Bytes returns the raw bytes of the public key. // Bytes returns the raw bytes of the public key.
func (pk *MLKEM768PublicKey) Bytes() []byte { return pk.data } func (pk *MLKEM768PublicKey) Bytes() []byte { return pk.data }
+10
View File
@@ -122,6 +122,16 @@ func (x *X25519Impl) PrivateKeySize() int { return 32 }
func (x *X25519Impl) CiphertextSize() int { return 32 } func (x *X25519Impl) CiphertextSize() int { return 32 }
func (x *X25519Impl) SharedSecretSize() int { return 32 } func (x *X25519Impl) SharedSecretSize() int { return 32 }
// ParseX25519PublicKey reconstructs an X25519PublicKey from 32 raw bytes.
func ParseX25519PublicKey(data []byte) (*X25519PublicKey, error) {
if len(data) != 32 {
return nil, errors.New("invalid X25519 public key size: expected 32 bytes")
}
pk := &X25519PublicKey{}
copy(pk.data[:], data)
return pk, nil
}
// Bytes returns the public key bytes // Bytes returns the public key bytes
func (pk *X25519PublicKey) Bytes() []byte { func (pk *X25519PublicKey) Bytes() []byte {
return pk.data[:] return pk.data[:]