From d2fcc50f2d86aa872cbf7866f544d2abc61077cd Mon Sep 17 00:00:00 2001
From: Tulir Asokan <tulir@maunium.net>
Date: Fri, 5 Aug 2022 18:08:44 +0300
Subject: [PATCH] Add device ID to appservice OTK count map

---
 appservice/http.go     | 17 ++++++++++-------
 appservice/protocol.go | 17 ++++++++++-------
 2 files changed, 20 insertions(+), 14 deletions(-)

diff --git a/appservice/http.go b/appservice/http.go
index 68d3395e..1ce7c31b 100644
--- a/appservice/http.go
+++ b/appservice/http.go
@@ -161,13 +161,16 @@ func (as *AppService) handleTransaction(id string, txn *Transaction) {
 	as.txnIDC.MarkProcessed(id)
 }
 
-func (as *AppService) handleOTKCounts(otks map[id.UserID]mautrix.OTKCount) {
-	for userID, otkCounts := range otks {
-		otkCounts.UserID = userID
-		select {
-		case as.OTKCounts <- &otkCounts:
-		default:
-			as.Log.Warnfln("Dropped OTK count update for %s because channel is full", userID)
+func (as *AppService) handleOTKCounts(otks OTKCountMap) {
+	for userID, devices := range otks {
+		for deviceID, otkCounts := range devices {
+			otkCounts.UserID = userID
+			otkCounts.DeviceID = deviceID
+			select {
+			case as.OTKCounts <- &otkCounts:
+			default:
+				as.Log.Warnfln("Dropped OTK count update for %s because channel is full", userID)
+			}
 		}
 	}
 }
diff --git a/appservice/protocol.go b/appservice/protocol.go
index f133e652..0b6cc98c 100644
--- a/appservice/protocol.go
+++ b/appservice/protocol.go
@@ -17,16 +17,19 @@ import (
 	"maunium.net/go/mautrix/id"
 )
 
+type OTKCountMap = map[id.UserID]map[id.DeviceID]mautrix.OTKCount
+
 // Transaction contains a list of events.
 type Transaction struct {
-	Events          []*event.Event                 `json:"events"`
-	EphemeralEvents []*event.Event                 `json:"ephemeral,omitempty"`
-	DeviceLists     *mautrix.DeviceLists           `json:"device_lists,omitempty"`
-	DeviceOTKCount  map[id.UserID]mautrix.OTKCount `json:"device_one_time_keys_count,omitempty"`
+	Events          []*event.Event `json:"events"`
+	EphemeralEvents []*event.Event `json:"ephemeral,omitempty"`
 
-	MSC2409EphemeralEvents []*event.Event                 `json:"de.sorunome.msc2409.ephemeral,omitempty"`
-	MSC3202DeviceLists     *mautrix.DeviceLists           `json:"org.matrix.msc3202.device_lists,omitempty"`
-	MSC3202DeviceOTKCount  map[id.UserID]mautrix.OTKCount `json:"org.matrix.msc3202.device_one_time_keys_count,omitempty"`
+	DeviceLists    *mautrix.DeviceLists `json:"device_lists,omitempty"`
+	DeviceOTKCount OTKCountMap          `json:"device_one_time_keys_count,omitempty"`
+
+	MSC2409EphemeralEvents []*event.Event       `json:"de.sorunome.msc2409.ephemeral,omitempty"`
+	MSC3202DeviceLists     *mautrix.DeviceLists `json:"org.matrix.msc3202.device_lists,omitempty"`
+	MSC3202DeviceOTKCount  OTKCountMap          `json:"org.matrix.msc3202.device_one_time_keys_count,omitempty"`
 }
 
 func (txn *Transaction) ContentString() string {