From 31799a3b2a733192a656cfa77662599bd1568bdc Mon Sep 17 00:00:00 2001
From: Neil Alexander <neilalexander@users.noreply.github.com>
Date: Fri, 29 Apr 2022 16:02:55 +0100
Subject: [PATCH] Device list display name fixes (#2405)

* Get device names from `unsigned` in `/user/devices`

* Fix display name updates

* Fix bug

* Fix another bug
---
 federationapi/routing/devices.go |  8 ++-
 keyserver/internal/internal.go   | 95 ++++++++++++++++----------------
 2 files changed, 54 insertions(+), 49 deletions(-)

diff --git a/federationapi/routing/devices.go b/federationapi/routing/devices.go
index 8890eac4..57286fa9 100644
--- a/federationapi/routing/devices.go
+++ b/federationapi/routing/devices.go
@@ -20,6 +20,7 @@ import (
 	keyapi "github.com/matrix-org/dendrite/keyserver/api"
 	"github.com/matrix-org/gomatrixserverlib"
 	"github.com/matrix-org/util"
+	"github.com/tidwall/gjson"
 )
 
 // GetUserDevices for the given user id
@@ -69,9 +70,14 @@ func GetUserDevices(
 			continue
 		}
 
+		displayName := dev.DisplayName
+		if displayName == "" {
+			displayName = gjson.GetBytes(dev.DeviceKeys.KeyJSON, "unsigned.device_display_name").Str
+		}
+
 		device := gomatrixserverlib.RespUserDevice{
 			DeviceID:    dev.DeviceID,
-			DisplayName: dev.DisplayName,
+			DisplayName: displayName,
 			Keys:        key,
 		}
 
diff --git a/keyserver/internal/internal.go b/keyserver/internal/internal.go
index 1677cf8e..e556f44b 100644
--- a/keyserver/internal/internal.go
+++ b/keyserver/internal/internal.go
@@ -632,43 +632,55 @@ func (a *KeyInternalAPI) uploadLocalDeviceKeys(ctx context.Context, req *api.Per
 	}
 
 	var keysToStore []api.DeviceMessage
-	// assert that the user ID / device ID are not lying for each key
-	for _, key := range req.DeviceKeys {
-		var serverName gomatrixserverlib.ServerName
-		_, serverName, err = gomatrixserverlib.SplitID('@', key.UserID)
-		if err != nil {
-			continue // ignore invalid users
-		}
-		if serverName != a.ThisServer {
-			continue // ignore remote users
-		}
-		if len(key.KeyJSON) == 0 {
-			keysToStore = append(keysToStore, key.WithStreamID(0))
-			continue // deleted keys don't need sanity checking
-		}
-		// check that the device in question actually exists in the user
-		// API before we try and store a key for it
-		if _, ok := existingDeviceMap[key.DeviceID]; !ok {
-			continue
-		}
-		gotUserID := gjson.GetBytes(key.KeyJSON, "user_id").Str
-		gotDeviceID := gjson.GetBytes(key.KeyJSON, "device_id").Str
-		if gotUserID == key.UserID && gotDeviceID == key.DeviceID {
-			keysToStore = append(keysToStore, key.WithStreamID(0))
-			continue
-		}
-
-		res.KeyError(key.UserID, key.DeviceID, &api.KeyError{
-			Err: fmt.Sprintf(
-				"user_id or device_id mismatch: users: %s - %s, devices: %s - %s",
-				gotUserID, key.UserID, gotDeviceID, key.DeviceID,
-			),
-		})
-	}
 
 	if req.OnlyDisplayNameUpdates {
-		// add the display name field from keysToStore into existingKeys
-		keysToStore = appendDisplayNames(existingKeys, keysToStore)
+		for _, existingKey := range existingKeys {
+			for _, newKey := range req.DeviceKeys {
+				switch {
+				case existingKey.UserID != newKey.UserID:
+					continue
+				case existingKey.DeviceID != newKey.DeviceID:
+					continue
+				case existingKey.DisplayName != newKey.DisplayName:
+					existingKey.DisplayName = newKey.DisplayName
+				}
+			}
+			keysToStore = append(keysToStore, existingKey)
+		}
+	} else {
+		// assert that the user ID / device ID are not lying for each key
+		for _, key := range req.DeviceKeys {
+			var serverName gomatrixserverlib.ServerName
+			_, serverName, err = gomatrixserverlib.SplitID('@', key.UserID)
+			if err != nil {
+				continue // ignore invalid users
+			}
+			if serverName != a.ThisServer {
+				continue // ignore remote users
+			}
+			if len(key.KeyJSON) == 0 {
+				keysToStore = append(keysToStore, key.WithStreamID(0))
+				continue // deleted keys don't need sanity checking
+			}
+			// check that the device in question actually exists in the user
+			// API before we try and store a key for it
+			if _, ok := existingDeviceMap[key.DeviceID]; !ok {
+				continue
+			}
+			gotUserID := gjson.GetBytes(key.KeyJSON, "user_id").Str
+			gotDeviceID := gjson.GetBytes(key.KeyJSON, "device_id").Str
+			if gotUserID == key.UserID && gotDeviceID == key.DeviceID {
+				keysToStore = append(keysToStore, key.WithStreamID(0))
+				continue
+			}
+
+			res.KeyError(key.UserID, key.DeviceID, &api.KeyError{
+				Err: fmt.Sprintf(
+					"user_id or device_id mismatch: users: %s - %s, devices: %s - %s",
+					gotUserID, key.UserID, gotDeviceID, key.DeviceID,
+				),
+			})
+		}
 	}
 
 	// store the device keys and emit changes
@@ -764,16 +776,3 @@ func emitDeviceKeyChanges(producer KeyChangeProducer, existing, new []api.Device
 	}
 	return producer.ProduceKeyChanges(keysAdded)
 }
-
-func appendDisplayNames(existing, new []api.DeviceMessage) []api.DeviceMessage {
-	for i, existingDevice := range existing {
-		for _, newDevice := range new {
-			if existingDevice.DeviceID != newDevice.DeviceID {
-				continue
-			}
-			existingDevice.DisplayName = newDevice.DisplayName
-			existing[i] = existingDevice
-		}
-	}
-	return existing
-}