Key Backups (3/3) : Implement querying keys and various bugfixes (#1946)

* Add querying device keys

Makes a bunch of sytests pass

* Apparently only the current version supports uploading keys

* Linting

2 files changed, 133 insertions, 25 deletions

diff --git a/clientapi/routing/key_backup.go b/clientapi/routing/key_backup.go
index dd21d482..ce62a047 100644
--- a/clientapi/routing/key_backup.go
+++ b/clientapi/routing/key_backup.go
@@ -37,7 +37,7 @@ type keyBackupVersionCreateResponse struct {
 type keyBackupVersionResponse struct {
 	Algorithm string          `json:"algorithm"`
 	AuthData  json.RawMessage `json:"auth_data"`
-	Count     int             `json:"count"`
+	Count     int64           `json:"count"`
 	ETag      string          `json:"etag"`
 	Version   string          `json:"version"`
 }
@@ -89,7 +89,10 @@ func CreateKeyBackupVersion(req *http.Request, userAPI userapi.UserInternalAPI,
 // Implements GET  /_matrix/client/r0/room_keys/version and GET /_matrix/client/r0/room_keys/version/{version}
 func KeyBackupVersion(req *http.Request, userAPI userapi.UserInternalAPI, device *userapi.Device, version string) util.JSONResponse {
 	var queryResp userapi.QueryKeyBackupResponse
-	userAPI.QueryKeyBackup(req.Context(), &userapi.QueryKeyBackupRequest{}, &queryResp)
+	userAPI.QueryKeyBackup(req.Context(), &userapi.QueryKeyBackupRequest{
+		UserID:  device.UserID,
+		Version: version,
+	}, &queryResp)
 	if queryResp.Error != "" {
 		return util.ErrorResponse(fmt.Errorf("QueryKeyBackup: %s", queryResp.Error))
 	}
@@ -216,3 +219,73 @@ func UploadBackupKeys(
 		},
 	}
 }
+
+// Get keys from a given backup version. Response returned varies depending on if roomID and sessionID are set.
+func GetBackupKeys(
+	req *http.Request, userAPI userapi.UserInternalAPI, device *userapi.Device, version, roomID, sessionID string,
+) util.JSONResponse {
+	var queryResp userapi.QueryKeyBackupResponse
+	userAPI.QueryKeyBackup(req.Context(), &userapi.QueryKeyBackupRequest{
+		UserID:           device.UserID,
+		Version:          version,
+		ReturnKeys:       true,
+		KeysForRoomID:    roomID,
+		KeysForSessionID: sessionID,
+	}, &queryResp)
+	if queryResp.Error != "" {
+		return util.ErrorResponse(fmt.Errorf("QueryKeyBackup: %s", queryResp.Error))
+	}
+	if !queryResp.Exists {
+		return util.JSONResponse{
+			Code: 404,
+			JSON: jsonerror.NotFound("version not found"),
+		}
+	}
+	if sessionID != "" {
+		// return the key itself if it was found
+		roomData, ok := queryResp.Keys[roomID]
+		if ok {
+			key, ok := roomData[sessionID]
+			if ok {
+				return util.JSONResponse{
+					Code: 200,
+					JSON: key,
+				}
+			}
+		}
+	} else if roomID != "" {
+		roomData, ok := queryResp.Keys[roomID]
+		if ok {
+			// wrap response in "sessions"
+			return util.JSONResponse{
+				Code: 200,
+				JSON: struct {
+					Sessions map[string]userapi.KeyBackupSession `json:"sessions"`
+				}{
+					Sessions: roomData,
+				},
+			}
+		}
+	} else {
+		// response is the same as the upload request
+		var resp keyBackupSessionRequest
+		resp.Rooms = make(map[string]struct {
+			Sessions map[string]userapi.KeyBackupSession `json:"sessions"`
+		})
+		for roomID, roomData := range queryResp.Keys {
+			resp.Rooms[roomID] = struct {
+				Sessions map[string]userapi.KeyBackupSession `json:"sessions"`
+			}{
+				Sessions: roomData,
+			}
+		}
+		return util.JSONResponse{
+			Code: 200,
+			JSON: resp,
+		}
+	}
+	return util.JSONResponse{
+		Code: 404,
+		JSON: jsonerror.NotFound("keys not found"),
+	}
+}
diff --git a/clientapi/routing/routing.go b/clientapi/routing/routing.go
index 3e0c53ee..cfc6c47f 100644
--- a/clientapi/routing/routing.go
+++ b/clientapi/routing/routing.go
@@ -896,11 +896,15 @@ func Setup(
 		}),
 	).Methods(http.MethodGet, http.MethodOptions)
 
