Added encryption for messages better API for sealing messages

3 files changed, 89 insertions, 21 deletions

diff --git a/whisper/envelope.go b/whisper/envelope.go
index 8e66a7bbb..eb80098ad 100644
--- a/whisper/envelope.go
+++ b/whisper/envelope.go
@@ -16,8 +16,8 @@ const (
 )
 
 type Envelope struct {
-	Expiry int32 // Whisper protocol specifies int32, really should be int64
-	Ttl    int32 // ^^^^^^
+	Expiry uint32 // Whisper protocol specifies int32, really should be int64
+	Ttl    uint32 // ^^^^^^
 	Topics [][]byte
 	Data   []byte
 	Nonce  uint32
@@ -52,11 +52,11 @@ func (self *Envelope) Hash() Hash {
 func NewEnvelope(ttl time.Duration, topics [][]byte, data *Message) *Envelope {
 	exp := time.Now().Add(ttl)
 
-	return &Envelope{int32(exp.Unix()), int32(ttl.Seconds()), topics, data.Bytes(), 0, Hash{}}
+	return &Envelope{uint32(exp.Unix()), uint32(ttl.Seconds()), topics, data.Bytes(), 0, Hash{}}
 }
 
-func (self *Envelope) Seal() {
-	self.proveWork(DefaultPow)
+func (self *Envelope) Seal(pow time.Duration) {
+	self.proveWork(pow)
 }
 
 func (self *Envelope) proveWork(dura time.Duration) {
diff --git a/whisper/message.go b/whisper/message.go
index 21cf163e6..408b9f7df 100644
--- a/whisper/message.go
+++ b/whisper/message.go
@@ -1,5 +1,11 @@
 package whisper
 
+import (
+	"time"
+
+	"github.com/ethereum/go-ethereum/crypto"
+)
+
 type Message struct {
 	Flags     byte
 	Signature []byte
@@ -10,6 +16,49 @@ func NewMessage(payload []byte) *Message {
 	return &Message{Flags: 0, Payload: payload}
 }
 
+func (self *Message) hash() []byte {
+	return crypto.Sha3(append([]byte{self.Flags}, self.Payload...))
+}
+
+func (self *Message) sign(key []byte) (err error) {
+	self.Flags = 1
+	self.Signature, err = crypto.Sign(self.hash(), key)
+	return
+}
+
+func (self *Message) Encrypt(from, to []byte) (err error) {
+	err = self.sign(from)
+	if err != nil {
+		return err
+	}
+
+	self.Payload, err = crypto.Encrypt(to, self.Payload)
+	if err != nil {
+		return err
+	}
+
+	return nil
+}
+
 func (self *Message) Bytes() []byte {
 	return append([]byte{self.Flags}, append(self.Signature, self.Payload...)...)
 }
+
+type Opts struct {
+	From, To []byte // private(sender), public(receiver) key
+	Ttl      time.Duration
+	Topics   [][]byte
+}
+
+func (self *Message) Seal(pow time.Duration, opts Opts) (*Envelope, error) {
+	if len(opts.To) > 0 && len(opts.From) > 0 {
+		if err := self.Encrypt(opts.From, opts.To); err != nil {
+			return nil, err
+		}
+	}
+
+	envelope := NewEnvelope(DefaultTtl, opts.Topics, self)
+	envelope.Seal(pow)
+
+	return envelope, nil
+}
diff --git a/whisper/whisper.go b/whisper/whisper.go
index 78e4d4848..b4e37b959 100644
--- a/whisper/whisper.go
+++ b/whisper/whisper.go
@@ -2,6 +2,7 @@ package whisper
 
 import (
 	"bytes"
+	"errors"
 	"fmt"
 	"sync"
 	"time"
@@ -35,7 +36,7 @@ const (
 	envelopesMsg = 0x01
 )
 
-const defaultTtl = 50 * time.Second
+const DefaultTtl = 50 * time.Second
 
 type Whisper struct {
 	pub, sec []byte
@@ -43,7 +44,7 @@ type Whisper struct {
 
 	mmu      sync.RWMutex
 	messages map[Hash]*Envelope
-	expiry   map[int32]*set.SetNonTS
+	expiry   map[uint32]*set.SetNonTS
 
 	quit chan struct{}
 }
@@ -53,12 +54,18 @@ func New(pub, sec []byte) *Whisper {
 		pub:      pub,
 		sec:      sec,
 		messages: make(map[Hash]*Envelope),
-		expiry:   make(map[int32]*set.SetNonTS),
+		expiry:   make(map[uint32]*set.SetNonTS),
 		quit:     make(chan struct{}),
 	}
 	go whisper.update()
 
-	whisper.Send(defaultTtl, nil, NewMessage([]byte("Hello world. This is whisper-go")))
+	msg := NewMessage([]byte(fmt.Sprintf("Hello world. This is whisper-go. Incase you're wondering; the time is %v", time.Now())))
+	envelope, _ := msg.Seal(DefaultPow, Opts{
+		Ttl: DefaultTtl,
+	})
+	if err := whisper.Send(envelope); err != nil {
+		fmt.Println(err)
+	}
 
 	// p2p whisper sub protocol handler
 	whisper.protocol = p2p.Protocol{
@@ -75,17 +82,14 @@ func (self *Whisper) Stop() {
 	close(self.quit)
 }
 
-func (self *Whisper) Send(ttl time.Duration, topics [][]byte, data *Message) {
-	envelope := NewEnvelope(ttl, topics, data)
-	envelope.Seal()
-
-	self.add(envelope)
+func (self *Whisper) Send(envelope *Envelope) error {
+	return self.add(envelope)
 }
 
 // Main handler for passing whisper messages to whisper peer objects
 func (self *Whisper) msgHandler(peer *p2p.Peer, ws p2p.MsgReadWriter) error {
 	wpeer := NewPeer(self, peer, ws)
-	// init whisper peer (handshake/status)
+	// initialise whisper peer (handshake/status)
 	if err := wpeer.init(); err != nil {
 		return err
 	}
@@ -106,22 +110,37 @@ func (self *Whisper) msgHandler(peer *p2p.Peer, ws p2p.MsgReadWriter) error {
 			continue
 		}
 
-		self.add(envelope)
+		if err := self.add(envelope); err != nil {
+			// TODO Punish peer here. Invalid envelope.
+			peer.Infoln(err)
+		}
 		wpeer.addKnown(envelope)
 	}
 }
 
 // takes care of adding envelopes to the messages pool. At this moment no sanity checks are being performed.
-func (self *Whisper) add(envelope *Envelope) {
+func (self *Whisper) add(envelope *Envelope) error {
+	if !envelope.valid() {
+		return errors.New("invalid pow for envelope")
+	}
+
 	self.mmu.Lock()
 	defer self.mmu.Unlock()
 
-	fmt.Println("add", envelope)
-	self.messages[envelope.Hash()] = envelope
+	hash := envelope.Hash()
+	self.messages[hash] = envelope
 	if self.expiry[envelope.Expiry] == nil {
 		self.expiry[envelope.Expiry] = set.NewNonTS()
 	}
-	self.expiry[envelope.Expiry].Add(envelope.Hash())
+
+	if !self.expiry[envelope.Expiry].Has(hash) {
+		self.expiry[envelope.Expiry].Add(hash)
+		// TODO notify listeners (given that we had any ...)
+	}
+
+	fmt.Println("add", envelope)
+
+	return nil
 }
 
 func (self *Whisper) update() {
@@ -141,7 +160,7 @@ func (self *Whisper) expire() {
 	self.mmu.Lock()
 	defer self.mmu.Unlock()
 
-	now := int32(time.Now().Unix())
+	now := uint32(time.Now().Unix())
 	for then, hashSet := range self.expiry {
 		if then > now {
 			continue