From 4dc796359e228bdb640abd190eee8df5adfd42bb Mon Sep 17 00:00:00 2001
From: vyzo <vyzo@hackzen.org>
Date: Sun, 14 Oct 2018 02:16:41 +0300
Subject: [PATCH] try to extract the key in order to decide whether to attach
 the key

---
 pubsub.go    | 2 +-
 sign.go      | 8 +++++---
 sign_test.go | 2 +-
 3 files changed, 7 insertions(+), 5 deletions(-)

diff --git a/pubsub.go b/pubsub.go
index 4c0ed55..0300cab 100644
--- a/pubsub.go
+++ b/pubsub.go
@@ -677,7 +677,7 @@ func (p *PubSub) Publish(topic string, data []byte) error {
 	}
 	if p.signKey != nil {
 		m.From = []byte(p.signID)
-		err := signMessage(p.signKey, m)
+		err := signMessage(p.signID, p.signKey, m)
 		if err != nil {
 			return err
 		}
diff --git a/sign.go b/sign.go
index 7d261d9..bd01c33 100644
--- a/sign.go
+++ b/sign.go
@@ -68,7 +68,7 @@ func messagePubKey(m *pb.Message) (crypto.PubKey, error) {
 	return pubk, nil
 }
 
-func signMessage(key crypto.PrivKey, m *pb.Message) error {
+func signMessage(pid peer.ID, key crypto.PrivKey, m *pb.Message) error {
 	bytes, err := m.Marshal()
 	if err != nil {
 		return err
@@ -82,14 +82,16 @@ func signMessage(key crypto.PrivKey, m *pb.Message) error {
 	}
 
 	m.Signature = sig
-	switch key.Type() {
-	case crypto.RSA:
+
+	pk, _ := pid.ExtractPublicKey()
+	if pk == nil {
 		pubk, err := key.GetPublic().Bytes()
 		if err != nil {
 			return err
 		}
 		m.Key = pubk
 	}
+
 	return nil
 }
 
diff --git a/sign_test.go b/sign_test.go
index 507a6c1..d12cdca 100644
--- a/sign_test.go
+++ b/sign_test.go
@@ -34,7 +34,7 @@ func testSignVerify(t *testing.T, privk crypto.PrivKey) {
 		From:     []byte(id),
 		Seqno:    []byte("123"),
 	}
-	signMessage(privk, &m)
+	signMessage(id, privk, &m)
 	err = verifyMessageSignature(&m)
 	if err != nil {
 		t.Fatal(err)