whisper: polish the Envelope a bit, prep for tests.

6 files changed, 75 insertions, 65 deletions

diff --git a/whisper/envelope.go b/whisper/envelope.go
index 65dc89936..6ad48dd25 100644
--- a/whisper/envelope.go
+++ b/whisper/envelope.go
@@ -1,3 +1,6 @@
+// Contains the Whisper protocol Envelope element. For formal details please see
+// the specs at https://github.com/ethereum/wiki/wiki/Whisper-PoC-1-Protocol-Spec#envelopes.
+
 package whisper
 
 import (
@@ -12,10 +15,8 @@ import (
 	"github.com/ethereum/go-ethereum/rlp"
 )
 
-const (
-	DefaultPow = 50 * time.Millisecond
-)
-
+// Envelope represents a clear-text data packet to transmit through the Whisper
+// network. Its contents may or may not be encrypted and signed.
 type Envelope struct {
 	Expiry uint32 // Whisper protocol specifies int32, really should be int64
 	TTL    uint32 // ^^^^^^
@@ -26,33 +27,60 @@ type Envelope struct {
 	hash common.Hash
 }
 
-func (self *Envelope) Hash() common.Hash {
-	if (self.hash == common.Hash{}) {
-		enc, _ := rlp.EncodeToBytes(self)
-		self.hash = crypto.Sha3Hash(enc)
-	}
-	return self.hash
-}
-
-func NewEnvelope(ttl time.Duration, topics [][]byte, data *Message) *Envelope {
-	exp := time.Now().Add(ttl)
+// NewEnvelope wraps a Whisper message with expiration and destination data
+// included into an envelope for network forwarding.
+func NewEnvelope(ttl time.Duration, topics [][]byte, msg *Message) *Envelope {
 	return &Envelope{
-		Expiry: uint32(exp.Unix()),
+		Expiry: uint32(time.Now().Add(ttl).Unix()),
 		TTL:    uint32(ttl.Seconds()),
 		Topics: topics,
-		Data:   data.bytes(),
+		Data:   msg.bytes(),
 		Nonce:  0,
 	}
 }
 
+// Seal closes the envelope by spending the requested amount of time as a proof
+// of work on hashing the data.
 func (self *Envelope) Seal(pow time.Duration) {
-	self.proveWork(pow)
+	d := make([]byte, 64)
+	copy(d[:32], self.rlpWithoutNonce())
+
+	finish, bestBit := time.Now().Add(pow).UnixNano(), 0
+	for nonce := uint32(0); time.Now().UnixNano() < finish; {
+		for i := 0; i < 1024; i++ {
+			binary.BigEndian.PutUint32(d[60:], nonce)
+
+			firstBit := common.FirstBitSet(common.BigD(crypto.Sha3(d)))
+			if firstBit > bestBit {
+				self.Nonce, bestBit = nonce, firstBit
+			}
+			nonce++
+		}
+	}
 }
 
+// Valid checks whether the claimed proof of work was indeed executed.
+// TODO: Is this really useful? Isn't this always true?
+func (self *Envelope) valid() bool {
+	d := make([]byte, 64)
+	copy(d[:32], self.rlpWithoutNonce())
+	binary.BigEndian.PutUint32(d[60:], self.Nonce)
+
+	return common.FirstBitSet(common.BigD(crypto.Sha3(d))) > 0
+}
+
+// RlpWithoutNonce returns the RLP encoded envelope contents, except the nonce.
+func (self *Envelope) rlpWithoutNonce() []byte {
+	enc, _ := rlp.EncodeToBytes([]interface{}{self.Expiry, self.TTL, self.Topics, self.Data})
+	return enc
+}
+
+// Open extracts the message contained within a potentially encrypted envelope.
 func (self *Envelope) Open(key *ecdsa.PrivateKey) (msg *Message, err error) {
+	// Split open the payload into a message construct
 	data := self.Data
 
-	message := Message{
+	message := &Message{
 		Flags: data[0],
 	}
 	data = data[1:]
@@ -65,57 +93,38 @@ func (self *Envelope) Open(key *ecdsa.PrivateKey) (msg *Message, err error) {
 	}
 	message.Payload = data
 
-	if key != nil {
-		message.Payload, err = crypto.Decrypt(key, message.Payload)
-		switch err {
-		case nil: // OK
-		case ecies.ErrInvalidPublicKey: // Payload isn't encrypted
-			return &message, err
-		default:
-			return nil, fmt.Errorf("unable to open envelope. Decrypt failed: %v", err)
-		}
+	// Short circuit if the encryption was requested
+	if key == nil {
+		return message, nil
 	}
-	return &message, nil
-}
-
-func (self *Envelope) proveWork(dura time.Duration) {
-	var bestBit int
-	d := make([]byte, 64)
-	enc, _ := rlp.EncodeToBytes(self.withoutNonce())
-	copy(d[:32], enc)
-
-	then := time.Now().Add(dura).UnixNano()
-	for n := uint32(0); time.Now().UnixNano() < then; {
-		for i := 0; i < 1024; i++ {
-			binary.BigEndian.PutUint32(d[60:], n)
+	// Otherwise try to decrypt the message
+	message.Payload, err = crypto.Decrypt(key, message.Payload)
+	switch err {
+	case nil:
+		return message, nil
 
-			fbs := common.FirstBitSet(common.BigD(crypto.Sha3(d)))
-			if fbs > bestBit {
-				bestBit = fbs
-				self.Nonce = n
-			}
+	case ecies.ErrInvalidPublicKey: // Payload isn't encrypted
+		return message, err
 
-			n++
-		}
+	default:
+		return nil, fmt.Errorf("unable to open envelope, decrypt failed: %v", err)
 	}
 }
 
-func (self *Envelope) valid() bool {
-	d := make([]byte, 64)
-	enc, _ := rlp.EncodeToBytes(self.withoutNonce())
-	copy(d[:32], enc)
-	binary.BigEndian.PutUint32(d[60:], self.Nonce)
-	return common.FirstBitSet(common.BigD(crypto.Sha3(d))) > 0
-}
-
-func (self *Envelope) withoutNonce() interface{} {
-	return []interface{}{self.Expiry, self.TTL, self.Topics, self.Data}
+// Hash returns the SHA3 hash of the envelope, calculating it if not yet done.
+func (self *Envelope) Hash() common.Hash {
+	if (self.hash == common.Hash{}) {
+		enc, _ := rlp.EncodeToBytes(self)
+		self.hash = crypto.Sha3Hash(enc)
+	}
+	return self.hash
 }
 
 // rlpenv is an Envelope but is not an rlp.Decoder.
 // It is used for decoding because we need to
 type rlpenv Envelope
 
+// DecodeRLP decodes an Envelope from an RLP data stream.
 func (self *Envelope) DecodeRLP(s *rlp.Stream) error {
 	raw, err := s.Raw()
 	if err != nil {
diff --git a/whisper/main.go b/whisper/main.go
index 6a089f894..422f0fa3b 100644
--- a/whisper/main.go
+++ b/whisper/main.go
@@ -69,7 +69,7 @@ func selfSend(shh *whisper.Whisper, payload []byte) error {
 	})
 	// Wrap the payload and encrypt it
 	msg := whisper.NewMessage(payload)
-	envelope, err := msg.Wrap(whisper.DefaultPow, whisper.Options{
+	envelope, err := msg.Wrap(whisper.DefaultProofOfWork, whisper.Options{
 		From: id,
 		To:   &id.PublicKey,
 		TTL:  whisper.DefaultTimeToLive,
diff --git a/whisper/message.go b/whisper/message.go
index 3bee83f39..255352ee8 100644
--- a/whisper/message.go
+++ b/whisper/message.go
@@ -17,7 +17,7 @@ import (
 // protocol. These are wrapped into Envelopes that need not be understood by
 // intermediate nodes, just forwarded.
 type Message struct {
-	Flags     byte // First bit it signature presence, rest reserved and should be random
+	Flags     byte // First bit is signature presence, rest reserved and should be random
 	Signature []byte
 	Payload   []byte
 	Sent      int64
diff --git a/whisper/message_test.go b/whisper/message_test.go
index ed24d37d8..8d4c5e990 100644
--- a/whisper/message_test.go
+++ b/whisper/message_test.go
@@ -13,7 +13,7 @@ func TestMessageSimpleWrap(t *testing.T) {
 	payload := []byte("hello world")
 
 	msg := NewMessage(payload)
-	if _, err := msg.Wrap(DefaultPow, Options{}); err != nil {
+	if _, err := msg.Wrap(DefaultProofOfWork, Options{}); err != nil {
 		t.Fatalf("failed to wrap message: %v", err)
 	}
 	if msg.Flags&128 != 0 {
@@ -36,7 +36,7 @@ func TestMessageCleartextSignRecover(t *testing.T) {
 	payload := []byte("hello world")
 
 	msg := NewMessage(payload)
-	if _, err := msg.Wrap(DefaultPow, Options{
+	if _, err := msg.Wrap(DefaultProofOfWork, Options{
 		From: key,
 	}); err != nil {
 		t.Fatalf("failed to sign message: %v", err)
@@ -69,7 +69,7 @@ func TestMessageAnonymousEncryptDecrypt(t *testing.T) {
 	payload := []byte("hello world")
 
 	msg := NewMessage(payload)
-	envelope, err := msg.Wrap(DefaultPow, Options{
+	envelope, err := msg.Wrap(DefaultProofOfWork, Options{
 		To: &key.PublicKey,
 	})
 	if err != nil {
@@ -104,7 +104,7 @@ func TestMessageFullCrypto(t *testing.T) {
 
 	payload := []byte("hello world")
 	msg := NewMessage(payload)
-	envelope, err := msg.Wrap(DefaultPow, Options{
+	envelope, err := msg.Wrap(DefaultProofOfWork, Options{
 		From: fromKey,
 		To:   &toKey.PublicKey,
 	})
diff --git a/whisper/whisper.go b/whisper/whisper.go
index 3cb92a07a..d803e27d4 100644
--- a/whisper/whisper.go
+++ b/whisper/whisper.go
@@ -29,7 +29,8 @@ type MessageEvent struct {
 }
 
 const (
-	DefaultTimeToLive = 50 * time.Second
+	DefaultTimeToLive  = 50 * time.Second
+	DefaultProofOfWork = 50 * time.Millisecond
 )
 
 type Whisper struct {
diff --git a/whisper/whisper_test.go b/whisper/whisper_test.go
index a3e0e03d2..b29e34a5e 100644
--- a/whisper/whisper_test.go
+++ b/whisper/whisper_test.go
@@ -18,7 +18,7 @@ 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.Wrap(DefaultPow, Options{
+	envelope, err := msg.Wrap(DefaultProofOfWork, Options{
 		TTL:  DefaultTimeToLive,
 		From: id,
 		To:   &id.PublicKey,