Refactor so both Messenger and MessageHandler could use code from a shsared source

protocol/identity_images.go

@@ -0,0 +1,100 @@
+package protocol
+
+import (
+	"crypto/ecdsa"
+	crand "crypto/rand"
+	"errors"
+
+	"github.com/status-im/status-go/protocol/common"
+	"github.com/status-im/status-go/protocol/protobuf"
+)
+
+func EncryptIdentityImagesWithContactPubKeys(iis map[string]*protobuf.IdentityImage, m *Messenger) error {
+	// Make AES key
+	AESKey := make([]byte, 32)
+	_, err := crand.Read(AESKey)
+	if err != nil {
+		return err
+	}
+
+	for _, ii := range iis {
+		// Encrypt image payload with the AES key
+		encryptedPayload, err := common.Encrypt(ii.Payload, AESKey, crand.Reader)
+		if err != nil {
+			return err
+		}
+
+		// Overwrite the unencrypted payload with the newly encrypted payload
+		ii.Payload = encryptedPayload
+		ii.Encrypted = true
+		for _, c := range m.allContacts {
+			if !c.IsAdded() {
+				continue
+			}
+
+			pubK, err := c.PublicKey()
+			if err != nil {
+				return err
+			}
+			// Generate a Diffie-Helman (DH) between the sender private key and the recipient's public key
+			sharedKey, err := common.MakeECDHSharedKey(m.identity, pubK)
+			if err != nil {
+				return err
+			}
+
+			// Encrypt the main AES key with AES encryption using the DH key
+			eAESKey, err := common.Encrypt(AESKey, sharedKey, crand.Reader)
+			if err != nil {
+				return err
+			}
+
+			// Append the the encrypted main AES key to the IdentityImage's EncryptionKeys slice.
+			ii.EncryptionKeys = append(ii.EncryptionKeys, eAESKey)
+		}
+	}
+
+	return nil
+}
+
+func DecryptIdentityImagesWithIdentityPrivateKey(iis map[string]*protobuf.IdentityImage, recipientIdentity *ecdsa.PrivateKey, senderPubKey *ecdsa.PublicKey) error {
+image:
+	for _, ii := range iis {
+		for _, empk := range ii.EncryptionKeys {
+			// Generate a Diffie-Helman (DH) between the recipient's private key and the sender's public key
+			sharedKey, err := common.MakeECDHSharedKey(recipientIdentity, senderPubKey)
+			if err != nil {
+				return err
+			}
+
+			// Decrypt the main encryption AES key with AES encryption using the DH key
+			dAESKey, err := common.Decrypt(empk, sharedKey)
+			if err != nil {
+				if err.Error() == "cipher: message authentication failed" {
+					continue
+				}
+				return err
+			}
+			if dAESKey == nil{
+				return errors.New("decrypting the payload encryption key resulted in no error and a nil key")
+			}
+
+			// Decrypt the payload with the newly decrypted main encryption AES key
+			payload, err := common.Decrypt(ii.Payload, dAESKey)
+			if err != nil {
+				return err
+			}
+			if payload == nil {
+				// TODO should this be a logger warn? A payload could theoretically be validly empty
+				return errors.New("decrypting the payload resulted in no error and a nil payload")
+			}
+
+			// Overwrite the payload with the decrypted data
+			ii.Payload = payload
+			ii.Encrypted = false
+			continue image
+		}
+	}
+
+	return nil
+}
+

protocol/messenger.go

