whisper: separate out magic number from the code

4 files changed, 29 insertions, 25 deletions

diff --git a/whisper/envelope.go b/whisper/envelope.go
index f35a40a42..c51c6e600 100644
--- a/whisper/envelope.go
+++ b/whisper/envelope.go
@@ -11,7 +11,6 @@ import (
 
 	"github.com/ethereum/go-ethereum/common"
 	"github.com/ethereum/go-ethereum/crypto"
-	"github.com/ethereum/go-ethereum/crypto/ecies"
 	"github.com/ethereum/go-ethereum/rlp"
 )
 
@@ -85,27 +84,22 @@ func (self *Envelope) Open(key *ecdsa.PrivateKey) (msg *Message, err error) {
 	}
 	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")
+	if message.Flags&signatureFlag == signatureFlag {
+		if len(data) < signatureLength {
+			return nil, fmt.Errorf("unable to open envelope. First bit set but len(data) < len(signature)")
 		}
-		message.Signature, data = data[:65], data[65:]
+		message.Signature, data = data[:signatureLength], data[signatureLength:]
 	}
 	message.Payload = data
 
-	// Short circuit if the encryption was requested
+	// Decrypt the message, if requested
 	if key == nil {
 		return message, nil
 	}
-	// Otherwise try to decrypt the message
-	message.Payload, err = crypto.Decrypt(key, message.Payload)
-	switch err {
+	switch message.decrypt(key) {
 	case nil:
 		return message, nil
 
-	case ecies.ErrInvalidPublicKey: // Payload isn't encrypted
-		return message, err
-
 	default:
 		return nil, fmt.Errorf("unable to open envelope, decrypt failed: %v", err)
 	}
diff --git a/whisper/message.go b/whisper/message.go
index 2666ee6e0..457cf6def 100644
--- a/whisper/message.go
+++ b/whisper/message.go
@@ -35,8 +35,9 @@ type Options struct {
 
 // NewMessage creates and initializes a non-signed, non-encrypted Whisper message.
 func NewMessage(payload []byte) *Message {
-	// Construct an initial flag set: bit #1 = 0 (no signature), rest random
-	flags := byte(rand.Intn(128))
+	// Construct an initial flag set: no signature, rest random
+	flags := byte(rand.Intn(256))
+	flags &= ^signatureFlag
 
 	// Assemble and return the message
 	return &Message{
@@ -84,7 +85,7 @@ func (self *Message) Wrap(pow time.Duration, options Options) (*Envelope, error)
 // 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 << 7
+	self.Flags |= signatureFlag
 	self.Signature, err = crypto.Sign(self.hash(), key)
 	return
 }
@@ -102,8 +103,14 @@ func (self *Message) Recover() *ecdsa.PublicKey {
 }
 
 // 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)
+func (self *Message) encrypt(key *ecdsa.PublicKey) (err error) {
+	self.Payload, err = crypto.Encrypt(key, self.Payload)
+	return
+}
+
+// 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
 }
 
diff --git a/whisper/message_test.go b/whisper/message_test.go
index 8d4c5e990..319bc6025 100644
--- a/whisper/message_test.go
+++ b/whisper/message_test.go
@@ -16,8 +16,8 @@ func TestMessageSimpleWrap(t *testing.T) {
 	if _, err := msg.Wrap(DefaultProofOfWork, 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 msg.Flags&signatureFlag != 0 {
+		t.Fatalf("signature flag mismatch: have %d, want %d", msg.Flags&signatureFlag, 0)
 	}
 	if len(msg.Signature) != 0 {
 		t.Fatalf("signature found for simple wrapping: 0x%x", msg.Signature)
@@ -41,8 +41,8 @@ func TestMessageCleartextSignRecover(t *testing.T) {
 	}); 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 msg.Flags&signatureFlag != signatureFlag {
+		t.Fatalf("signature flag mismatch: have %d, want %d", msg.Flags&signatureFlag, signatureFlag)
 	}
 	if bytes.Compare(msg.Payload, payload) != 0 {
 		t.Fatalf("payload mismatch after signing: have 0x%x, want 0x%x", msg.Payload, payload)
@@ -75,8 +75,8 @@ func TestMessageAnonymousEncryptDecrypt(t *testing.T) {
 	if err != nil {
 		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 msg.Flags&signatureFlag != 0 {
+		t.Fatalf("signature flag mismatch: have %d, want %d", msg.Flags&signatureFlag, 0)
 	}
 	if len(msg.Signature) != 0 {
 		t.Fatalf("signature found for anonymous message: 0x%x", msg.Signature)
@@ -111,8 +111,8 @@ func TestMessageFullCrypto(t *testing.T) {
 	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 msg.Flags&signatureFlag != signatureFlag {
+		t.Fatalf("signature flag mismatch: have %d, want %d", msg.Flags&signatureFlag, signatureFlag)
 	}
 	if len(msg.Signature) == 0 {
 		t.Fatalf("no signature found for signed message")
diff --git a/whisper/whisper.go b/whisper/whisper.go
index d803e27d4..ad29fe16a 100644
--- a/whisper/whisper.go
+++ b/whisper/whisper.go
@@ -20,6 +20,9 @@ const (
 	statusMsg      = 0x0
 	envelopesMsg   = 0x01
 	whisperVersion = 0x02
+
+	signatureFlag   = byte(1 << 7)
+	signatureLength = 65
 )
 
 type MessageEvent struct {