tuwunel_service/admin/
context.rs

1use std::{fmt, fmt::Debug, time::SystemTime};
2
3use futures::{FutureExt, lock::Mutex};
4use ruma::UserId;
5use tokio::time::Instant;
6use tuwunel_core::{Err, Result};
7
8use crate::Services;
9
10/// Ceiling on a single command's accumulated output; a handler that writes past
11/// it aborts rather than letting the buffer grow without bound.
12const OUTPUT_MAX_BYTES: usize = 64 * 1024 * 1024;
13
14/// What a command handler runs with, built afresh for each command.
15///
16/// The `body` holds the lines below the command line, and the processor takes
17/// whatever the handler wrote to `output` once the handler returns.
18pub struct Context<'a> {
19	pub services: &'a Services,
20	pub body: &'a [&'a str],
21
22	/// The user who issued the command, or `None` for the operator.
23	pub sender: Option<&'a UserId>,
24
25	pub timer: SystemTime,
26	pub output: Mutex<String>,
27}
28
29impl Context<'_> {
30	pub async fn write_timed_query<F, T>(&self, query: F) -> Result
31	where
32		F: Future<Output = T>,
33		T: Debug,
34	{
35		let timer = Instant::now();
36		let result = query.await;
37		let query_time = timer.elapsed();
38
39		self.write_string(format!(
40			"Query completed in {query_time:?}:\n\n```rs\n{result:#?}\n```"
41		))
42		.await
43	}
44
45	pub async fn write_timed_query_try<F, T>(&self, query: F) -> Result
46	where
47		F: Future<Output = Result<T>>,
48		T: Debug,
49	{
50		let timer = Instant::now();
51		let result = query.await?;
52		let query_time = timer.elapsed();
53
54		self.write_string(format!(
55			"Query completed in {query_time:?}:\n\n```rs\n{result:#?}\n```"
56		))
57		.await
58	}
59
60	pub fn write_fmt(
61		&self,
62		arguments: fmt::Arguments<'_>,
63	) -> impl Future<Output = Result> + Send + '_ + use<'_> {
64		let buf = format!("{arguments}");
65		self.write_string(buf)
66	}
67
68	#[inline]
69	pub async fn write_string(&self, s: String) -> Result { self.write_str(&s).await }
70
71	pub fn write_str<'a>(&'a self, s: &'a str) -> impl Future<Output = Result> + Send + 'a {
72		self.output.lock().map(move |mut output| {
73			if output.len().saturating_add(s.len()) > OUTPUT_MAX_BYTES {
74				return Err!("Command output exceeded the maximum size and was aborted.");
75			}
76
77			output.push_str(s);
78			Ok(())
79		})
80	}
81}
tuwunel_service/admin/context.rs

tuwunel_service/admin/
context.rs
