From 1357ae52489c18ea6a02f5375a051a72f8970bd9 Mon Sep 17 00:00:00 2001
From: Neil Alexander <neilalexander@users.noreply.github.com>
Date: Wed, 16 Feb 2022 13:11:51 +0000
Subject: [PATCH] Fix UserExists flag, device check

---
 keyserver/internal/internal.go | 26 ++++++++++----------------
 userapi/internal/api.go        |  1 +
 2 files changed, 11 insertions(+), 16 deletions(-)

diff --git a/keyserver/internal/internal.go b/keyserver/internal/internal.go
index d9cac894f..b19b19b1d 100644
--- a/keyserver/internal/internal.go
+++ b/keyserver/internal/internal.go
@@ -558,6 +558,12 @@ func (a *KeyInternalAPI) uploadLocalDeviceKeys(ctx context.Context, req *api.Per
 		}
 		return
 	}
+	if !uapidevices.UserExists {
+		res.Error = &api.KeyError{
+			Err: fmt.Sprintf("user %q does not exist", req.UserID),
+		}
+		return
+	}
 
 	// Get all of the user existing device keys so we can check for changes.
 	existingKeys, err := a.DB.DeviceKeysForUser(ctx, req.UserID, nil)
@@ -568,8 +574,8 @@ func (a *KeyInternalAPI) uploadLocalDeviceKeys(ctx context.Context, req *api.Per
 		return
 	}
 	existingDeviceMap := make(map[string]struct{}, len(existingKeys))
-	for _, k := range existingKeys {
-		existingDeviceMap[k.DeviceID] = struct{}{}
+	for _, key := range uapidevices.Devices {
+		existingDeviceMap[key.ID] = struct{}{}
 	}
 
 	// Work out whether we have device keys in the keyserver for devices that
@@ -577,13 +583,7 @@ func (a *KeyInternalAPI) uploadLocalDeviceKeys(ctx context.Context, req *api.Per
 	// that we keep some integrity between the two.
 	var toClean []gomatrixserverlib.KeyID
 	for _, k := range existingKeys {
-		found := false
-		for _, d := range uapidevices.Devices {
-			if k.UserID == d.UserID && k.DeviceID == d.ID {
-				found = true
-			}
-		}
-		if !found {
+		if _, ok := existingDeviceMap[k.DeviceID]; !ok {
 			toClean = append(toClean, gomatrixserverlib.KeyID(k.DeviceID))
 		}
 	}
@@ -613,13 +613,7 @@ func (a *KeyInternalAPI) uploadLocalDeviceKeys(ctx context.Context, req *api.Per
 		}
 		// check that the device in question actually exists in the user
 		// API before we try and store a key for it
-		foundDevice := false
-		for _, d := range uapidevices.Devices {
-			if d.ID == key.DeviceID {
-				foundDevice = true
-			}
-		}
-		if !foundDevice {
+		if _, ok := existingDeviceMap[key.DeviceID]; !ok {
 			continue
 		}
 		gotUserID := gjson.GetBytes(key.KeyJSON, "user_id").Str
diff --git a/userapi/internal/api.go b/userapi/internal/api.go
index fb8d42f98..865fc2c34 100644
--- a/userapi/internal/api.go
+++ b/userapi/internal/api.go
@@ -315,6 +315,7 @@ func (a *UserInternalAPI) QueryDevices(ctx context.Context, req *api.QueryDevice
 	if err != nil {
 		return err
 	}
+	res.UserExists = true
 	res.Devices = devs
 	return nil
 }