try to extract the key in order to decide whether to attach the key

libp2p · Oct 13, 2018 · c4b3bd1 · c4b3bd1
1 parent cb7679f
commit c4b3bd1
Show file tree

Hide file tree

Showing 3 changed files with 7 additions and 5 deletions.
diff --git 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
@@ -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
@@ -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)