From 496e58075393ac5083fbef2743f06248fff70083 Mon Sep 17 00:00:00 2001
From: David Stainton <dstainton415@gmail.com>
Date: Wed, 24 Apr 2024 12:21:56 -0700
Subject: [PATCH] Add some improvements to mod p dh

---
 nike/diffiehellman/dh.go     | 51 ++++++++++++++++++------------------
 nike/schemes/schemes.go      |  2 ++
 nike/schemes/schemes_test.go |  4 +++
 nike/x448/x448.go            |  5 ++--
 4 files changed, 35 insertions(+), 27 deletions(-)

diff --git a/nike/diffiehellman/dh.go b/nike/diffiehellman/dh.go
index 9509a9c..6fa4248 100644
--- a/nike/diffiehellman/dh.go
+++ b/nike/diffiehellman/dh.go
@@ -16,16 +16,11 @@ import (
 )
 
 const (
-	// GroupElementLength is the length of a ECDH group element in bytes.
-	// XXX wrong size FIXME: fix the serialization so that key blobs are
-	// the same size as group element size. i think gob encoding adds extra bytes.
-	GroupElementLength = 56
+	// PublicKeySize is the size of a serialized compressed PublicKey in bytes.
+	PublicKeySize = 407
 
-	// PublicKeySize is the size of a serialized PublicKey in bytes.
-	PublicKeySize = GroupElementLength
-
-	// PrivateKeySize is the size of a serialized PrivateKey in bytes.
-	PrivateKeySize = GroupElementLength
+	// PrivateKeySize is the size of a serialized compressed PrivateKey in bytes.
+	PrivateKeySize = 40
 )
 
 var (
@@ -184,19 +179,20 @@ func (p *PrivateKey) Public() nike.PublicKey {
 }
 
 func (p *PrivateKey) Reset() {
-	// no op
+	p.privKey.Reset()
 }
 
 func (p *PrivateKey) Bytes() []byte {
-	blob, err := p.privKey.GobEncode()
-	if err != nil {
-		panic(err)
-	}
-	return blob
+	return p.privKey.BinaryEncode()
 }
 
 func (p *PrivateKey) FromBytes(data []byte) error {
-	return p.privKey.GobDecode(data)
+	if len(data) != PrivateKeySize {
+		return errInvalidKey
+	}
+
+	p.privKey = new(cyclic.Int)
+	return p.privKey.BinaryDecode(data)
 }
 
 func (p *PrivateKey) MarshalBinary() ([]byte, error) {
@@ -225,20 +221,20 @@ type PublicKey struct {
 }
 
 func (p *PublicKey) Blind(blindingFactor nike.PrivateKey) error {
-	// FIX ME
-	return nil
+	_, ok := blindingFactor.(*PrivateKey)
+	if !ok {
+		return errors.New("blindingFactor nike.PrivateKey must be the same concrete type as diffiehellman.PublicKey")
+	}
+	pubBytes := Exp(blindingFactor.(*PrivateKey).privKey, p.pubKey)
+	return p.FromBytes(pubBytes)
 }
 
 func (p *PublicKey) Reset() {
-	// no op
+	p.pubKey.Reset()
 }
 
 func (p *PublicKey) Bytes() []byte {
-	blob, err := p.pubKey.GobEncode()
-	if err != nil {
-		panic(err)
-	}
-	return blob
+	return p.pubKey.BinaryEncode()
 }
 
 func (p *PublicKey) rebuildB64String() {
@@ -246,7 +242,12 @@ func (p *PublicKey) rebuildB64String() {
 }
 
 func (p *PublicKey) FromBytes(data []byte) error {
-	err := p.pubKey.GobDecode(data)
+	if len(data) != PublicKeySize {
+		return errInvalidKey
+	}
+
+	p.pubKey = new(cyclic.Int)
+	err := p.pubKey.BinaryDecode(data)
 	if err != nil {
 		return err
 	}
diff --git a/nike/schemes/schemes.go b/nike/schemes/schemes.go
index bf23ba2..8a7a4bb 100644
--- a/nike/schemes/schemes.go
+++ b/nike/schemes/schemes.go
@@ -8,6 +8,7 @@ import (
 	"github.com/katzenpost/hpqc/nike/ctidh/ctidh2048"
 	"github.com/katzenpost/hpqc/nike/ctidh/ctidh511"
 	"github.com/katzenpost/hpqc/nike/ctidh/ctidh512"
+	"github.com/katzenpost/hpqc/nike/diffiehellman"
 	"github.com/katzenpost/hpqc/nike/hybrid"
 	"github.com/katzenpost/hpqc/nike/x25519"
 	"github.com/katzenpost/hpqc/nike/x448"
@@ -19,6 +20,7 @@ var allSchemes = [...]nike.Scheme{
 	// classical NIKE schemes
 	x25519.Scheme(rand.Reader),
 	x448.Scheme(rand.Reader),
+	diffiehellman.Scheme(rand.Reader),
 
 	// post quantum NIKE schemes
 	ctidh511.Scheme(),
diff --git a/nike/schemes/schemes_test.go b/nike/schemes/schemes_test.go
index 2454fc3..62d3717 100644
--- a/nike/schemes/schemes_test.go
+++ b/nike/schemes/schemes_test.go
@@ -22,6 +22,8 @@ func TestNIKEUnmarshaling(t *testing.T) {
 		pubkey1Blob, err := pubkey1.MarshalBinary()
 		require.NoError(t, err)
 
+		t.Logf("pubkey1Blob is len %d", len(pubkey1Blob))
+
 		pubkey2, err := s.UnmarshalBinaryPublicKey(pubkey1Blob)
 		require.NoError(t, err)
 
@@ -30,6 +32,8 @@ func TestNIKEUnmarshaling(t *testing.T) {
 		privkey1blob, err := privkey1.MarshalBinary()
 		require.NoError(t, err)
 
+		t.Logf("privkey1blob is len %d", len(privkey1blob))
+
 		privkey2, err := s.UnmarshalBinaryPrivateKey(privkey1blob)
 		require.NoError(t, err)
 
diff --git a/nike/x448/x448.go b/nike/x448/x448.go
index 6a1cf71..4037ea5 100644
--- a/nike/x448/x448.go
+++ b/nike/x448/x448.go
@@ -223,8 +223,9 @@ type PublicKey struct {
 }
 
 func (p *PublicKey) Blind(blindingFactor nike.PrivateKey) error {
-	if len(blindingFactor.Bytes()) != GroupElementLength {
-		return ErrBlindDataSizeInvalid
+	_, ok := blindingFactor.(*PrivateKey)
+	if !ok {
+		return errors.New("blindingFactor nike.PrivateKey must be the same concrete type as x448.PublicKey")
 	}
 	pubBytes := Exp(p.pubBytes, blindingFactor.(*PrivateKey).privBytes)
 	copy(p.pubBytes[:], pubBytes)