From 0796e170bd910b0b5ef92a92c95f55b86f70b77c Mon Sep 17 00:00:00 2001
From: Till Faelligen <2353100+S7evinK@users.noreply.github.com>
Date: Sun, 11 Sep 2022 16:51:37 +0200
Subject: [PATCH] Use json.RawMessage directly; try more devices, if available

---
 clientapi/producers/syncapi.go     | 32 ++++++++++++++----------------
 federationapi/producers/syncapi.go | 24 ++++++++++------------
 2 files changed, 25 insertions(+), 31 deletions(-)

diff --git a/clientapi/producers/syncapi.go b/clientapi/producers/syncapi.go
index 5933ce1a8..2dc0c4843 100644
--- a/clientapi/producers/syncapi.go
+++ b/clientapi/producers/syncapi.go
@@ -21,12 +21,13 @@ import (
 	"strconv"
 	"time"
 
-	"github.com/matrix-org/dendrite/setup/jetstream"
-	"github.com/matrix-org/dendrite/syncapi/types"
-	userapi "github.com/matrix-org/dendrite/userapi/api"
 	"github.com/matrix-org/gomatrixserverlib"
 	"github.com/nats-io/nats.go"
 	log "github.com/sirupsen/logrus"
+
+	"github.com/matrix-org/dendrite/setup/jetstream"
+	"github.com/matrix-org/dendrite/syncapi/types"
+	userapi "github.com/matrix-org/dendrite/userapi/api"
 )
 
 // SyncAPIProducer produces events for the sync API server to consume
@@ -61,7 +62,7 @@ func (p *SyncAPIProducer) SendReceipt(
 
 func (p *SyncAPIProducer) SendToDevice(
 	ctx context.Context, sender, userID, deviceID, eventType string,
-	message interface{},
+	message json.RawMessage,
 ) error {
 	devices := []string{}
 	_, domain, err := gomatrixserverlib.SplitID('@', userID)
@@ -89,24 +90,19 @@ func (p *SyncAPIProducer) SendToDevice(
 		devices = append(devices, deviceID)
 	}
 
-	js, err := json.Marshal(message)
-	if err != nil {
-		return err
-	}
-
 	log.WithFields(log.Fields{
 		"user_id":     userID,
 		"num_devices": len(devices),
 		"type":        eventType,
 	}).Tracef("Producing to topic '%s'", p.TopicSendToDeviceEvent)
-	for _, device := range devices {
+	for i, device := range devices {
 		ote := &types.OutputSendToDeviceEvent{
 			UserID:   userID,
 			DeviceID: device,
 			SendToDeviceEvent: gomatrixserverlib.SendToDeviceEvent{
 				Sender:  sender,
 				Type:    eventType,
-				Content: js,
+				Content: message,
 			},
 		}
 
@@ -115,15 +111,17 @@ func (p *SyncAPIProducer) SendToDevice(
 			log.WithError(err).Error("sendToDevice failed json.Marshal")
 			return err
 		}
-		m := &nats.Msg{
-			Subject: p.TopicSendToDeviceEvent,
-			Data:    eventJSON,
-			Header:  nats.Header{},
-		}
+		m := nats.NewMsg(p.TopicSendToDeviceEvent)
+		m.Data = eventJSON
 		m.Header.Set("sender", sender)
 		m.Header.Set(jetstream.UserID, userID)
+
 		if _, err = p.JetStream.PublishMsg(m, nats.Context(ctx)); err != nil {
-			log.WithError(err).Error("sendToDevice failed t.Producer.SendMessage")
+			if i < len(devices)-1 {
+				log.WithError(err).Warn("sendToDevice failed to PublishMsg, trying further devices")
+				continue
+			}
+			log.WithError(err).Error("sendToDevice failed to PublishMsg for all devices")
 			return err
 		}
 	}
diff --git a/federationapi/producers/syncapi.go b/federationapi/producers/syncapi.go
index 86c8c10a3..4abd3fbe5 100644
--- a/federationapi/producers/syncapi.go
+++ b/federationapi/producers/syncapi.go
@@ -64,7 +64,7 @@ func (p *SyncAPIProducer) SendReceipt(
 
 func (p *SyncAPIProducer) SendToDevice(
 	ctx context.Context, sender, userID, deviceID, eventType string,
-	message interface{},
+	message json.RawMessage,
 ) error {
 	devices := []string{}
 	_, domain, err := gomatrixserverlib.SplitID('@', userID)
@@ -92,24 +92,19 @@ func (p *SyncAPIProducer) SendToDevice(
 		devices = append(devices, deviceID)
 	}
 
-	js, err := json.Marshal(message)
-	if err != nil {
-		return err
-	}
-
 	log.WithFields(log.Fields{
 		"user_id":     userID,
 		"num_devices": len(devices),
 		"type":        eventType,
 	}).Tracef("Producing to topic '%s'", p.TopicSendToDeviceEvent)
-	for _, device := range devices {
+	for i, device := range devices {
 		ote := &types.OutputSendToDeviceEvent{
 			UserID:   userID,
 			DeviceID: device,
 			SendToDeviceEvent: gomatrixserverlib.SendToDeviceEvent{
 				Sender:  sender,
 				Type:    eventType,
-				Content: js,
+				Content: message,
 			},
 		}
 
@@ -118,16 +113,17 @@ func (p *SyncAPIProducer) SendToDevice(
 			log.WithError(err).Error("sendToDevice failed json.Marshal")
 			return err
 		}
-		m := &nats.Msg{
-			Subject: p.TopicSendToDeviceEvent,
-			Data:    eventJSON,
-			Header:  nats.Header{},
-		}
+		m := nats.NewMsg(p.TopicSendToDeviceEvent)
+		m.Data = eventJSON
 		m.Header.Set("sender", sender)
 		m.Header.Set(jetstream.UserID, userID)
 
 		if _, err = p.JetStream.PublishMsg(m, nats.Context(ctx)); err != nil {
-			log.WithError(err).Error("sendToDevice failed t.Producer.SendMessage")
+			if i < len(devices)-1 {
+				log.WithError(err).Warn("sendToDevice failed to PublishMsg, trying further devices")
+				continue
+			}
+			log.WithError(err).Error("sendToDevice failed to PublishMsg for all devices")
 			return err
 		}
 	}