1use std::collections::HashSet;
10use std::path::Path;
11
12use anyhow::{Context, Result, bail};
13use serde::{Deserialize, Serialize};
14
15use crate::noise::Psk;
16
17pub const DEFAULT_MAX_CONNS: usize = 64;
18
19fn default_max_conns() -> usize {
20 DEFAULT_MAX_CONNS
21}
22
23pub fn normalize_host(host: &str) -> String {
26 let host = host.trim();
27 host.strip_suffix('.').unwrap_or(host).to_ascii_lowercase()
28}
29
30#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
31pub struct Tenant {
32 pub name: String,
33 pub psk: Psk,
34 pub sni: String,
35 #[serde(default = "default_max_conns")]
36 pub max_conns: usize,
37 #[serde(default)]
38 pub disabled: bool,
39}
40
41#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
43pub struct EnrollCode {
44 pub psk: Psk,
45 pub sni: String,
46 #[serde(default = "default_max_conns")]
47 pub max_conns: usize,
48 #[serde(default)]
50 pub expires_at: Option<u64>,
51}
52
53#[derive(Debug, Clone, PartialEq, Eq)]
55pub enum Identity {
56 Tenant(Tenant),
57 Code(EnrollCode),
58}
59
60#[derive(Debug, Clone, Default)]
62pub struct Candidates {
63 pub psks: Vec<Psk>,
64 ids: Vec<Identity>,
65}
66
67impl Candidates {
68 pub fn identity(&self, index: usize) -> &Identity {
69 &self.ids[index]
70 }
71}
72
73#[derive(Debug, Clone, Default, PartialEq, Eq)]
75pub struct TenantSet {
76 tenants: Vec<Tenant>,
77 codes: Vec<EnrollCode>,
78}
79
80impl TenantSet {
81 pub fn new(tenants: Vec<Tenant>, codes: Vec<EnrollCode>) -> Result<Self> {
85 let mut names = HashSet::new();
86 let mut hosts = HashSet::new();
87 let mut psks = HashSet::new();
88 for t in &tenants {
89 if t.name.is_empty() {
90 bail!("tenant has an empty name");
91 }
92 if !names.insert(t.name.clone()) {
93 bail!("duplicate tenant name '{}'", t.name);
94 }
95 check_entry(&t.sni, t.max_conns, &t.name)?;
96 if !hosts.insert(normalize_host(&t.sni)) {
97 bail!("duplicate sni '{}' (tenant '{}')", t.sni, t.name);
98 }
99 if !psks.insert(t.psk) {
100 bail!("tenant '{}' shares a PSK with another entry", t.name);
101 }
102 }
103 for c in &codes {
104 check_entry(&c.sni, c.max_conns, "enrollment code")?;
105 if !psks.insert(c.psk) {
106 bail!("an enrollment code shares a PSK with another entry");
107 }
108 }
109 Ok(Self { tenants, codes })
110 }
111
112 pub fn tenants(&self) -> &[Tenant] {
113 &self.tenants
114 }
115
116 pub fn codes(&self) -> &[EnrollCode] {
117 &self.codes
118 }
119
120 pub fn tenant(&self, name: &str) -> Option<&Tenant> {
121 self.tenants.iter().find(|t| t.name == name)
122 }
123
124 pub fn tenant_by_sni(&self, sni: &str) -> Option<&Tenant> {
126 let wanted = normalize_host(sni);
127 self.tenants
128 .iter()
129 .find(|t| !t.disabled && normalize_host(&t.sni) == wanted)
130 }
131
132 pub fn candidates(&self, now_unix: u64) -> Candidates {
135 let mut out = Candidates::default();
136 for t in self.tenants.iter().filter(|t| !t.disabled) {
137 out.psks.push(t.psk);
138 out.ids.push(Identity::Tenant(t.clone()));
139 }
140 for c in self
141 .codes
142 .iter()
143 .filter(|c| c.expires_at.is_none_or(|at| now_unix < at))
144 {
145 out.psks.push(c.psk);
146 out.ids.push(Identity::Code(c.clone()));
147 }
148 out
149 }
150}
151
152fn check_entry(sni: &str, max_conns: usize, what: &str) -> Result<()> {
153 if normalize_host(sni).is_empty() {
154 bail!("{what} has an empty sni");
155 }
156 if max_conns == 0 {
157 bail!("{what} has max_conns = 0");
158 }
159 Ok(())
160}
161
162fn random_hex(bytes: usize) -> String {
164 let mut buf = vec![0u8; bytes];
165 rand::fill(&mut buf[..]);
166 buf.iter().map(|b| format!("{b:02x}")).collect()
167}
168
169pub fn new_tenant_name() -> String {
172 format!("t-{}", random_hex(6))
173}
174
175pub fn new_token() -> String {
177 random_hex(32)
178}
179
180pub fn new_enroll_code(
184 sni: &str,
185 max_conns: usize,
186 expires_at: Option<u64>,
187) -> (String, EnrollCode) {
188 let code = random_hex(16);
189 let entry = EnrollCode {
190 psk: crate::noise::derive_psk(&code),
191 sni: sni.to_string(),
192 max_conns,
193 expires_at,
194 };
195 (code, entry)
196}
197
198impl TenantSet {
200 pub fn with_tenant(&self, tenant: Tenant) -> Result<Self> {
201 let mut tenants = self.tenants.clone();
202 tenants.push(tenant);
203 Self::new(tenants, self.codes.clone())
204 }
205
206 pub fn with_code(&self, code: EnrollCode) -> Result<Self> {
207 let host = normalize_host(&code.sni);
209 if let Some(t) = self.tenants.iter().find(|t| normalize_host(&t.sni) == host) {
210 bail!("sni '{}' already belongs to tenant '{}'", code.sni, t.name);
211 }
212 let mut codes = self.codes.clone();
213 codes.push(code);
214 Self::new(self.tenants.clone(), codes)
215 }
216
217 pub fn without_tenant(&self, name: &str) -> Result<Self> {
218 if self.tenant(name).is_none() {
219 bail!("no tenant named '{name}'");
220 }
221 let tenants = self
222 .tenants
223 .iter()
224 .filter(|t| t.name != name)
225 .cloned()
226 .collect();
227 Self::new(tenants, self.codes.clone())
228 }
229
230 pub fn with_disabled(&self, name: &str, disabled: bool) -> Result<Self> {
231 if self.tenant(name).is_none() {
232 bail!("no tenant named '{name}'");
233 }
234 let mut tenants = self.tenants.clone();
235 for t in tenants.iter_mut().filter(|t| t.name == name) {
236 t.disabled = disabled;
237 }
238 Self::new(tenants, self.codes.clone())
239 }
240
241 pub fn redeem(
245 &self,
246 code_psk: &Psk,
247 now_unix: u64,
248 name: String,
249 psk: Psk,
250 ) -> Result<(Self, Tenant)> {
251 let code = self
252 .codes
253 .iter()
254 .find(|c| c.psk == *code_psk && c.expires_at.is_none_or(|at| now_unix < at))
255 .context("unknown, used or expired enrollment code")?;
256 let tenant = Tenant {
257 name,
258 psk,
259 sni: code.sni.clone(),
260 max_conns: code.max_conns,
261 disabled: false,
262 };
263 let codes = self
264 .codes
265 .iter()
266 .filter(|c| c.psk != *code_psk && c.expires_at.is_none_or(|at| now_unix < at))
268 .cloned()
269 .collect();
270 let mut tenants = self.tenants.clone();
271 tenants.push(tenant.clone());
272 Ok((Self::new(tenants, codes)?, tenant))
273 }
274}
275
276#[derive(Debug, Default, Serialize, Deserialize)]
278struct TenantFile {
279 #[serde(default)]
280 tenant: Vec<Tenant>,
281 #[serde(default)]
282 code: Vec<EnrollCode>,
283}
284
285impl TenantSet {
286 pub fn from_toml(text: &str) -> Result<Self> {
287 let file: TenantFile = toml::from_str(text).context("parsing tenant file")?;
288 Self::new(file.tenant, file.code)
289 }
290
291 pub fn to_toml(&self) -> Result<String> {
292 let file = TenantFile {
293 tenant: self.tenants.clone(),
294 code: self.codes.clone(),
295 };
296 toml::to_string_pretty(&file).context("serialising tenant file")
297 }
298
299 pub fn load(path: &Path) -> Result<Self> {
302 match std::fs::read_to_string(path) {
303 Ok(text) => Self::from_toml(&text)
304 .with_context(|| format!("loading tenant file {}", path.display())),
305 Err(e) if e.kind() == std::io::ErrorKind::NotFound => Ok(Self::default()),
306 Err(e) => Err(e).with_context(|| format!("reading {}", path.display())),
307 }
308 }
309}
310
311#[cfg(test)]
313mod tests {
314 use super::*;
315 use crate::noise::derive_psk;
316
317 fn tenant(name: &str, token: &str, sni: &str) -> Tenant {
318 Tenant {
319 name: name.into(),
320 psk: derive_psk(token),
321 sni: sni.into(),
322 max_conns: 8,
323 disabled: false,
324 }
325 }
326
327 fn code(token: &str, sni: &str, expires_at: Option<u64>) -> EnrollCode {
328 EnrollCode {
329 psk: derive_psk(token),
330 sni: sni.into(),
331 max_conns: 8,
332 expires_at,
333 }
334 }
335
336 #[test]
337 fn sni_lookup_ignores_case_and_trailing_dot() {
338 let set = TenantSet::new(vec![tenant("a", "ta", "Abc.Example.com")], vec![]).unwrap();
339 assert_eq!(set.tenant_by_sni("abc.example.COM.").unwrap().name, "a");
340 assert!(set.tenant_by_sni("other.example.com").is_none());
341 }
342
343 #[test]
344 fn disabled_tenants_are_neither_routed_nor_candidates() {
345 let mut t = tenant("a", "ta", "a.example.com");
346 t.disabled = true;
347 let set = TenantSet::new(vec![t], vec![]).unwrap();
348 assert!(set.tenant_by_sni("a.example.com").is_none());
349 assert!(set.candidates(0).psks.is_empty());
350 }
351
352 #[test]
353 fn expired_codes_are_not_candidates() {
354 let set = TenantSet::new(
355 vec![],
356 vec![
357 code("fresh", "a.x", Some(200)),
358 code("stale", "b.x", Some(100)),
359 ],
360 )
361 .unwrap();
362 let c = set.candidates(150);
363 assert_eq!(c.psks, vec![derive_psk("fresh")]);
364 assert!(matches!(c.identity(0), Identity::Code(_)));
365 assert!(set.candidates(100).psks.len() == 1, "expiry is exclusive");
366 }
367
368 #[test]
369 fn candidates_map_back_to_their_identity() {
370 let set = TenantSet::new(
371 vec![tenant("a", "ta", "a.x"), tenant("b", "tb", "b.x")],
372 vec![code("c1", "c.x", None)],
373 )
374 .unwrap();
375 let c = set.candidates(0);
376 let at = |token: &str| c.psks.iter().position(|p| *p == derive_psk(token)).unwrap();
377 assert!(matches!(c.identity(at("tb")), Identity::Tenant(t) if t.name == "b"));
378 assert!(matches!(c.identity(at("c1")), Identity::Code(_)));
379 }
380
381 #[test]
382 fn ambiguous_sets_are_rejected() {
383 let dup_name = TenantSet::new(
384 vec![tenant("a", "t1", "a.x"), tenant("a", "t2", "b.x")],
385 vec![],
386 );
387 let dup_host = TenantSet::new(
388 vec![tenant("a", "t1", "A.x"), tenant("b", "t2", "a.X")],
389 vec![],
390 );
391 let dup_psk = TenantSet::new(
392 vec![tenant("a", "same", "a.x"), tenant("b", "same", "b.x")],
393 vec![],
394 );
395 let code_psk = TenantSet::new(
396 vec![tenant("a", "same", "a.x")],
397 vec![code("same", "z.x", None)],
398 );
399 let empty_sni = TenantSet::new(vec![tenant("a", "t", " ")], vec![]);
400 let zero = {
401 let mut t = tenant("a", "t", "a.x");
402 t.max_conns = 0;
403 TenantSet::new(vec![t], vec![])
404 };
405 for (what, r) in [
406 ("duplicate name", dup_name),
407 ("duplicate host", dup_host),
408 ("duplicate psk", dup_psk),
409 ("code psk clash", code_psk),
410 ("empty sni", empty_sni),
411 ("zero max_conns", zero),
412 ] {
413 assert!(r.is_err(), "{what} should be rejected");
414 }
415 }
416
417 #[test]
418 fn toml_round_trip_stores_the_psk_as_hex() {
419 #[derive(Serialize, Deserialize, PartialEq, Debug)]
420 struct File {
421 tenant: Vec<Tenant>,
422 }
423 let original = File {
424 tenant: vec![tenant("a", "ta", "a.x")],
425 };
426 let text = toml::to_string(&original).unwrap();
427 assert!(text.contains(&format!(
428 "psk = \"{}\"",
429 hex_of(derive_psk("ta").as_bytes())
430 )));
431 assert!(!text.contains("\"ta\""), "the token must never be stored");
432 assert_eq!(toml::from_str::<File>(&text).unwrap(), original);
433 }
434
435 #[test]
436 fn generated_names_tokens_and_codes_are_well_formed_and_unique() {
437 let mut seen = HashSet::new();
438 for _ in 0..1000 {
439 let name = new_tenant_name();
440 assert!(name.starts_with("t-") && name.len() == 14, "{name}");
441 let token = new_token();
442 assert_eq!(token.len(), 64);
443 let (code, entry) = new_enroll_code("a.x", 4, None);
444 assert_eq!(code.len(), 32);
445 assert_eq!(entry.psk, derive_psk(&code));
446 assert!(seen.insert(name) && seen.insert(token) && seen.insert(code));
447 }
448 }
449
450 #[test]
451 fn malformed_psk_is_rejected_on_load() {
452 let bad = "name = \"a\"\npsk = \"zz\"\nsni = \"a.x\"\n";
453 assert!(toml::from_str::<Tenant>(bad).is_err());
454 }
455
456 fn hex_of(b: &[u8; 32]) -> String {
457 b.iter().map(|x| format!("{x:02x}")).collect()
458 }
459}
460
461#[cfg(test)]
462mod redaction_tests {
463 use super::*;
464
465 #[test]
466 fn debug_output_never_contains_the_psk() {
467 let tenant = Tenant {
468 name: "t".into(),
469 psk: Psk::from_bytes([0xab; 32]),
470 sni: "a.example.com".into(),
471 max_conns: 1,
472 disabled: false,
473 };
474 let code = EnrollCode {
475 psk: Psk::from_bytes([0xcd; 32]),
476 sni: "a.example.com".into(),
477 max_conns: 1,
478 expires_at: None,
479 };
480 let shown = format!("{tenant:?} {code:?}");
481 assert!(shown.contains("<redacted>"));
482 assert!(!shown.contains("171") && !shown.contains("205"), "{shown}");
483 }
484}
485
486#[cfg(test)]
487mod prune_tests {
488 use super::*;
489
490 #[test]
491 fn redeeming_drops_other_expired_codes() {
492 let code = |b: u8, expires_at| EnrollCode {
493 psk: Psk::from_bytes([b; 32]),
494 sni: format!("h{b}.example.com"),
495 max_conns: 1,
496 expires_at,
497 };
498 let set = TenantSet::new(
499 vec![],
500 vec![code(1, None), code(2, Some(50)), code(3, Some(500))],
501 )
502 .unwrap();
503 let (next, _) = set
504 .redeem(
505 &Psk::from_bytes([1; 32]),
506 100,
507 "t-x".into(),
508 Psk::from_bytes([9; 32]),
509 )
510 .unwrap();
511 let left: Vec<u8> = next.codes().iter().map(|c| c.psk.as_bytes()[0]).collect();
512 assert_eq!(left, vec![3]);
513 }
514}