whisper: polish the messages, fix some bugs, tests

Bugs fixed:
  - Use randomly generated flags as the spec required.
  - During envelope opening check the first bit only for signature.

6 files changed, 209 insertions, 88 deletions

diff --git a/whisper/envelope.go b/whisper/envelope.go
index 20e3e6d39..65dc89936 100644
--- a/whisper/envelope.go
+++ b/whisper/envelope.go
@@ -40,7 +40,7 @@ func NewEnvelope(ttl time.Duration, topics [][]byte, data *Message) *Envelope {
 		Expiry: uint32(exp.Unix()),
 		TTL:    uint32(ttl.Seconds()),
 		Topics: topics,
-		Data:   data.Bytes(),
+		Data:   data.bytes(),
 		Nonce:  0,
 	}
 }
@@ -49,32 +49,32 @@ func (self *Envelope) Seal(pow time.Duration) {
 	self.proveWork(pow)
 }
 
-func (self *Envelope) Open(prv *ecdsa.PrivateKey) (msg *Message, err error) {
+func (self *Envelope) Open(key *ecdsa.PrivateKey) (msg *Message, err error) {
 	data := self.Data
-	var message Message
-	dataStart := 1
-	if data[0] > 0 {
-		if len(data) < 66 {
-			return nil, fmt.Errorf("unable to open envelope. First bit set but len(data) < 66")
+
+	message := Message{
+		Flags: data[0],
+	}
+	data = data[1:]
+
+	if message.Flags&128 == 128 {
+		if len(data) < 65 {
+			return nil, fmt.Errorf("unable to open envelope. First bit set but len(data) < 65")
 		}
-		dataStart = 66
-		message.Flags = data[0]
-		message.Signature = data[1:66]
+		message.Signature, data = data[:65], data[65:]
 	}
+	message.Payload = data
 
-	payload := data[dataStart:]
-	if prv != nil {
-		message.Payload, err = crypto.Decrypt(prv, payload)
+	if key != nil {
+		message.Payload, err = crypto.Decrypt(key, message.Payload)
 		switch err {
 		case nil: // OK
 		case ecies.ErrInvalidPublicKey: // Payload isn't encrypted
-			message.Payload = payload
 			return &message, err
 		default:
 			return nil, fmt.Errorf("unable to open envelope. Decrypt failed: %v", err)
 		}
 	}
-
 	return &message, nil
 }
 
diff --git a/whisper/main.go b/whisper/main.go
index 643644ee6..6a089f894 100644
--- a/whisper/main.go
+++ b/whisper/main.go
@@ -69,10 +69,10 @@ func selfSend(shh *whisper.Whisper, payload []byte) error {
 	})
 	// Wrap the payload and encrypt it
 	msg := whisper.NewMessage(payload)
-	envelope, err := msg.Seal(whisper.DefaultPow, whisper.Opts{
-		Ttl:  whisper.DefaultTtl,
+	envelope, err := msg.Wrap(whisper.DefaultPow, whisper.Options{
 		From: id,
 		To:   &id.PublicKey,
+		TTL:  whisper.DefaultTimeToLive,
 	})
 	if err != nil {
 		return fmt.Errorf("failed to seal message: %v", err)
diff --git a/whisper/message.go b/whisper/message.go
index ad6a1bcff..3bee83f39 100644
--- a/whisper/message.go
+++ b/whisper/message.go
@@ -1,7 +1,11 @@
+// Contains the Whisper protocol Message element. For formal details please see
+// the specs at https://github.com/ethereum/wiki/wiki/Whisper-PoC-1-Protocol-Spec#messages.
+
 package whisper
 
 import (
 	"crypto/ecdsa"
+	"math/rand"
 	"time"
 
 	"github.com/ethereum/go-ethereum/crypto"
@@ -9,8 +13,11 @@ import (
 	"github.com/ethereum/go-ethereum/logger/glog"
 )
 
+// Message represents an end-user data packet to trasmit through the Whisper
+// protocol. These are wrapped into Envelopes that need not be understood by
+// intermediate nodes, just forwarded.
 type Message struct {
-	Flags     byte
+	Flags     byte // First bit it signature presence, rest reserved and should be random
 	Signature []byte
 	Payload   []byte
 	Sent      int64
@@ -18,71 +25,95 @@ type Message struct {
 	To *ecdsa.PublicKey
 }
 
+// Options specifies the exact way a message should be wrapped into an Envelope.
+type Options struct {
+	From   *ecdsa.PrivateKey
+	To     *ecdsa.PublicKey
+	TTL    time.Duration
+	Topics [][]byte
+}
+
+// NewMessage creates and initializes a non-signed, non-encrypted Whisper message.
 func NewMessage(payload []byte) *Message {
-	return &Message{Flags: 0, Payload: payload, Sent: time.Now().Unix()}
+	// Construct an initial flag set: bit #1 = 0 (no signature), rest random
+	flags := byte(rand.Intn(128))
+
+	// Assemble and return the message
+	return &Message{
+		Flags:   flags,
+		Payload: payload,
+		Sent:    time.Now().Unix(),
+	}
 }
 
-func (self *Message) hash() []byte {
-	return crypto.Sha3(append([]byte{self.Flags}, self.Payload...))
+// Wrap bundles the message into an Envelope to transmit over the network.
+//
+// Pov (Proof Of Work) controls how much time to spend on hashing the message,
+// inherently controlling its priority through the network (smaller hash, bigger
+// priority).
+//
+// The user can control the amount of identity, privacy and encryption through
+// the options parameter as follows:
+//   - options.From == nil && options.To == nil: anonymous broadcast
+//   - options.From != nil && options.To == nil: signed broadcast (known sender)
+//   - options.From == nil && options.To != nil: encrypted anonymous message
+//   - options.From != nil && options.To != nil: encrypted signed message
+func (self *Message) Wrap(pow time.Duration, options Options) (*Envelope, error) {
+	// Use the default TTL if non was specified
+	if options.TTL == 0 {
+		options.TTL = DefaultTimeToLive
+	}
+	// Sign and encrypt the message if requested
+	if options.From != nil {
+		if err := self.sign(options.From); err != nil {
+			return nil, err
+		}
+	}
+	if options.To != nil {
+		if err := self.encrypt(options.To); err != nil {
+			return nil, err
+		}
+	}
+	// Wrap the processed message, seal it and return
+	envelope := NewEnvelope(options.TTL, options.Topics, self)
+	envelope.Seal(pow)
+
+	return envelope, nil
 }
 
+// Sign calculates and sets the cryptographic signature for the message , also
+// setting the sign flag.
 func (self *Message) sign(key *ecdsa.PrivateKey) (err error) {
-	self.Flags = 1
+	self.Flags |= 1 << 7
 	self.Signature, err = crypto.Sign(self.hash(), key)
 	return
 }
 
+// Recover retrieves the public key of the message signer.
 func (self *Message) Recover() *ecdsa.PublicKey {
-	defer func() { recover() }() // in case of invalid sig
+	defer func() { recover() }() // in case of invalid signature
+
 	pub, err := crypto.SigToPub(self.hash(), self.Signature)
 	if err != nil {
-		glog.V(logger.Error).Infof("Could not get pubkey from signature: ", err)
+		glog.V(logger.Error).Infof("Could not get public key from signature: %v", err)
 		return nil
 	}
 	return pub
 }
 
-func (self *Message) Encrypt(to *ecdsa.PublicKey) (err error) {
+// Encrypt encrypts a message payload with a public key.
+func (self *Message) encrypt(to *ecdsa.PublicKey) (err error) {
 	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...)...)
+	return
 }
 
-type Opts struct {
-	From   *ecdsa.PrivateKey
-	To     *ecdsa.PublicKey
-	Ttl    time.Duration
-	Topics [][]byte
+// Hash calculates the SHA3 checksum of the message flags and payload.
+func (self *Message) hash() []byte {
+	return crypto.Sha3(append([]byte{self.Flags}, self.Payload...))
 }
 
-func (self *Message) Seal(pow time.Duration, opts Opts) (*Envelope, error) {
-	if opts.From != nil {
-		err := self.sign(opts.From)
-		if err != nil {
-			return nil, err
-		}
-	}
-
-	if opts.To != nil {
-		err := self.Encrypt(opts.To)
-		if err != nil {
-			return nil, err
-		}
-	}
-
-	if opts.Ttl == 0 {
-		opts.Ttl = DefaultTtl
-	}
-
-	envelope := NewEnvelope(opts.Ttl, opts.Topics, self)
-	envelope.Seal(pow)
-
-	return envelope, nil
+// Bytes flattens the message contents (flags, signature and payload) into a
+// single binary blob.
+func (self *Message) bytes() []byte {
+	return append([]byte{self.Flags}, append(self.Signature, self.Payload...)...)
 }
diff --git a/whisper/message_test.go b/whisper/message_test.go
index 93caa31b3..ed24d37d8 100644
--- a/whisper/message_test.go
+++ b/whisper/message_test.go
@@ -3,48 +3,136 @@ package whisper
 import (
 	"bytes"
 	"crypto/elliptic"
-	"fmt"
 	"testing"
 
 	"github.com/ethereum/go-ethereum/crypto"
 )
 
-func TestSign(t *testing.T) {
-	prv, _ := crypto.GenerateKey()
-	msg := NewMessage([]byte("hello world"))
-	msg.sign(prv)
+// Tests whether a message can be wrapped without any identity or encryption.
+func TestMessageSimpleWrap(t *testing.T) {
+	payload := []byte("hello world")
+
+	msg := NewMessage(payload)
+	if _, err := msg.Wrap(DefaultPow, Options{}); err != nil {
+		t.Fatalf("failed to wrap message: %v", err)
+	}
+	if msg.Flags&128 != 0 {
+		t.Fatalf("signature flag mismatch: have %d, want %d", (msg.Flags&128)>>7, 0)
+	}
+	if len(msg.Signature) != 0 {
+		t.Fatalf("signature found for simple wrapping: 0x%x", msg.Signature)
+	}
+	if bytes.Compare(msg.Payload, payload) != 0 {
+		t.Fatalf("payload mismatch after wrapping: have 0x%x, want 0x%x", msg.Payload, payload)
+	}
+}
+
+// Tests whether a message can be signed, and wrapped in plain-text.
+func TestMessageCleartextSignRecover(t *testing.T) {
+	key, err := crypto.GenerateKey()
+	if err != nil {
+		t.Fatalf("failed to create crypto key: %v", err)
+	}
+	payload := []byte("hello world")
+
+	msg := NewMessage(payload)
+	if _, err := msg.Wrap(DefaultPow, Options{
+		From: key,
+	}); err != nil {
+		t.Fatalf("failed to sign message: %v", err)
+	}
+	if msg.Flags&128 != 128 {
+		t.Fatalf("signature flag mismatch: have %d, want %d", (msg.Flags&128)>>7, 1)
+	}
+	if bytes.Compare(msg.Payload, payload) != 0 {
+		t.Fatalf("payload mismatch after signing: have 0x%x, want 0x%x", msg.Payload, payload)
+	}
 
 	pubKey := msg.Recover()
-	p1 := elliptic.Marshal(crypto.S256(), prv.PublicKey.X, prv.PublicKey.Y)
+	if pubKey == nil {
+		t.Fatalf("failed to recover public key")
+	}
+	p1 := elliptic.Marshal(crypto.S256(), key.PublicKey.X, key.PublicKey.Y)
 	p2 := elliptic.Marshal(crypto.S256(), pubKey.X, pubKey.Y)
-
 	if !bytes.Equal(p1, p2) {
-		t.Error("recovered pub key did not match")
+		t.Fatalf("public key mismatch: have 0x%x, want 0x%x", p2, p1)
 	}
 }
 
-func TestMessageEncryptDecrypt(t *testing.T) {
-	prv1, _ := crypto.GenerateKey()
-	prv2, _ := crypto.GenerateKey()
+// Tests whether a message can be encrypted and decrypted using an anonymous
+// sender (i.e. no signature).
+func TestMessageAnonymousEncryptDecrypt(t *testing.T) {
+	key, err := crypto.GenerateKey()
+	if err != nil {
+		t.Fatalf("failed to create recipient crypto key: %v", err)
+	}
+	payload := []byte("hello world")
 
-	data := []byte("hello world")
-	msg := NewMessage(data)
-	envelope, err := msg.Seal(DefaultPow, Opts{
-		From: prv1,
-		To:   &prv2.PublicKey,
+	msg := NewMessage(payload)
+	envelope, err := msg.Wrap(DefaultPow, Options{
+		To: &key.PublicKey,
 	})
 	if err != nil {
-		fmt.Println(err)
-		t.FailNow()
+		t.Fatalf("failed to encrypt message: %v", err)
+	}
+	if msg.Flags&128 != 0 {
+		t.Fatalf("signature flag mismatch: have %d, want %d", (msg.Flags&128)>>7, 0)
+	}
+	if len(msg.Signature) != 0 {
+		t.Fatalf("signature found for anonymous message: 0x%x", msg.Signature)
 	}
 
-	msg1, err := envelope.Open(prv2)
+	out, err := envelope.Open(key)
 	if err != nil {
-		t.Error(err)
-		t.FailNow()
+		t.Fatalf("failed to open encrypted message: %v", err)
+	}
+	if !bytes.Equal(out.Payload, payload) {
+		t.Error("payload mismatch: have 0x%x, want 0x%x", out.Payload, payload)
+	}
+}
+
+// Tests whether a message can be properly signed and encrypted.
+func TestMessageFullCrypto(t *testing.T) {
+	fromKey, err := crypto.GenerateKey()
+	if err != nil {
+		t.Fatalf("failed to create sender crypto key: %v", err)
+	}
+	toKey, err := crypto.GenerateKey()
+	if err != nil {
+		t.Fatalf("failed to create recipient crypto key: %v", err)
 	}
 
-	if !bytes.Equal(msg1.Payload, data) {
-		t.Error("encryption error. data did not match")
+	payload := []byte("hello world")
+	msg := NewMessage(payload)
+	envelope, err := msg.Wrap(DefaultPow, Options{
+		From: fromKey,
+		To:   &toKey.PublicKey,
+	})
+	if err != nil {
+		t.Fatalf("failed to encrypt message: %v", err)
+	}
+	if msg.Flags&128 != 128 {
+		t.Fatalf("signature flag mismatch: have %d, want %d", (msg.Flags&128)>>7, 1)
+	}
+	if len(msg.Signature) == 0 {
+		t.Fatalf("no signature found for signed message")
+	}
+
+	out, err := envelope.Open(toKey)
+	if err != nil {
+		t.Fatalf("failed to open encrypted message: %v", err)
+	}
+	if !bytes.Equal(out.Payload, payload) {
+		t.Error("payload mismatch: have 0x%x, want 0x%x", out.Payload, payload)
+	}
+
+	pubKey := out.Recover()
+	if pubKey == nil {
+		t.Fatalf("failed to recover public key")
+	}
+	p1 := elliptic.Marshal(crypto.S256(), fromKey.PublicKey.X, fromKey.PublicKey.Y)
+	p2 := elliptic.Marshal(crypto.S256(), pubKey.X, pubKey.Y)
+	if !bytes.Equal(p1, p2) {
+		t.Fatalf("public key mismatch: have 0x%x, want 0x%x", p2, p1)
 	}
 }
diff --git a/whisper/whisper.go b/whisper/whisper.go
index 00dcb1932..3cb92a07a 100644
--- a/whisper/whisper.go
+++ b/whisper/whisper.go
@@ -28,7 +28,9 @@ type MessageEvent struct {
 	Message *Message
 }
 
-const DefaultTtl = 50 * time.Second
+const (
+	DefaultTimeToLive = 50 * time.Second
+)
 
 type Whisper struct {
 	protocol p2p.Protocol
diff --git a/whisper/whisper_test.go b/whisper/whisper_test.go
index 3e3945a0a..a3e0e03d2 100644
--- a/whisper/whisper_test.go
+++ b/whisper/whisper_test.go
@@ -18,8 +18,8 @@ func TestEvent(t *testing.T) {
 	})
 
 	msg := NewMessage([]byte(fmt.Sprintf("Hello world. This is whisper-go. Incase you're wondering; the time is %v", time.Now())))
-	envelope, err := msg.Seal(DefaultPow, Opts{
-		Ttl:  DefaultTtl,
+	envelope, err := msg.Wrap(DefaultPow, Options{
+		TTL:  DefaultTimeToLive,
 		From: id,
 		To:   &id.PublicKey,
 	})