whisper: fix payload loss in case of plaintext decrypt

1 files changed, 103 insertions, 1 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)
+	}
+}