From ba4f2a29954c4ee5dc37b866c52337cd1af6f6e8 Mon Sep 17 00:00:00 2001
From: Hanzo AI <dev@hanzo.ai>
Date: Sat, 27 Dec 2025 18:12:27 -0800
Subject: [PATCH] fix: kem cleanup

---
 kem/hybrid.go | 27 +++++++++++++++++++++++++++
 kem/mlkem.go  | 15 +++++++++++++++
 kem/x25519.go | 10 ++++++++++
 3 files changed, 52 insertions(+)

diff --git a/kem/hybrid.go b/kem/hybrid.go
index 8d5c075..69e13d7 100644
--- a/kem/hybrid.go
+++ b/kem/hybrid.go
@@ -3,6 +3,7 @@ package kem
 import (
 	"crypto/sha256"
 	"errors"
+	"fmt"
 	"io"
 
 	"golang.org/x/crypto/hkdf"
@@ -159,6 +160,32 @@ func (h *HybridKEMImpl) SharedSecretSize() int {
 	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
 func (pk *HybridPublicKey) Bytes() []byte {
 	return append(pk.X25519PK.Bytes(), pk.MLKEMPK.Bytes()...)
diff --git a/kem/mlkem.go b/kem/mlkem.go
index 9112074..136cb2b 100644
--- a/kem/mlkem.go
+++ b/kem/mlkem.go
@@ -7,6 +7,7 @@ import (
 	"crypto/mlkem"
 	"crypto/subtle"
 	"errors"
+	"fmt"
 )
 
 // 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.
 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.
 func (pk *MLKEM768PublicKey) Bytes() []byte { return pk.data }
 
diff --git a/kem/x25519.go b/kem/x25519.go
index d1a9bdf..e3bcf11 100644
--- a/kem/x25519.go
+++ b/kem/x25519.go
@@ -122,6 +122,16 @@ func (x *X25519Impl) PrivateKeySize() int   { return 32 }
 func (x *X25519Impl) CiphertextSize() 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
 func (pk *X25519PublicKey) Bytes() []byte {
 	return pk.data[:]