tuwunel_service/registration_tokens/
mod.rs

1mod data;
2
3use std::{collections::HashSet, fmt::Display, sync::Arc};
4
5use data::Data;
6pub use data::{DatabaseTokenInfo, TokenExpires};
7use futures::{Stream, StreamExt, pin_mut};
8use tuwunel_core::{
9	Err, Result, error,
10	utils::{IterStream, random_string},
11};
12
13const RANDOM_TOKEN_LENGTH: usize = 16;
14
15pub struct Service {
16	db: Data,
17	services: Arc<crate::services::OnceServices>,
18}
19
20/// A validated registration token which may be used to create an account.
21#[derive(Debug)]
22pub struct ValidToken {
23	pub token: String,
24	pub info: TokenInfo,
25}
26
27impl Display for ValidToken {
28	fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
29		write!(f, "`{}` --- {}", self.token, self.info)
30	}
31}
32
33impl PartialEq<str> for ValidToken {
34	fn eq(&self, other: &str) -> bool { self.token == other }
35}
36
37#[derive(Clone, Copy, Debug)]
38pub enum TokenInfo {
39	/// The static token set in the homeserver's config file, which is
40	/// always valid.
41	Config,
42	/// A database token which has been checked to be valid.
43	Database(DatabaseTokenInfo),
44}
45
46impl Display for TokenInfo {
47	fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
48		match self {
49			| Self::Config => write!(f, "Token defined in config file"),
50			| Self::Database(info) => info.fmt(f),
51		}
52	}
53}
54
55impl crate::Service for Service {
56	fn build(args: &crate::Args<'_>) -> Result<Arc<Self>> {
57		Ok(Arc::new(Self {
58			db: Data::new(args.db),
59			services: args.services.clone(),
60		}))
61	}
62
63	fn name(&self) -> &str { crate::service::make_name(std::module_path!()) }
64}
65
66impl Service {
67	/// Create a registration token, using the caller's token or generating a
68	/// random one of `length` characters (default `RANDOM_TOKEN_LENGTH`). A
69	/// token that already exists is rejected.
70	pub async fn create_token(
71		&self,
72		token: Option<&str>,
73		length: Option<usize>,
74		expires: TokenExpires,
75	) -> Result<(String, DatabaseTokenInfo)> {
76		let token = token.map(ToOwned::to_owned).unwrap_or_else(|| {
77			let length = length.unwrap_or(RANDOM_TOKEN_LENGTH);
78
79			random_string(length)
80		});
81
82		let info = self.db.save_token(&token, expires).await?;
83
84		Ok((token, info))
85	}
86
87	/// Look up a token's stored metadata, returning `None` when it is absent.
88	pub async fn get_token_info(&self, token: &str) -> Result<TokenInfo> {
89		if self.get_config_tokens().await.contains(token) {
90			return Ok(TokenInfo::Config);
91		}
92
93		self.db
94			.get_token_info(token)
95			.await
96			.map(TokenInfo::Database)
97	}
98
99	/// Replace a token's expiry, preserving its use counter. Returns a `404`
100	/// when the token is unknown.
101	pub async fn update_token(
102		&self,
103		token: &str,
104		expires: TokenExpires,
105	) -> Result<DatabaseTokenInfo> {
106		if self.get_config_tokens().await.contains(token) {
107			return Err!(Request(Forbidden(
108				"The token set in the config file cannot be updated"
109			)));
110		}
111
112		self.db.update_token(token, expires).await
113	}
114
115	pub async fn is_enabled(&self) -> bool {
116		let stream = self.iterate_tokens().await;
117
118		pin_mut!(stream);
119
120		stream.next().await.is_some()
121	}
122
123	pub async fn get_config_tokens(&self) -> HashSet<String> {
124		let mut tokens = HashSet::new();
125
126		if let Some(file) = &self.services.config.registration_token_file {
127			match tokio::fs::read_to_string(file).await {
128				| Err(e) => error!("Failed to read the registration token file: {e}"),
129				| Ok(text) => tokens.extend(
130					text.split_ascii_whitespace()
131						.map(ToOwned::to_owned),
132				),
133			}
134		}
135
136		if let Some(token) = &self.services.config.registration_token {
137			tokens.insert(token.to_owned());
138		}
139
140		tokens
141	}
142
143	pub async fn is_token_valid(&self, token: &str) -> Result { self.check(token, false).await }
144
145	pub async fn try_consume(&self, token: &str) -> Result { self.check(token, true).await }
146
147	async fn check(&self, token: &str, consume: bool) -> Result {
148		if self.get_config_tokens().await.contains(token)
149			|| self.db.check_token(token, consume).await
150		{
151			return Ok(());
152		}
153
154		Err!(Request(Forbidden("Registration token not valid")))
155	}
156
157	/// Try to revoke a valid token.
158	///
159	/// Note that tokens set in the config file cannot be revoked.
160	pub async fn revoke_token(&self, token: &str) -> Result {
161		if self.get_config_tokens().await.contains(token) {
162			return Err!(Request(Forbidden(
163				"The token set in the config file cannot be revoked. Edit the config file to \
164				 change it."
165			)));
166		}
167
168		self.db.revoke_token(token).await
169	}
170
171	/// Iterate over all valid registration tokens.
172	pub async fn iterate_tokens(&self) -> impl Stream<Item = ValidToken> + Send + '_ {
173		let config_tokens = self
174			.get_config_tokens()
175			.await
176			.into_iter()
177			.map(|token| ValidToken { token, info: TokenInfo::Config })
178			.stream();
179
180		let db_tokens = self
181			.db
182			.iterate_and_clean_tokens()
183			.map(|(token, info)| ValidToken {
184				token: token.to_owned(),
185				info: TokenInfo::Database(info),
186			});
187
188		config_tokens.chain(db_tokens)
189	}
190}
tuwunel_service/registration_tokens/mod.rs

tuwunel_service/registration_tokens/
mod.rs
