send: Acquire LIDs from PNs by USync during message enc

send.go

@@ -1246,6 +1246,53 @@ func (cli *Client) encryptMessageForDevices(
 	}
 
 	var retryDevices, retryEncryptionIdentities []types.JID
+	jidLidMap := make(map[types.JID]types.JID)
+	usyncDeviceList := make(map[types.JID][]uint16) // bare JID -> device IDs
+
+	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)
+			}
+
+			if !lidForPN.IsEmpty() {
+				jidLidMap[jid] = lidForPN
+			} else {
+				bare := jid.ToNonAD()
+				usyncDeviceList[bare] = append(usyncDeviceList[bare], jid.Device)
+			}
+
+		}
+	}
+
+	if len(usyncDeviceList) > 0 {
+		usyncJids := make([]types.JID, 0, len(usyncDeviceList))
+		for bare := range usyncDeviceList {
+			usyncJids = append(usyncJids, bare)
+		}
+
+		info, err := cli.GetUserInfo(usyncJids)
+		if err != nil {
+			cli.Log.Warnf("Failed to get LID info from USync, err: %v", err)
+		} else {
+			for bare, userInfo := range info {
+				if userInfo.LID.IsEmpty() {
+					continue
+				}
+				for _, deviceID := range usyncDeviceList[bare] {
+					jid := bare
+					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 {
@@ -1254,7 +1301,14 @@ func (cli *Client) encryptMessageForDevices(
 			}
 			plaintext = dsmPlaintext
 		}
-		encryptionIdentity := encryptionIdentities[jid]
+
+		encryptionIdentity := jid
+		if jid.Server == types.DefaultUserServer {
+			if lid, ok := jidLidMap[jid]; ok && !lid.IsEmpty() {
+				cli.migrateSessionStore(ctx, jid, lid)
+				encryptionIdentity = lid
+			}
+		}
 
 		encrypted, isPreKey, err := cli.encryptMessageForDeviceAndWrap(
 			ctx, plaintext, jid, encryptionIdentity, nil, encAttrs, existingSessions,