tuwunel_service/admin/
context.rs

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

tuwunel_service/admin/
context.rs
