whisper: fix payload loss in case of plaintext decrypt

2 files changed, 109 insertions, 4 deletions

diff --git a/whisper/envelope_test.go b/whisper/envelope_test.go
index 3117284f1..b64767b2e 100644
--- a/whisper/envelope_test.go
+++ b/whisper/envelope_test.go
@@ -4,6 +4,9 @@ import (
 	"bytes"
 	"testing"
 	"time"
+
+	"github.com/ethereum/go-ethereum/crypto"
+	"github.com/ethereum/go-ethereum/crypto/ecies"
 )
 
 func TestEnvelopeOpen(t *testing.T) {
@@ -16,7 +19,7 @@ func TestEnvelopeOpen(t *testing.T) {
 	}
 	opened, err := envelope.Open(nil)
 	if err != nil {
-		t.Fatalf("failed to open envelope: %v.", err)
+		t.Fatalf("failed to open envelope: %v", err)
 	}
 	if opened.Flags != message.Flags {
 		t.Fatalf("flags mismatch: have %d, want %d", opened.Flags, message.Flags)
@@ -38,3 +41,102 @@ func TestEnvelopeOpen(t *testing.T) {
 		t.Fatalf("message hash mismatch: have 0x%x, want 0x%x", opened.Hash, envelope.Hash())
 	}
 }
+
+func TestEnvelopeAnonymousOpenUntargeted(t *testing.T) {
+	payload := []byte("hello envelope")
+	envelope, err := NewMessage(payload).Wrap(DefaultPoW, Options{})
+	if err != nil {
+		t.Fatalf("failed to wrap message: %v", err)
+	}
+	opened, err := envelope.Open(nil)
+	if err != nil {
+		t.Fatalf("failed to open envelope: %v", err)
+	}
+	if opened.To != nil {
+		t.Fatalf("recipient mismatch: have 0x%x, want nil", opened.To)
+	}
+	if bytes.Compare(opened.Payload, payload) != 0 {
+		t.Fatalf("payload mismatch: have 0x%x, want 0x%x", opened.Payload, payload)
+	}
+}
+
+func TestEnvelopeAnonymousOpenTargeted(t *testing.T) {
+	key, err := crypto.GenerateKey()
+	if err != nil {
+		t.Fatalf("failed to generate test identity: %v", err)
+	}
+
+	payload := []byte("hello envelope")
+	envelope, err := NewMessage(payload).Wrap(DefaultPoW, Options{
+		To: &key.PublicKey,
+	})
+	if err != nil {
+		t.Fatalf("failed to wrap message: %v", err)
+	}
+	opened, err := envelope.Open(nil)
+	if err != nil {
+		t.Fatalf("failed to open envelope: %v", err)
+	}
+	if opened.To != nil {
+		t.Fatalf("recipient mismatch: have 0x%x, want nil", opened.To)
+	}
+	if bytes.Compare(opened.Payload, payload) == 0 {
+		t.Fatalf("payload match, should have been encrypted: 0x%x", opened.Payload)
+	}
+}
+
+func TestEnvelopeIdentifiedOpenUntargeted(t *testing.T) {
+	key, err := crypto.GenerateKey()
+	if err != nil {
+		t.Fatalf("failed to generate test identity: %v", err)
+	}
+
+	payload := []byte("hello envelope")
+	envelope, err := NewMessage(payload).Wrap(DefaultPoW, Options{})
+	if err != nil {
+		t.Fatalf("failed to wrap message: %v", err)
+	}
+	opened, err := envelope.Open(key)
+	switch err {
+	case nil:
+		t.Fatalf("envelope opened with bad key: %v", opened)
+
+	case ecies.ErrInvalidPublicKey:
+		// Ok, key mismatch but opened
+
+	default:
+		t.Fatalf("failed to open envelope: %v", err)
+	}
+
+	if opened.To != nil {
+		t.Fatalf("recipient mismatch: have 0x%x, want nil", opened.To)
+	}
+	if bytes.Compare(opened.Payload, payload) != 0 {
+		t.Fatalf("payload mismatch: have 0x%x, want 0x%x", opened.Payload, payload)
+	}
+}
+
+func TestEnvelopeIdentifiedOpenTargeted(t *testing.T) {
+	key, err := crypto.GenerateKey()
+	if err != nil {
+		t.Fatalf("failed to generate test identity: %v", err)
+	}
+
+	payload := []byte("hello envelope")
+	envelope, err := NewMessage(payload).Wrap(DefaultPoW, Options{
+		To: &key.PublicKey,
+	})
+	if err != nil {
+		t.Fatalf("failed to wrap message: %v", err)
+	}
+	opened, err := envelope.Open(key)
+	if err != nil {
+		t.Fatalf("failed to open envelope: %v", err)
+	}
+	if opened.To != nil {
+		t.Fatalf("recipient mismatch: have 0x%x, want nil", opened.To)
+	}
+	if bytes.Compare(opened.Payload, payload) != 0 {
+		t.Fatalf("payload mismatch: have 0x%x, want 0x%x", opened.Payload, payload)
+	}
+}
diff --git a/whisper/message.go b/whisper/message.go
index 2b92d515c..a80380a92 100644
--- a/whisper/message.go
+++ b/whisper/message.go
@@ -120,9 +120,12 @@ func (self *Message) encrypt(key *ecdsa.PublicKey) (err error) {
 }
 
 // decrypt decrypts an encrypted payload with a private key.
-func (self *Message) decrypt(key *ecdsa.PrivateKey) (err error) {
-	self.Payload, err = crypto.Decrypt(key, self.Payload)
-	return
+func (self *Message) decrypt(key *ecdsa.PrivateKey) error {
+	cleartext, err := crypto.Decrypt(key, self.Payload)
+	if err == nil {
+		self.Payload = cleartext
+	}
+	return err
 }
 
 // hash calculates the SHA3 checksum of the message flags and payload.