1use std::borrow::Cow;
2use std::env::VarError;
3use std::error::Error;
4use std::fmt;
5use std::path::Path;
6use std::sync::Arc;
7
8mod kv {
9 use hightower_kv::{
10 AuthService, Error as KvError, KvEngine, SingleNodeEngine, StoreConfig,
11 command::Command,
12 crypto::{AesGcmEncryptor, Argon2SecretHasher},
13 };
14 use rand::RngCore;
15 use std::error::Error;
16 use std::fmt;
17 use std::fs;
18 use std::path::{Path, PathBuf};
19 use std::sync::Arc;
20 use std::time::{SystemTime, UNIX_EPOCH};
21 use tempfile::{Builder as TempDirBuilder, TempDir};
22 use tracing::{debug, info, warn};
23
24 const AUTH_MASTER_KEY: &[u8] = b"secrets/auth_master_key";
25
26 type SharedEngine = Arc<SingleNodeEngine>;
27 type SharedAuthService = AuthService<SharedEngine, Argon2SecretHasher, AesGcmEncryptor>;
28
29 pub type GatewayAuthService = SharedAuthService;
30
31 #[derive(Debug)]
32 pub enum KvInitError {
33 TempDir(std::io::Error),
34 CreateDir(std::io::Error),
35 Store(KvError),
36 }
37
38 impl fmt::Display for KvInitError {
39 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
40 match self {
41 KvInitError::TempDir(err) => {
42 write!(f, "failed to create temporary directory: {}", err)
43 }
44 KvInitError::CreateDir(err) => {
45 write!(f, "failed to create data directory: {}", err)
46 }
47 KvInitError::Store(err) => {
48 write!(f, "failed to start key-value engine: {}", err)
49 }
50 }
51 }
52 }
53
54 impl Error for KvInitError {
55 fn source(&self) -> Option<&(dyn Error + 'static)> {
56 match self {
57 KvInitError::TempDir(err) | KvInitError::CreateDir(err) => Some(err),
58 KvInitError::Store(err) => Some(err),
59 }
60 }
61 }
62
63 pub struct KvHandle {
64 engine: SharedEngine,
65 auth: Arc<SharedAuthService>,
66 #[allow(dead_code)]
67 data_dir: PathBuf,
68 #[allow(dead_code)]
69 temp_dir: Option<TempDir>,
70 }
71
72 impl fmt::Debug for KvHandle {
73 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
74 f.debug_struct("KvHandle")
75 .field("data_dir", &self.data_dir)
76 .field(
77 "temp_dir",
78 &self.temp_dir.as_ref().map(|dir| dir.path().to_path_buf()),
79 )
80 .finish()
81 }
82 }
83
84 impl KvHandle {
85 pub fn data_dir(&self) -> &Path {
86 &self.data_dir
87 }
88
89 #[cfg(test)]
90 pub fn temp_dir_path(&self) -> Option<&Path> {
91 self.temp_dir.as_ref().map(|dir| dir.path())
92 }
93
94 pub fn put_bytes(&self, key: &[u8], value: &[u8]) -> Result<(), KvError> {
95 let (version, timestamp) = monotonic_version_timestamp();
96 self.engine.submit(Command::Set {
97 key: key.to_vec(),
98 value: value.to_vec(),
99 version,
100 timestamp,
101 })?;
102 Ok(())
103 }
104
105 pub fn put_secret(&self, key: &[u8], value: &[u8]) {
106 if let Err(err) = self.put_bytes(key, value) {
107 warn!(?err, "Failed to persist secret to KV");
108 }
109 }
110
111 pub fn get_bytes(&self, key: &[u8]) -> Result<Option<Vec<u8>>, KvError> {
112 self.engine.get(key)
113 }
114
115 pub fn auth(&self) -> Arc<SharedAuthService> {
116 Arc::clone(&self.auth)
117 }
118
119 pub fn get_prefix(&self, prefix: &[u8]) -> Result<Vec<(Vec<u8>, Vec<u8>)>, KvError> {
120 self.engine.get_prefix(prefix)
121 }
122 }
123
124 pub fn initialize(dir: Option<&Path>) -> Result<KvHandle, KvInitError> {
125 let (data_dir, temp_dir) = match dir {
126 Some(path) => {
127 fs::create_dir_all(path).map_err(KvInitError::CreateDir)?;
128 (path.to_path_buf(), None)
129 }
130 None => {
131 let temp = TempDirBuilder::new()
132 .prefix("hightower-kv")
133 .tempdir()
134 .map_err(KvInitError::TempDir)?;
135 (temp.path().to_path_buf(), Some(temp))
136 }
137 };
138
139 let mut config = StoreConfig::default();
140 config.data_dir = data_dir.to_string_lossy().into_owned();
141
142 let engine = SingleNodeEngine::with_config(config).map_err(KvInitError::Store)?;
143 let master_key = load_or_initialize_master_key(&engine).map_err(KvInitError::Store)?;
144 let (engine, auth) = engine.into_argon2_hasher_aes_gcm_auth_service(master_key);
145 let auth = Arc::new(auth);
146 let temporary = temp_dir.is_some();
147 debug!(path = %data_dir.display(), temporary, "Initialized key-value store");
148
149 Ok(KvHandle {
150 engine,
151 auth,
152 data_dir,
153 temp_dir,
154 })
155 }
156
157 fn load_or_initialize_master_key(engine: &SingleNodeEngine) -> Result<[u8; 32], KvError> {
158 match engine.get(AUTH_MASTER_KEY)? {
159 Some(bytes) => match <[u8; 32]>::try_from(bytes.as_slice()) {
160 Ok(key) => Ok(key),
161 Err(_) => {
162 warn!(
163 stored_len = bytes.len(),
164 "Stored auth master key has unexpected length; generating a new key"
165 );
166 generate_and_persist_master_key(engine)
167 }
168 },
169 None => generate_and_persist_master_key(engine),
170 }
171 }
172
173 fn generate_and_persist_master_key(engine: &SingleNodeEngine) -> Result<[u8; 32], KvError> {
174 let mut key = [0u8; 32];
175 rand::thread_rng().fill_bytes(&mut key);
176 let (version, timestamp) = monotonic_version_timestamp();
177 engine.submit(Command::Set {
178 key: AUTH_MASTER_KEY.to_vec(),
179 value: key.to_vec(),
180 version,
181 timestamp,
182 })?;
183 info!("Provisioned new auth master key");
184 Ok(key)
185 }
186
187 fn monotonic_version_timestamp() -> (u64, i64) {
188 let now = SystemTime::now()
189 .duration_since(UNIX_EPOCH)
190 .unwrap_or_default();
191 let millis = now.as_millis();
192 let version = millis.min(u64::MAX as u128) as u64;
193 let timestamp = millis.min(i64::MAX as u128) as i64;
194 (version, timestamp)
195 }
196
197 #[cfg(test)]
198 mod tests {
199 use super::*;
200
201 #[test]
202 fn initialize_uses_provided_directory() {
203 let temp_root = TempDirBuilder::new()
204 .prefix("hightower-test")
205 .tempdir()
206 .unwrap();
207 let target = temp_root.path().join("kv-store");
208
209 let handle = initialize(Some(&target)).expect("kv init succeeds");
210
211 assert_eq!(handle.data_dir(), target.as_path());
212 assert!(target.exists());
213 assert!(handle.temp_dir_path().is_none());
214 }
215
216 #[test]
217 fn initialize_falls_back_to_temp_directory() {
218 let handle = initialize(None).expect("kv init succeeds");
219
220 let temp_dir = handle.temp_dir_path().expect("temp dir is retained");
221 assert!(temp_dir.exists());
222 assert!(handle.data_dir().starts_with(temp_dir));
223 }
224
225 #[test]
226 fn put_bytes_persists_values() {
227 let handle = initialize(None).expect("kv init succeeds");
228 let key = b"kv-tests/cert";
229 let value = b"payload";
230
231 handle.put_bytes(key, value).expect("write succeeds");
232
233 let stored = handle.get_bytes(key).expect("read succeeds");
234 assert_eq!(stored, Some(value.to_vec()));
235 }
236
237 #[test]
238 fn put_secret_does_not_panic_on_failure() {
239 let handle = initialize(None).expect("kv init succeeds");
240 let key = b"kv-tests/secret";
241
242 handle.put_secret(key, b"secret");
243 let stored = handle.get_bytes(key).expect("read succeeds");
244 assert!(stored.is_some());
245 }
246 }
247}
248
249mod token {
250 use std::env::VarError;
251 use std::fmt;
252
253 #[derive(Debug, PartialEq, Eq)]
254 pub enum TokenError {
255 Missing,
256 }
257
258 impl fmt::Display for TokenError {
259 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
260 match self {
261 TokenError::Missing => {
262 write!(f, "HT_AUTH_KEY environment variable must be set.")
263 }
264 }
265 }
266 }
267
268 impl std::error::Error for TokenError {}
269
270 pub fn fetch<F>(mut lookup: F) -> Result<String, TokenError>
271 where
272 F: FnMut(&str) -> Result<String, VarError>,
273 {
274 lookup("HT_AUTH_KEY").map_err(|_| TokenError::Missing)
275 }
276
277 #[cfg(test)]
278 mod tests {
279 use super::*;
280
281 #[test]
282 fn fetch_returns_token_when_present() {
283 let token = fetch(|_| Ok(String::from("123"))).expect("token should be present");
284
285 assert_eq!(token, "123");
286 }
287
288 #[test]
289 fn fetch_reports_missing_token() {
290 let error =
291 fetch(|_| Err(VarError::NotPresent)).expect_err("token should be missing");
292
293 assert_eq!(error, TokenError::Missing);
294 }
295 }
296}
297
298pub use kv::{GatewayAuthService, KvHandle, KvInitError, initialize as initialize_kv};
299pub use token::{TokenError, fetch as fetch_token};
300
301pub const NODE_NAME_KEY: &[u8] = b"nodes/name";
302pub const HT_AUTH_KEY: &[u8] = b"secrets/ht_auth_key";
303const DEFAULT_AUTH_USERNAME_ENV: &str = "HT_DEFAULT_USER";
304const DEFAULT_AUTH_PASSWORD_ENV: &str = "HT_DEFAULT_PASSWORD";
305const DEFAULT_AUTH_USERNAME: &str = "admin";
306const DEFAULT_AUTH_PASSWORD: &str = "admin";
307
308#[derive(Clone)]
309pub struct CommonContext {
310 pub kv: NamespacedKv,
311 pub auth: Arc<GatewayAuthService>,
312}
313
314impl CommonContext {
315 pub fn new(kv: KvHandle) -> Self {
316 let kv = NamespacedKv::from_handle(kv);
317 let auth = kv.auth_service();
318 Self { kv, auth }
319 }
320
321 pub fn namespaced(&self, prefix: &[u8]) -> Self {
322 Self {
323 kv: self.kv.clone_with_additional_prefix(prefix),
324 auth: Arc::clone(&self.auth),
325 }
326 }
327}
328
329impl fmt::Debug for CommonContext {
330 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
331 f.debug_struct("CommonContext")
332 .field("kv", &self.kv)
333 .finish()
334 }
335}
336
337#[derive(Clone, Debug)]
338pub struct NamespacedKv {
339 inner: Arc<KvHandle>,
340 prefix: Option<Vec<u8>>,
341}
342
343impl NamespacedKv {
344 pub fn from_handle(kv: KvHandle) -> Self {
345 Self {
346 inner: Arc::new(kv),
347 prefix: None,
348 }
349 }
350
351 fn with_parts(inner: Arc<KvHandle>, prefix: Option<Vec<u8>>) -> Self {
352 Self { inner, prefix }
353 }
354
355 pub fn clone_with_additional_prefix(&self, additional: &[u8]) -> Self {
356 let prefix = match &self.prefix {
357 Some(existing) => {
358 let mut composed = existing.clone();
359 if !existing.is_empty() {
360 composed.push(b'/');
361 }
362 composed.extend_from_slice(additional);
363 composed
364 }
365 None => additional.to_vec(),
366 };
367
368 Self::with_parts(Arc::clone(&self.inner), Some(prefix))
369 }
370
371 pub fn put_bytes(&self, key: &[u8], value: &[u8]) -> Result<(), hightower_kv::Error> {
372 let key = self.prefixed_key(key);
373 self.inner.put_bytes(key.as_ref(), value)
374 }
375
376 pub fn put_secret(&self, key: &[u8], value: &[u8]) {
377 let key = self.prefixed_key(key);
378 self.inner.put_secret(key.as_ref(), value);
379 }
380
381 pub fn get_bytes(&self, key: &[u8]) -> Result<Option<Vec<u8>>, hightower_kv::Error> {
382 let key = self.prefixed_key(key);
383 self.inner.get_bytes(key.as_ref())
384 }
385
386 pub fn auth_service(&self) -> Arc<GatewayAuthService> {
387 self.inner.auth()
388 }
389
390 pub fn list_by_prefix(
391 &self,
392 prefix: &[u8],
393 ) -> Result<Vec<(Vec<u8>, Vec<u8>)>, hightower_kv::Error> {
394 let full_prefix = match self.prefixed_key(prefix) {
395 Cow::Borrowed(bytes) => bytes.to_vec(),
396 Cow::Owned(bytes) => bytes,
397 };
398
399 let entries = self.inner.get_prefix(&full_prefix)?;
400
401 let mut results = Vec::new();
402 for (key, value) in entries {
403 let mut remainder = &key[full_prefix.len()..];
404 if remainder.first() == Some(&b'/') {
405 remainder = &remainder[1..];
406 }
407 results.push((remainder.to_vec(), value));
408 }
409
410 Ok(results)
411 }
412
413 fn prefixed_key<'a>(&self, key: &'a [u8]) -> Cow<'a, [u8]> {
414 match &self.prefix {
415 Some(prefix) if !prefix.is_empty() => {
416 let mut composed = Vec::with_capacity(prefix.len() + 1 + key.len());
417 composed.extend_from_slice(prefix);
418 composed.push(b'/');
419 composed.extend_from_slice(key);
420 Cow::Owned(composed)
421 }
422 _ => Cow::Borrowed(key),
423 }
424 }
425}
426
427#[derive(Debug)]
428pub enum ContextError {
429 Token(TokenError),
430 Kv(KvInitError),
431 Auth(hightower_kv::Error),
432}
433
434impl fmt::Display for ContextError {
435 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
436 match self {
437 ContextError::Token(err) => write!(f, "failed to read HT_AUTH_KEY: {}", err),
438 ContextError::Kv(err) => write!(f, "failed to initialize key-value store: {}", err),
439 ContextError::Auth(err) => write!(f, "failed to bootstrap auth service: {}", err),
440 }
441 }
442}
443
444impl Error for ContextError {
445 fn source(&self) -> Option<&(dyn Error + 'static)> {
446 match self {
447 ContextError::Token(err) => Some(err),
448 ContextError::Kv(err) => Some(err),
449 ContextError::Auth(err) => Some(err),
450 }
451 }
452}
453
454pub fn initialize_with_token_source<F>(
455 kv_path: Option<&Path>,
456 mut lookup: F,
457) -> Result<CommonContext, ContextError>
458where
459 F: FnMut(&str) -> Result<String, VarError>,
460{
461 let token = token::fetch(|key| lookup(key)).map_err(ContextError::Token)?;
462 initialize_with_token(kv_path, token)
463}
464
465pub fn initialize_with_token(
466 kv_path: Option<&Path>,
467 token: String,
468) -> Result<CommonContext, ContextError> {
469 let kv = kv::initialize(kv_path).map_err(ContextError::Kv)?;
470 let context = CommonContext::new(kv);
471 context.kv.put_secret(HT_AUTH_KEY, token.as_bytes());
472 bootstrap_default_user(&context).map_err(ContextError::Auth)?;
473 Ok(context)
474}
475
476fn bootstrap_default_user(context: &CommonContext) -> Result<(), hightower_kv::Error> {
477 let username = std::env::var(DEFAULT_AUTH_USERNAME_ENV)
478 .unwrap_or_else(|_| DEFAULT_AUTH_USERNAME.to_string())
479 .trim()
480 .to_owned();
481
482 if username.is_empty() {
483 tracing::warn!("Skipping default auth bootstrap; username is empty");
484 return Ok(());
485 }
486
487 let password = std::env::var(DEFAULT_AUTH_PASSWORD_ENV)
488 .unwrap_or_else(|_| DEFAULT_AUTH_PASSWORD.to_string());
489
490 if password.trim().is_empty() {
491 tracing::warn!("Skipping default auth bootstrap; password is empty");
492 return Ok(());
493 }
494
495 match context.auth.create_user(&username, &password) {
496 Ok(_) => {
497 tracing::info!(username = %username, "Bootstrapped default auth user");
498 Ok(())
499 }
500 Err(hightower_kv::Error::Conflict(_)) => {
501 tracing::debug!(username = %username, "Default auth user already exists");
502 Ok(())
503 }
504 Err(err) => Err(err),
505 }
506}
507
508#[cfg(test)]
509pub mod fixtures {
510 use crate::context::{CommonContext, initialize_kv};
511 use tempfile::TempDir;
512
513 pub fn context() -> CommonContext {
514 let temp = TempDir::new().expect("tempdir");
515 let kv = initialize_kv(Some(temp.path())).expect("kv init");
516 CommonContext::new(kv)
517 }
518}
519
520pub mod env {
521 use std::sync::{Mutex, MutexGuard, OnceLock};
522
523 pub const DISABLE_GATEWAY_REGISTRATION_ENV: &str = "HT_DISABLE_ROOT_REGISTRATION";
524
525 fn registration_lock() -> &'static Mutex<()> {
526 static LOCK: OnceLock<Mutex<()>> = OnceLock::new();
527 LOCK.get_or_init(|| Mutex::new(()))
528 }
529
530 pub struct RegistrationGuard {
531 lock: Option<MutexGuard<'static, ()>>,
532 }
533
534 impl RegistrationGuard {
535 fn new(lock: MutexGuard<'static, ()>) -> Self {
536 Self { lock: Some(lock) }
537 }
538 }
539
540 impl Drop for RegistrationGuard {
541 fn drop(&mut self) {
542 unsafe {
543 std::env::remove_var(DISABLE_GATEWAY_REGISTRATION_ENV);
544 }
545 drop(self.lock.take());
546 }
547 }
548
549 pub fn disable_gateway_registration() -> RegistrationGuard {
550 let guard = registration_lock().lock().expect("registration lock");
551 unsafe {
552 std::env::set_var(DISABLE_GATEWAY_REGISTRATION_ENV, "1");
553 }
554 RegistrationGuard::new(guard)
555 }
556
557 pub struct RegistrationEnableGuard {
558 lock: Option<MutexGuard<'static, ()>>,
559 previous: Option<String>,
560 }
561
562 impl RegistrationEnableGuard {
563 fn new(lock: MutexGuard<'static, ()>, previous: Option<String>) -> Self {
564 Self {
565 lock: Some(lock),
566 previous,
567 }
568 }
569 }
570
571 impl Drop for RegistrationEnableGuard {
572 fn drop(&mut self) {
573 unsafe {
574 if let Some(value) = self.previous.take() {
575 std::env::set_var(DISABLE_GATEWAY_REGISTRATION_ENV, value);
576 } else {
577 std::env::remove_var(DISABLE_GATEWAY_REGISTRATION_ENV);
578 }
579 }
580 drop(self.lock.take());
581 }
582 }
583
584 pub fn enable_gateway_registration() -> RegistrationEnableGuard {
585 let guard = registration_lock().lock().expect("registration lock");
586 let previous = std::env::var(DISABLE_GATEWAY_REGISTRATION_ENV).ok();
587 unsafe {
588 std::env::remove_var(DISABLE_GATEWAY_REGISTRATION_ENV);
589 }
590 RegistrationEnableGuard::new(guard, previous)
591 }
592}
593
594#[cfg(test)]
595mod tests {
596 use super::*;
597 use std::env::VarError;
598 use tempfile::TempDir;
599
600 #[test]
601 fn initialize_with_token_source_persists_token() {
602 let temp = TempDir::new().expect("tempdir");
603 let context =
604 initialize_with_token_source(Some(temp.path()), |_| Ok("test-auth".into()))
605 .expect("initialize");
606
607 let stored = context
608 .kv
609 .get_bytes(HT_AUTH_KEY)
610 .expect("kv read")
611 .expect("value present");
612 assert_eq!(stored, b"test-auth");
613 }
614
615 #[test]
616 fn initialize_with_token_source_reports_missing_token() {
617 let error = initialize_with_token_source(None, |_| Err(VarError::NotPresent))
618 .expect_err("missing token");
619 assert!(matches!(error, ContextError::Token(TokenError::Missing)));
620 }
621
622 #[test]
623 fn initialize_with_token_bootstraps_default_user() {
624 use std::sync::{Mutex, OnceLock};
625
626 fn env_lock() -> &'static Mutex<()> {
627 static LOCK: OnceLock<Mutex<()>> = OnceLock::new();
628 LOCK.get_or_init(|| Mutex::new(()))
629 }
630
631 let _guard = env_lock().lock().expect("env lock");
632
633 let previous_user = std::env::var(DEFAULT_AUTH_USERNAME_ENV).ok();
634 let previous_password = std::env::var(DEFAULT_AUTH_PASSWORD_ENV).ok();
635 unsafe {
636 std::env::remove_var(DEFAULT_AUTH_USERNAME_ENV);
637 std::env::remove_var(DEFAULT_AUTH_PASSWORD_ENV);
638 }
639
640 let temp = TempDir::new().expect("tempdir");
641 let context = initialize_with_token(Some(temp.path()), "token".into())
642 .expect("context initialized");
643
644 assert!(
645 context
646 .auth
647 .verify_password(DEFAULT_AUTH_USERNAME, DEFAULT_AUTH_PASSWORD)
648 .expect("default password verification")
649 );
650
651 match previous_user {
652 Some(value) => unsafe { std::env::set_var(DEFAULT_AUTH_USERNAME_ENV, value) },
653 None => unsafe { std::env::remove_var(DEFAULT_AUTH_USERNAME_ENV) },
654 }
655
656 match previous_password {
657 Some(value) => unsafe { std::env::set_var(DEFAULT_AUTH_PASSWORD_ENV, value) },
658 None => unsafe { std::env::remove_var(DEFAULT_AUTH_PASSWORD_ENV) },
659 }
660 }
661
662 #[test]
663 fn context_creates_isolated_store() {
664 let ctx_a = fixtures::context();
665 let ctx_b = fixtures::context();
666
667 ctx_a
668 .kv
669 .put_bytes(b"test/key", b"value-a")
670 .expect("store a");
671
672 let stored = ctx_b.kv.get_bytes(b"test/key").expect("read b");
673 assert!(stored.is_none(), "unexpected shared state");
674 }
675
676 #[test]
677 fn disable_gateway_registration_sets_env() {
678 {
679 let _guard = env::disable_gateway_registration();
680 let value = std::env::var(env::DISABLE_GATEWAY_REGISTRATION_ENV).unwrap();
681 assert_eq!(value, "1");
682 }
683
684 assert!(std::env::var(env::DISABLE_GATEWAY_REGISTRATION_ENV).is_err());
685 }
686
687 #[test]
688 fn enable_gateway_registration_restores_previous_state() {
689 use std::sync::{Mutex, OnceLock};
690
691 fn test_env_lock() -> &'static Mutex<()> {
692 static LOCK: OnceLock<Mutex<()>> = OnceLock::new();
693 LOCK.get_or_init(|| Mutex::new(()))
694 }
695
696 let _test_lock = test_env_lock().lock().expect("test env lock");
697
698 let previous = std::env::var(env::DISABLE_GATEWAY_REGISTRATION_ENV).ok();
699
700 unsafe {
701 std::env::set_var(env::DISABLE_GATEWAY_REGISTRATION_ENV, "1");
702 }
703
704 {
705 let _guard = env::enable_gateway_registration();
706 assert!(
707 std::env::var(env::DISABLE_GATEWAY_REGISTRATION_ENV).is_err(),
708 "env var should be cleared while guard is active"
709 );
710 }
711
712 let restored =
713 std::env::var(env::DISABLE_GATEWAY_REGISTRATION_ENV).expect("env var restored");
714 assert_eq!(restored, "1");
715
716 match previous {
717 Some(value) => unsafe { std::env::set_var(env::DISABLE_GATEWAY_REGISTRATION_ENV, value) },
718 None => unsafe { std::env::remove_var(env::DISABLE_GATEWAY_REGISTRATION_ENV) },
719 }
720 }
721}