From e7b2c1428011e0f32353bbd42b8765b21511abd3 Mon Sep 17 00:00:00 2001
From: Jason Volk <jason@zemos.net>
Date: Wed, 28 Aug 2024 07:05:13 +0000
Subject: [PATCH] tag admin room m.server_notice on grant

Signed-off-by: Jason Volk <jason@zemos.net>
---
 src/core/config/mod.rs     |  5 +++++
 src/service/admin/grant.rs | 44 +++++++++++++++++++++++++++++++++++---
 src/service/admin/mod.rs   |  4 +++-
 3 files changed, 49 insertions(+), 4 deletions(-)

diff --git a/src/core/config/mod.rs b/src/core/config/mod.rs
index 1f78dea6..12b36a65 100644
--- a/src/core/config/mod.rs
+++ b/src/core/config/mod.rs
@@ -342,6 +342,8 @@ pub struct Config {
 	pub admin_execute_errors_ignore: bool,
 	#[serde(default = "default_admin_log_capture")]
 	pub admin_log_capture: String,
+	#[serde(default = "default_admin_room_tag")]
+	pub admin_room_tag: String,
 
 	#[serde(default)]
 	pub sentry: bool,
@@ -608,6 +610,7 @@ impl fmt::Display for Config {
 			&self.admin_execute_errors_ignore.to_string(),
 		);
 		line("Filter for admin command log capture", &self.admin_log_capture);
+		line("Admin room tag", &self.admin_room_tag);
 		line("Allow outgoing federated typing", &self.allow_outgoing_typing.to_string());
 		line("Allow incoming federated typing", &self.allow_incoming_typing.to_string());
 		line(
@@ -1068,3 +1071,5 @@ fn default_sentry_filter() -> String { "info".to_owned() }
 fn default_startup_netburst_keep() -> i64 { 50 }
 
 fn default_admin_log_capture() -> String { "debug".to_owned() }
+
+fn default_admin_room_tag() -> String { "m.server_notice".to_owned() }
diff --git a/src/service/admin/grant.rs b/src/service/admin/grant.rs
index 44e11a3e..8cbee604 100644
--- a/src/service/admin/grant.rs
+++ b/src/service/admin/grant.rs
@@ -1,6 +1,6 @@
 use std::collections::BTreeMap;
 
-use conduit::Result;
+use conduit::{error, implement, Result};
 use ruma::{
 	events::{
 		room::{
@@ -8,9 +8,10 @@ use ruma::{
 			message::RoomMessageEventContent,
 			power_levels::RoomPowerLevelsEventContent,
 		},
-		TimelineEventType,
+		tag::{TagEvent, TagEventContent, TagInfo},
+		RoomAccountDataEventType, TimelineEventType,
 	},
-	UserId,
+	RoomId, UserId,
 };
 use serde_json::value::to_raw_value;
 
@@ -110,6 +111,14 @@ impl super::Service {
 			)
 			.await?;
 
+		// Set room tag
+		let room_tag = &self.services.server.config.admin_room_tag;
+		if !room_tag.is_empty() {
+			if let Err(e) = self.set_room_tag(&room_id, user_id, room_tag) {
+				error!(?room_id, ?user_id, ?room_tag, ?e, "Failed to set tag for admin grant");
+			}
+		}
+
 		// Send welcome message
 		self.services.timeline.build_and_append_pdu(
   			PduBuilder {
@@ -132,3 +141,32 @@ impl super::Service {
 		Ok(())
 	}
 }
+
+#[implement(super::Service)]
+fn set_room_tag(&self, room_id: &RoomId, user_id: &UserId, tag: &str) -> Result<()> {
+	let mut event = self
+		.services
+		.account_data
+		.get(Some(room_id), user_id, RoomAccountDataEventType::Tag)?
+		.map(|event| serde_json::from_str(event.get()))
+		.and_then(Result::ok)
+		.unwrap_or_else(|| TagEvent {
+			content: TagEventContent {
+				tags: BTreeMap::new(),
+			},
+		});
+
+	event
+		.content
+		.tags
+		.insert(tag.to_owned().into(), TagInfo::new());
+
+	self.services.account_data.update(
+		Some(room_id),
+		user_id,
+		RoomAccountDataEventType::Tag,
+		&serde_json::to_value(event)?,
+	)?;
+
+	Ok(())
+}
diff --git a/src/service/admin/mod.rs b/src/service/admin/mod.rs
index bd89d5a3..3274249e 100644
--- a/src/service/admin/mod.rs
+++ b/src/service/admin/mod.rs
@@ -23,7 +23,7 @@ use ruma::{
 use serde_json::value::to_raw_value;
 use tokio::sync::{Mutex, RwLock};
 
-use crate::{globals, rooms, rooms::state::RoomMutexGuard, Dep};
+use crate::{account_data, globals, rooms, rooms::state::RoomMutexGuard, Dep};
 
 pub struct Service {
 	services: Services,
@@ -42,6 +42,7 @@ struct Services {
 	timeline: Dep<rooms::timeline::Service>,
 	state: Dep<rooms::state::Service>,
 	state_cache: Dep<rooms::state_cache::Service>,
+	account_data: Dep<account_data::Service>,
 	services: StdRwLock<Option<Weak<crate::Services>>>,
 }
 
@@ -87,6 +88,7 @@ impl crate::Service for Service {
 				timeline: args.depend::<rooms::timeline::Service>("rooms::timeline"),
 				state: args.depend::<rooms::state::Service>("rooms::state"),
 				state_cache: args.depend::<rooms::state_cache::Service>("rooms::state_cache"),
+				account_data: args.depend::<account_data::Service>("account_data"),
 				services: None.into(),
 			},
 			sender,