mkem: cleanup, remove pub/priv key wrapper types

katzenpost · Oct 1, 2024 · 2a1635f · 2a1635f
1 parent b16316b
commit 2a1635f
Show file tree

Hide file tree

Showing 3 changed files with 36 additions and 61 deletions.
diff --git a/kem/mkem/ciphertext.go b/kem/mkem/ciphertext.go
@@ -4,15 +4,19 @@
 // Package mkem provides multiparty KEM construction.
 package mkem
 
-import "github.com/fxamacker/cbor/v2"
+import (
+	"github.com/fxamacker/cbor/v2"
+
+	"github.com/katzenpost/hpqc/nike"
+)
 
 var (
 	// Create reusable EncMode interface with immutable options, safe for concurrent use.
 	ccbor cbor.EncMode
 )
 
 type Ciphertext struct {
-	EphemeralPublicKey *PublicKey
+	EphemeralPublicKey nike.PublicKey
 	DEKCiphertexts     [][]byte
 	Envelope           []byte
 }
@@ -50,18 +54,16 @@ func CiphertextFromBytes(scheme *Scheme, b []byte) (*Ciphertext, error) {
 		return nil, err
 	}
 	c := &Ciphertext{
-		EphemeralPublicKey: &PublicKey{
-			publicKey: pubkey,
-		},
-		DEKCiphertexts: ic.DEKCiphertexts,
-		Envelope:       ic.Envelope,
+		EphemeralPublicKey: pubkey,
+		DEKCiphertexts:     ic.DEKCiphertexts,
+		Envelope:           ic.Envelope,
 	}
 	return c, nil
 }
 
 func (m *Ciphertext) Marshal() []byte {
 	ic := &IntermediaryCiphertext{
-		EphemeralPublicKey: m.EphemeralPublicKey.publicKey.Bytes(),
+		EphemeralPublicKey: m.EphemeralPublicKey.Bytes(),
 		DEKCiphertexts:     m.DEKCiphertexts,
 		Envelope:           m.Envelope,
 	}

diff --git a/kem/mkem/mkem.go b/kem/mkem/mkem.go
@@ -12,34 +12,26 @@ import (
 	"github.com/katzenpost/chacha20poly1305"
 	"github.com/katzenpost/hpqc/hash"
 	"github.com/katzenpost/hpqc/nike"
+	"github.com/katzenpost/hpqc/nike/hybrid"
 )
 
 // Scheme is an MKEM scheme.
 type Scheme struct {
 	nike nike.Scheme
 }
 
-// FromNIKE creates a new KEM adapter Scheme
-// using the given NIKE Scheme.
-func FromNIKE(nike nike.Scheme) *Scheme {
-	if nike == nil {
-		panic("NIKE is nil")
-	}
+func NewScheme() *Scheme {
 	return &Scheme{
-		nike: nike,
+		nike: hybrid.CTIDH1024X25519,
 	}
 }
 
-func (s *Scheme) GenerateKeyPair() (*PublicKey, *PrivateKey, error) {
+func (s *Scheme) GenerateKeyPair() (nike.PublicKey, nike.PrivateKey, error) {
 	pubkey, privkey, err := s.nike.GenerateKeyPair()
 	if err != nil {
 		return nil, nil, err
 	}
-	return &PublicKey{
-			publicKey: pubkey,
-		}, &PrivateKey{
-			privateKey: privkey,
-		}, nil
+	return pubkey, privkey, nil
 }
 
 func (s *Scheme) createCipher(key []byte) cipher.AEAD {
@@ -67,42 +59,40 @@ func (s *Scheme) decrypt(key []byte, ciphertext []byte) ([]byte, error) {
 	return aead.Open(nil, nonce, ciphertext, nil)
 }
 
-func (s *Scheme) EnvelopeReply(privkey *PrivateKey, pubkey *PublicKey, plaintext []byte) []byte {
-	secret := hash.Sum256(s.nike.DeriveSecret(privkey.privateKey, pubkey.publicKey))
+func (s *Scheme) EnvelopeReply(privkey nike.PrivateKey, pubkey nike.PublicKey, plaintext []byte) []byte {
+	secret := hash.Sum256(s.nike.DeriveSecret(privkey, pubkey))
 	ciphertext := s.encrypt(secret[:], plaintext)
 	c := &Ciphertext{
-		EphemeralPublicKey: &PublicKey{
-			publicKey: pubkey.publicKey,
-		},
-		DEKCiphertexts: nil,
-		Envelope:       ciphertext,
+		EphemeralPublicKey: pubkey,
+		DEKCiphertexts:     nil,
+		Envelope:           ciphertext,
 	}
 	return c.Marshal()
 }
 
-func (s *Scheme) DecryptEnvelope(privkey *PrivateKey, pubkey *PublicKey, ciphertext []byte) ([]byte, error) {
+func (s *Scheme) DecryptEnvelope(privkey nike.PrivateKey, pubkey nike.PublicKey, ciphertext []byte) ([]byte, error) {
 	c, err := CiphertextFromBytes(s, ciphertext)
 	if err != nil {
 		return nil, err
 	}
 
-	secret := hash.Sum256(s.nike.DeriveSecret(privkey.privateKey, pubkey.publicKey))
+	secret := hash.Sum256(s.nike.DeriveSecret(privkey, pubkey))
 	plaintext, err := s.decrypt(secret[:], c.Envelope)
 	if err != nil {
 		return nil, err
 	}
 	return plaintext, nil
 }
 
-func (s *Scheme) Encapsulate(keys []*PublicKey, payload []byte) (*PrivateKey, []byte) {
+func (s *Scheme) Encapsulate(keys []nike.PublicKey, payload []byte) (nike.PrivateKey, []byte) {
 	ephPub, ephPriv, err := s.nike.GenerateKeyPair()
 	if err != nil {
 		panic(err)
 	}
 
 	secrets := make([][hash.HashSize]byte, len(keys))
 	for i := 0; i < len(keys); i++ {
-		secrets[i] = hash.Sum256(s.nike.DeriveSecret(ephPriv, keys[i].publicKey))
+		secrets[i] = hash.Sum256(s.nike.DeriveSecret(ephPriv, keys[i]))
 	}
 
 	msgKey := make([]byte, 32)
@@ -118,24 +108,20 @@ func (s *Scheme) Encapsulate(keys []*PublicKey, payload []byte) (*PrivateKey, []
 	}
 
 	c := &Ciphertext{
-		EphemeralPublicKey: &PublicKey{
-			publicKey: ephPub,
-		},
-		DEKCiphertexts: outCiphertexts,
-		Envelope:       ciphertext,
+		EphemeralPublicKey: ephPub,
+		DEKCiphertexts:     outCiphertexts,
+		Envelope:           ciphertext,
 	}
-	return &PrivateKey{
-		privateKey: ephPriv,
-	}, c.Marshal()
+	return ephPriv, c.Marshal()
 }
 
-func (s *Scheme) Decapsulate(privkey *PrivateKey, ciphertext []byte) ([]byte, error) {
+func (s *Scheme) Decapsulate(privkey nike.PrivateKey, ciphertext []byte) ([]byte, error) {
 	c, err := CiphertextFromBytes(s, ciphertext)
 	if err != nil {
 		return nil, err
 	}
 
-	ephSecret := hash.Sum256(s.nike.DeriveSecret(privkey.privateKey, c.EphemeralPublicKey.publicKey))
+	ephSecret := hash.Sum256(s.nike.DeriveSecret(privkey, c.EphemeralPublicKey))
 	for i := 0; i < len(c.DEKCiphertexts); i++ {
 		msgKey, err := s.decrypt(ephSecret[:], c.DEKCiphertexts[i])
 		if err != nil {
@@ -145,13 +131,3 @@ func (s *Scheme) Decapsulate(privkey *PrivateKey, ciphertext []byte) ([]byte, er
 	}
 	return nil, errors.New("failed to trial decrypt")
 }
-
-// PrivateKey is an MKEM private key.
-type PrivateKey struct {
-	privateKey nike.PrivateKey
-}
-
-// PublicKey is an MKEM public key.
-type PublicKey struct {
-	publicKey nike.PublicKey
-}
diff --git a/kem/mkem/mkem_test.go b/kem/mkem/mkem_test.go
@@ -8,8 +8,9 @@ import (
 	"crypto/rand"
 	"testing"
 
-	"github.com/katzenpost/hpqc/nike/schemes"
 	"github.com/stretchr/testify/require"
+
+	"github.com/katzenpost/hpqc/nike"
 )
 
 func TestCiphertextMarshaling(t *testing.T) {
@@ -34,9 +35,7 @@ func TestCiphertextMarshaling(t *testing.T) {
 }
 
 func TestMKEMCorrectness(t *testing.T) {
-	nikeName := "x25519"
-	nike := schemes.ByName(nikeName)
-	s := FromNIKE(nike)
+	s := NewScheme()
 
 	replica1pub, replica1priv, err := s.GenerateKeyPair()
 	require.NoError(t, err)
@@ -48,7 +47,7 @@ func TestMKEMCorrectness(t *testing.T) {
 	_, err = rand.Reader.Read(secret)
 	require.NoError(t, err)
 
-	_, ciphertext := s.Encapsulate([]*PublicKey{replica1pub, replica2pub}, secret)
+	_, ciphertext := s.Encapsulate([]nike.PublicKey{replica1pub, replica2pub}, secret)
 
 	ciphertext2, err := CiphertextFromBytes(s, ciphertext)
 	require.NoError(t, err)
@@ -67,9 +66,7 @@ func TestMKEMCorrectness(t *testing.T) {
 }
 
 func TestMKEMProtocol(t *testing.T) {
-	nikeName := "x25519"
-	nike := schemes.ByName(nikeName)
-	s := FromNIKE(nike)
+	s := NewScheme()
 
 	// replicas create their keys and publish them
 	replica1pub, replica1priv, err := s.GenerateKeyPair()
@@ -81,7 +78,7 @@ func TestMKEMProtocol(t *testing.T) {
 	request := make([]byte, 32)
 	_, err = rand.Reader.Read(request)
 	require.NoError(t, err)
-	privKey1, envelopeRaw := s.Encapsulate([]*PublicKey{replica1pub, replica2pub}, request)
+	privKey1, envelopeRaw := s.Encapsulate([]nike.PublicKey{replica1pub, replica2pub}, request)
 	envelope1, err := CiphertextFromBytes(s, envelopeRaw)
 	require.NoError(t, err)