send: get LID for pn during message enc

Author: purpshell

send.go

@@ -14,6 +14,7 @@ import (
 	"encoding/hex"
 	"errors"
 	"fmt"
+	"slices"
 	"sort"
 	"strconv"
 	"strings"
@@ -1171,6 +1172,46 @@ func (cli *Client) encryptMessageForDevices(
 	includeIdentity := false
 	participantNodes := make([]waBinary.Node, 0, len(allDevices))
 	var retryDevices, retryEncryptionIdentities []types.JID
+	jidLidMap := make(map[types.JID]types.JID)
+	usyncJids := make([]types.JID, 0)
+	usyncDeviceList := make(map[types.JID][]uint16, 0)
+
+	for _, jid := range allDevices {
+		if jid.Server == types.DefaultUserServer {
+			lidForPN, err := cli.Store.LIDs.GetLIDForPN(ctx, jid)
+			if err != nil {
+				cli.Log.Warnf("Failed to get LID for %s: %v", jid, err)
+			} else if !lidForPN.IsEmpty() {
+				jidLidMap[jid] = lidForPN
+			} else {
+				if !slices.Contains(usyncJids, jid.ToNonAD()) {
+					usyncJids = append(usyncJids, jid.ToNonAD())
+				}
+
+				if _, ok := usyncDeviceList[jid.ToNonAD()]; ok {
+					usyncDeviceList[jid.ToNonAD()] = append(usyncDeviceList[jid.ToNonAD()], jid.Device)
+				}
+			}
+		}
+	}
+
+	if len(usyncJids) > 0 {
+		info, err := cli.GetUserInfo(usyncJids)
+		if err != nil {
+			cli.Log.Warnf("Failed to get LID info from USync, err: %v", err)
+		}
+		for jid, userInfo := range info {
+			if !userInfo.LID.IsEmpty() {
+				for _, deviceID := range usyncDeviceList[jid] {
+					jid.Device = deviceID
+					lid := userInfo.LID
+					lid.Device = deviceID
+					jidLidMap[jid] = lid
+				}
+			}
+		}
+	}
+
 	for _, jid := range allDevices {
 		plaintext := msgPlaintext
 		if (jid.User == ownJID.User || jid.User == ownLID.User) && dsmPlaintext != nil {
@@ -1181,10 +1222,8 @@ func (cli *Client) encryptMessageForDevices(
 		}
 		encryptionIdentity := jid
 		if jid.Server == types.DefaultUserServer {
-			lidForPN, err := cli.Store.LIDs.GetLIDForPN(ctx, jid)
-			if err != nil {
-				cli.Log.Warnf("Failed to get LID for %s: %v", jid, err)
-			} else if !lidForPN.IsEmpty() {
+			lidForPN := jidLidMap[jid]
+			if !lidForPN.IsEmpty() {
 				cli.migrateSessionStore(ctx, jid, lidForPN)
 				encryptionIdentity = lidForPN
 			}