-	// Key Backup Versions
-	r0mux.Handle("/room_keys/version/{versionID}",
+	// Key Backup Versions (Metadata)
+
+	r0mux.Handle("/room_keys/version/{version}",
 		httputil.MakeAuthAPI("get_backup_keys_version", userAPI, func(req *http.Request, device *userapi.Device) util.JSONResponse {
-			version := req.URL.Query().Get("version")
-			return KeyBackupVersion(req, userAPI, device, version)
+			vars, err := httputil.URLDecodeMapValues(mux.Vars(req))
+			if err != nil {
+				return util.ErrorResponse(err)
+			}
+			return KeyBackupVersion(req, userAPI, device, vars["version"])
 		}),
 	).Methods(http.MethodGet, http.MethodOptions)
 	r0mux.Handle("/room_keys/version",
@@ -908,28 +912,22 @@ func Setup(
 			return KeyBackupVersion(req, userAPI, device, "")
 		}),
 	).Methods(http.MethodGet, http.MethodOptions)
-	r0mux.Handle("/room_keys/version/{versionID}",
+	r0mux.Handle("/room_keys/version/{version}",
 		httputil.MakeAuthAPI("put_backup_keys_version", userAPI, func(req *http.Request, device *userapi.Device) util.JSONResponse {
-			version := req.URL.Query().Get("version")
-			if version == "" {
-				return util.JSONResponse{
-					Code: 400,
-					JSON: jsonerror.InvalidArgumentValue("version must be specified"),
-				}
+			vars, err := httputil.URLDecodeMapValues(mux.Vars(req))
+			if err != nil {
+				return util.ErrorResponse(err)
 			}
-			return ModifyKeyBackupVersionAuthData(req, userAPI, device, version)
+			return ModifyKeyBackupVersionAuthData(req, userAPI, device, vars["version"])
 		}),
 	).Methods(http.MethodPut)
-	r0mux.Handle("/room_keys/version/{versionID}",
+	r0mux.Handle("/room_keys/version/{version}",
 		httputil.MakeAuthAPI("delete_backup_keys_version", userAPI, func(req *http.Request, device *userapi.Device) util.JSONResponse {
-			version := req.URL.Query().Get("version")
-			if version == "" {
-				return util.JSONResponse{
-					Code: 400,
-					JSON: jsonerror.InvalidArgumentValue("version must be specified"),
-				}
+			vars, err := httputil.URLDecodeMapValues(mux.Vars(req))
+			if err != nil {
+				return util.ErrorResponse(err)
 			}
-			return DeleteKeyBackupVersion(req, userAPI, device, version)
+			return DeleteKeyBackupVersion(req, userAPI, device, vars["version"])
 		}),
 	).Methods(http.MethodDelete)
 	r0mux.Handle("/room_keys/version",
@@ -938,7 +936,8 @@ func Setup(
 		}),
 	).Methods(http.MethodPost, http.MethodOptions)
 
-	// E2E Backup Keys
+	// Inserting E2E Backup Keys
+
 	// Bulk room and session
 	r0mux.Handle("/room_keys/keys",
 		httputil.MakeAuthAPI("put_backup_keys", userAPI, func(req *http.Request, device *userapi.Device) util.JSONResponse {
@@ -973,6 +972,9 @@ func Setup(
 			}
 			roomID := vars["roomID"]
 			var reqBody keyBackupSessionRequest
+			reqBody.Rooms = make(map[string]struct {
+				Sessions map[string]userapi.KeyBackupSession `json:"sessions"`
+			})
 			reqBody.Rooms[roomID] = struct {
 				Sessions map[string]userapi.KeyBackupSession `json:"sessions"`
 			}{
@@ -989,7 +991,7 @@ func Setup(
 	).Methods(http.MethodPut)
 	// Single room, single session
 	r0mux.Handle("/room_keys/keys/{roomID}/{sessionID}",
-		httputil.MakeAuthAPI("put_backup_keys_room", userAPI, func(req *http.Request, device *userapi.Device) util.JSONResponse {
+		httputil.MakeAuthAPI("put_backup_keys_room_session", userAPI, func(req *http.Request, device *userapi.Device) util.JSONResponse {
 			vars, err := httputil.URLDecodeMapValues(mux.Vars(req))
 			if err != nil {
 				return util.ErrorResponse(err)
@@ -1009,14 +1011,47 @@ func Setup(
 			roomID := vars["roomID"]
 			sessionID := vars["sessionID"]
 			var keyReq keyBackupSessionRequest
+			keyReq.Rooms = make(map[string]struct {
+				Sessions map[string]userapi.KeyBackupSession `json:"sessions"`
+			})
 			keyReq.Rooms[roomID] = struct {
 				Sessions map[string]userapi.KeyBackupSession `json:"sessions"`
-			}{}
+			}{
+				Sessions: make(map[string]userapi.KeyBackupSession),
+			}
 			keyReq.Rooms[roomID].Sessions[sessionID] = reqBody
 			return UploadBackupKeys(req, userAPI, device, version, &keyReq)
 		}),
 	).Methods(http.MethodPut)
 
+	// Querying E2E Backup Keys
+
+	r0mux.Handle("/room_keys/keys",
+		httputil.MakeAuthAPI("get_backup_keys", userAPI, func(req *http.Request, device *userapi.Device) util.JSONResponse {
+			return GetBackupKeys(req, userAPI, device, req.URL.Query().Get("version"), "", "")
+		}),
+	).Methods(http.MethodGet, http.MethodOptions)
+	r0mux.Handle("/room_keys/keys/{roomID}",
+		httputil.MakeAuthAPI("get_backup_keys_room", userAPI, func(req *http.Request, device *userapi.Device) util.JSONResponse {
+			vars, err := httputil.URLDecodeMapValues(mux.Vars(req))
+			if err != nil {
+				return util.ErrorResponse(err)
+			}
+			return GetBackupKeys(req, userAPI, device, req.URL.Query().Get("version"), vars["roomID"], "")
+		}),
+	).Methods(http.MethodGet, http.MethodOptions)
+	r0mux.Handle("/room_keys/keys/{roomID}/{sessionID}",
+		httputil.MakeAuthAPI("get_backup_keys_room_session", userAPI, func(req *http.Request, device *userapi.Device) util.JSONResponse {
+			vars, err := httputil.URLDecodeMapValues(mux.Vars(req))
+			if err != nil {
+				return util.ErrorResponse(err)
+			}
+			return GetBackupKeys(req, userAPI, device, req.URL.Query().Get("version"), vars["roomID"], vars["sessionID"])
+		}),
+	).Methods(http.MethodGet, http.MethodOptions)
+
+	// Deleting E2E Backup Keys
+
 	// Supplying a device ID is deprecated.
 	r0mux.Handle("/keys/upload/{deviceID}",
 		httputil.MakeAuthAPI("keys_upload", userAPI, func(req *http.Request, device *userapi.Device) util.JSONResponse {