From 2051c22a281b6daeda252e75c1332d26edd6f48f Mon Sep 17 00:00:00 2001
From: Jason Volk <jason@zemos.net>
Date: Mon, 28 Apr 2025 01:32:13 +0000
Subject: [PATCH] Support optional device_id's in lazy-loading context.

Co-authored-by: Jade Ellis <jade@ellis.link>
Signed-off-by: Jason Volk <jason@zemos.net>
---
 src/api/client/context.rs             |  2 +-
 src/api/client/message.rs             | 28 +++++++++++++--------------
 src/api/client/sync/v3.rs             |  2 +-
 src/service/rooms/lazy_loading/mod.rs |  4 ++--
 4 files changed, 17 insertions(+), 19 deletions(-)

diff --git a/src/api/client/context.rs b/src/api/client/context.rs
index dbc2a22f..ca787a16 100644
--- a/src/api/client/context.rs
+++ b/src/api/client/context.rs
@@ -111,7 +111,7 @@ pub(crate) async fn get_context_route(
 
 	let lazy_loading_context = lazy_loading::Context {
 		user_id: sender_user,
-		device_id: sender_device,
+		device_id: Some(sender_device),
 		room_id,
 		token: Some(base_count.into_unsigned()),
 		options: Some(&filter.lazy_load_options),
diff --git a/src/api/client/message.rs b/src/api/client/message.rs
index e442850b..7a87a9b0 100644
--- a/src/api/client/message.rs
+++ b/src/api/client/message.rs
@@ -1,5 +1,3 @@
-use core::panic;
-
 use axum::extract::State;
 use conduwuit::{
 	Err, Result, at,
@@ -34,6 +32,7 @@ use ruma::{
 	},
 	serde::Raw,
 };
+use tracing::warn;
 
 use crate::Ruma;
 
@@ -73,7 +72,7 @@ pub(crate) async fn get_message_events_route(
 ) -> Result<get_message_events::v3::Response> {
 	debug_assert!(IGNORED_MESSAGE_TYPES.is_sorted(), "IGNORED_MESSAGE_TYPES is not sorted");
 	let sender_user = body.sender_user();
-	let sender_device = body.sender_device.as_ref();
+	let sender_device = body.sender_device.as_deref();
 	let room_id = &body.room_id;
 	let filter = &body.filter;
 
@@ -137,18 +136,17 @@ pub(crate) async fn get_message_events_route(
 
 	let lazy_loading_context = lazy_loading::Context {
 		user_id: sender_user,
-		device_id: match sender_device {
-			| Some(device_id) => device_id,
-			| None =>
-				if let Some(registration) = body.appservice_info.as_ref() {
-					<&DeviceId>::from(registration.registration.id.as_str())
-				} else {
-					panic!(
-						"No device_id provided and no appservice registration found, this \
-						 should be unreachable"
-					);
-				},
-		},
+		device_id: sender_device.or_else(|| {
+			if let Some(registration) = body.appservice_info.as_ref() {
+				Some(<&DeviceId>::from(registration.registration.id.as_str()))
+			} else {
+				warn!(
+					"No device_id provided and no appservice registration found, this should be \
+					 unreachable"
+				);
+				None
+			}
+		}),
 		room_id,
 		token: Some(from.into_unsigned()),
 		options: Some(&filter.lazy_load_options),
diff --git a/src/api/client/sync/v3.rs b/src/api/client/sync/v3.rs
index 7bc74c95..feaf8689 100644
--- a/src/api/client/sync/v3.rs
+++ b/src/api/client/sync/v3.rs
@@ -645,7 +645,7 @@ async fn load_joined_room(
 
 	let lazy_loading_context = &lazy_loading::Context {
 		user_id: sender_user,
-		device_id: sender_device,
+		device_id: Some(sender_device),
 		room_id,
 		token: Some(since),
 		options: Some(&filter.room.state.lazy_load_options),
diff --git a/src/service/rooms/lazy_loading/mod.rs b/src/service/rooms/lazy_loading/mod.rs
index 346314d1..61f081a9 100644
--- a/src/service/rooms/lazy_loading/mod.rs
+++ b/src/service/rooms/lazy_loading/mod.rs
@@ -27,7 +27,7 @@ pub trait Options: Send + Sync {
 #[derive(Clone, Debug)]
 pub struct Context<'a> {
 	pub user_id: &'a UserId,
-	pub device_id: &'a DeviceId,
+	pub device_id: Option<&'a DeviceId>,
 	pub room_id: &'a RoomId,
 	pub token: Option<u64>,
 	pub options: Option<&'a LazyLoadOptions>,
@@ -40,7 +40,7 @@ pub enum Status {
 }
 
 pub type Witness = HashSet<OwnedUserId>;
-type Key<'a> = (&'a UserId, &'a DeviceId, &'a RoomId, &'a UserId);
+type Key<'a> = (&'a UserId, Option<&'a DeviceId>, &'a RoomId, &'a UserId);
 
 impl crate::Service for Service {
 	fn build(args: crate::Args<'_>) -> Result<Arc<Self>> {