From 8e13fb89d044379bd919ba03513847d65904a8d2 Mon Sep 17 00:00:00 2001
From: Neil Alexander <neilalexander@users.noreply.github.com>
Date: Fri, 29 Apr 2022 13:55:30 +0100
Subject: [PATCH] Fix display name updates

---
 federationapi/routing/devices.go |  2 +-
 keyserver/internal/internal.go   | 96 ++++++++++++++++----------------
 2 files changed, 49 insertions(+), 49 deletions(-)

diff --git a/federationapi/routing/devices.go b/federationapi/routing/devices.go
index 9be0df988..57286fa90 100644
--- a/federationapi/routing/devices.go
+++ b/federationapi/routing/devices.go
@@ -71,7 +71,7 @@ func GetUserDevices(
 		}
 
 		displayName := dev.DisplayName
-		if dev.DisplayName == "" {
+		if displayName == "" {
 			displayName = gjson.GetBytes(dev.DeviceKeys.KeyJSON, "unsigned.device_display_name").Str
 		}
 
diff --git a/keyserver/internal/internal.go b/keyserver/internal/internal.go
index 1677cf8e3..d836cffa6 100644
--- a/keyserver/internal/internal.go
+++ b/keyserver/internal/internal.go
@@ -632,43 +632,56 @@ 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)
+		keysToStore = append(keysToStore, existingKeys...)
+		for _, existingKey := range existingKeys {
+			for _, newKey := range req.DeviceKeys {
+				switch {
+				case existingKey.UserID != newKey.DeviceID:
+					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 +777,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
-}