@@ -4,7 +4,6 @@ import (
 	"bytes"
 	"context"
 	"crypto/ecdsa"
-	crand "crypto/rand"
 	"database/sql"
 	"encoding/hex"
 	"fmt"
@@ -820,7 +819,7 @@ func (m *Messenger) attachIdentityImagesToChatIdentity(context chatContext, ci *
 	}
 
 	if s.ProfilePicturesShowTo == accounts.ProfilePicturesShowToContactsOnly {
-		err := m.encryptIdentityImagesWithContactPubKeys(ci.Images)
+		err := EncryptIdentityImagesWithContactPubKeys(ci.Images, m)
 		if err != nil {
 			return err
 		}
@@ -829,101 +828,6 @@ func (m *Messenger) attachIdentityImagesToChatIdentity(context chatContext, ci *
 	return nil
 }
 
-func (m *Messenger) encryptIdentityImagesWithContactPubKeys(iis map[string]*protobuf.IdentityImage) (err error) {
-	// Make AES key
-	AESKey := make([]byte, 32)
-	_, err = crand.Read(AESKey)
-	if err != nil {
-		return err
-	}
-
-	for _, ii := range iis {
-		// Encrypt image payload with the AES key
-		encryptedPayload, err := common.Encrypt(ii.Payload, AESKey, crand.Reader)
-		if err != nil {
-			return err
-		}
-
-		// Overwrite the unencrypted payload with the newly encrypted payload
-		ii.Payload = encryptedPayload
-		ii.Encrypted = true
-		m.allContacts.Range(func(contactID string, c *Contact) (shouldContinue bool) {
-			if !c.IsAdded() {
-				return false
-			}
-
-			pubK, err := c.PublicKey()
-			if err != nil {
-				return false
-			}
-			// Generate a Diffie-Helman (DH) between the sender private key and the recipient's public key
-			sharedKey, err := common.MakeECDHSharedKey(m.identity, pubK)
-			if err != nil {
-				return false
-			}
-
-			// Encrypt the main AES key with AES encryption using the DH key
-			eAESKey, err := common.Encrypt(AESKey, sharedKey, crand.Reader)
-			if err != nil {
-				return false
-			}
-
-			// Append the the encrypted main AES key to the IdentityImage's EncryptionKeys slice.
-			ii.EncryptionKeys = append(ii.EncryptionKeys, eAESKey)
-
-			return true
-		})
-
-		if err != nil {
-			return err
-		}
-	}
-
-	return nil
-}
-
-func (m *Messenger) decryptIdentityImagesWithIdentityPrivateKey(iis map[string]*protobuf.IdentityImage, senderPubKey *ecdsa.PublicKey) error {
-image:
-	for _, ii := range iis {
-		for _, empk := range ii.EncryptionKeys {
-			// Generate a Diffie-Helman (DH) between the recipient's private key and the sender's public key
-			sharedKey, err := common.MakeECDHSharedKey(m.identity, senderPubKey)
-			if err != nil {
-				return err
-			}
-
-			// Decrypt the main encryption AES key with AES encryption using the DH key
-			dAESKey, err := common.Decrypt(empk, sharedKey)
-			if err != nil {
-				if err.Error() == "cipher: message authentication failed" {
-					continue
-				}
-				return err
-			}
-			if dAESKey == nil{
-				return errors.New("decrypting the payload encryption key resulted in no error and a nil key")
-			}
-
-			// Decrypt the payload with the newly decrypted main encryption AES key
-			payload, err := common.Decrypt(ii.Payload, dAESKey)
-			if err != nil {
-				return err
-			}
-			if payload == nil {
-				// TODO should this be a logger warn? A payload could theoretically be validly empty
-				return errors.New("decrypting the payload resulted in no error and a nil payload")
-			}
-
-			// Overwrite the payload with the decrypted data
-			ii.Payload = payload
-			ii.Encrypted = false // TODO handle the encryption state in the consuming code
-			continue image
-		}
-	}
-
-	return nil
-}
-
 // handleSharedSecrets process the negotiated secrets received from the encryption layer
 func (m *Messenger) handleSharedSecrets(secrets []*sharedsecret.Secret) error {
 	for _, secret := range secrets {

protocol/messenger_handler.go

@@ -1320,6 +1320,18 @@ func (m *Messenger) HandleChatIdentity(state *ReceivedMessageState, ci protobuf.
 	}
 	contact := state.CurrentMessageState.Contact
 
+	err := DecryptIdentityImagesWithIdentityPrivateKey(ci.Images, m.identity, state.CurrentMessageState.PublicKey)
+	if err != nil {
+		return err
+	}
+
+	// Remove any images still encrypted after the decryption process
+	for name, image := range ci.Images {
+		if image.Encrypted {
+			delete(ci.Images, name)
+		}
+	}
+
 	logger.Info("Handling contact update")
 	newImages, err := m.persistence.SaveContactChatIdentity(contact.ID, &ci)
 	if err != nil {