From 60a06ce9616ac5ca4bac4c8f4c743d7cf1d592c8 Mon Sep 17 00:00:00 2001
From: vyzo <vyzo@hackzen.org>
Date: Sun, 26 Aug 2018 13:57:14 +0300
Subject: [PATCH] implement signMessage and verifyMessageSignature

---
 pubsub.go |  5 +++-
 sign.go   | 74 ++++++++++++++++++++++++++++++++++++++++++++++++++++++-
 2 files changed, 77 insertions(+), 2 deletions(-)

diff --git a/pubsub.go b/pubsub.go
index c1e5cc4..cd1291a 100644
--- a/pubsub.go
+++ b/pubsub.go
@@ -664,7 +664,10 @@ func (p *PubSub) Publish(topic string, data []byte) error {
 		Seqno:    seqno,
 	}
 	if p.signKey != nil {
-		signMessage(p.signKey, m)
+		err := signMessage(p.signKey, m)
+		if err != nil {
+			return err
+		}
 	}
 	p.publish <- &Message{m}
 	return nil
diff --git a/sign.go b/sign.go
index 5b55aad..8dd9998 100644
--- a/sign.go
+++ b/sign.go
@@ -1,15 +1,87 @@
 package floodsub
 
 import (
+	"fmt"
+
 	pb "github.com/libp2p/go-floodsub/pb"
 
 	crypto "github.com/libp2p/go-libp2p-crypto"
+	peer "github.com/libp2p/go-libp2p-peer"
 )
 
 func verifyMessageSignature(m *pb.Message) error {
+	var pubk crypto.PubKey
+
+	pid, err := peer.IDFromBytes(m.From)
+	if err != nil {
+		return err
+	}
+
+	if m.Key == nil {
+		// no attached key, it must be extractable from the source ID
+		pubk, err = pid.ExtractPublicKey()
+		if err != nil {
+			return err
+		}
+	} else {
+		pubk, err = crypto.UnmarshalPublicKey(m.Key)
+		if err != nil {
+			return err
+		}
+
+		// verify that the source ID matches the attached key
+		xpid, err := peer.IDFromPublicKey(pubk)
+		if err != nil {
+			return err
+		}
+
+		if pid != xpid {
+			return fmt.Errorf("bad signing key; source ID/key mismatch: %s %s", pid, xpid)
+		}
+	}
+
+	xm := pb.Message{
+		Data:     m.Data,
+		TopicIDs: m.TopicIDs,
+		From:     m.From,
+		Seqno:    m.Seqno,
+	}
+	bytes, err := xm.Marshal()
+	if err != nil {
+		return err
+	}
+
+	valid, err := pubk.Verify(bytes, m.Signature)
+	if err != nil {
+		return err
+	}
+
+	if !valid {
+		return fmt.Errorf("invalid signature")
+	}
+
 	return nil
 }
 
-func signMessage(key crypto.PrivKey, m *pb.Message) {
+func signMessage(key crypto.PrivKey, m *pb.Message) error {
+	bytes, err := m.Marshal()
+	if err != nil {
+		return err
+	}
 
+	sig, err := key.Sign(bytes)
+	if err != nil {
+		return err
+	}
+
+	m.Signature = sig
+	switch key.Type() {
+	case crypto.RSA:
+		pubk, err := key.GetPublic().Bytes()
+		if err != nil {
+			return err
+		}
+		m.Key = pubk
+	}
+	return nil
 }