Skip to main content

reinhardt_testkit/auth/
secondary.rs

1use async_trait::async_trait;
2
3use super::error::TestAuthError;
4use super::identity::SessionIdentity;
5use crate::client::APIClient;
6
7/// Secondary (multi-factor) authentication layer.
8///
9/// Open trait — implement this for custom secondary auth methods
10/// (e.g., PassKey, FIDO2) that compose with a primary auth method.
11#[async_trait]
12pub trait SecondaryAuth: Send + Sync {
13	/// Apply this secondary auth layer to an HTTP test client.
14	///
15	/// Typically sets additional headers or cookies on the client.
16	async fn apply_to_client(
17		&self,
18		client: &APIClient,
19		primary: &SessionIdentity,
20	) -> Result<(), TestAuthError>;
21}
22
23/// TOTP MFA secondary authentication for tests.
24///
25/// When applied, sets the `X-MFA-Code` header with either an auto-generated
26/// valid TOTP code or a manually specified code.
27#[cfg(native)]
28pub struct TotpSecondaryAuth {
29	manager: reinhardt_auth::mfa::MFAAuthentication,
30	code_override: Option<String>,
31}
32
33#[cfg(native)]
34impl TotpSecondaryAuth {
35	/// Create a new TOTP secondary auth using the given MFA manager.
36	pub fn new(manager: reinhardt_auth::mfa::MFAAuthentication) -> Self {
37		Self {
38			manager,
39			code_override: None,
40		}
41	}
42
43	/// Create a TOTP secondary auth with a pre-generated code.
44	///
45	/// Use this when you don't have access to the `MFAAuthentication` manager,
46	/// or when testing with a specific code (including invalid codes).
47	pub fn with_code_only(code: impl Into<String>) -> Self {
48		Self {
49			manager: reinhardt_auth::mfa::MFAAuthentication::new("test"),
50			code_override: Some(code.into()),
51		}
52	}
53
54	/// Override the TOTP code with a specific value.
55	///
56	/// Useful for testing invalid codes or expired codes.
57	pub fn with_code(mut self, code: impl Into<String>) -> Self {
58		self.code_override = Some(code.into());
59		self
60	}
61}
62
63#[cfg(native)]
64#[async_trait]
65impl SecondaryAuth for TotpSecondaryAuth {
66	async fn apply_to_client(
67		&self,
68		client: &APIClient,
69		primary: &SessionIdentity,
70	) -> Result<(), TestAuthError> {
71		let code = match &self.code_override {
72			Some(c) => c.clone(),
73			None => {
74				let secret = self
75					.manager
76					.get_secret(&primary.user_id)
77					.await
78					.ok_or_else(|| TestAuthError::MfaUserNotRegistered(primary.user_id.clone()))?;
79				generate_totp_code(&secret)
80					.map_err(|e| TestAuthError::SecondaryAuthError(e.to_string()))?
81			}
82		};
83		client
84			.set_header("X-MFA-Code", &code)
85			.await
86			.map_err(|e| TestAuthError::ClientError(e.to_string()))?;
87		Ok(())
88	}
89}
90
91/// Generate a TOTP code from a base32-encoded secret using the same
92/// algorithm as `MFAAuthentication` (SHA-256, 6 digits, 30s window).
93#[cfg(native)]
94fn generate_totp_code(secret: &str) -> Result<String, String> {
95	use std::time::{SystemTime, UNIX_EPOCH};
96
97	let secret_bytes = base32_decode(secret).ok_or("invalid base32 secret")?;
98	let time = SystemTime::now()
99		.duration_since(UNIX_EPOCH)
100		.map_err(|e| e.to_string())?
101		.as_secs();
102	let step = time / 30;
103	Ok(totp_lite::totp_custom::<totp_lite::Sha256>(
104		30,
105		6,
106		&secret_bytes,
107		step,
108	))
109}
110
111/// Minimal base32 decoder (RFC 4648, no padding required).
112#[cfg(native)]
113fn base32_decode(input: &str) -> Option<Vec<u8>> {
114	const ALPHABET: &[u8] = b"ABCDEFGHIJKLMNOPQRSTUVWXYZ234567";
115	let mut bits = 0u64;
116	let mut bit_count = 0u32;
117	let mut result = Vec::new();
118
119	for &byte in input.as_bytes() {
120		if byte == b'=' {
121			break;
122		}
123		let val = ALPHABET
124			.iter()
125			.position(|&c| c == byte.to_ascii_uppercase())? as u64;
126		bits = (bits << 5) | val;
127		bit_count += 5;
128		if bit_count >= 8 {
129			bit_count -= 8;
130			result.push((bits >> bit_count) as u8);
131			bits &= (1 << bit_count) - 1;
132		}
133	}
134	Some(result)
135}
136
137#[cfg(test)]
138mod tests {
139	#[cfg(native)]
140	mod totp_tests {
141		use super::super::*;
142		use rstest::*;
143
144		#[rstest]
145		fn base32_decode_known_value() {
146			let decoded = base32_decode("JBSWY3DPEHPK3PXP");
147			assert!(decoded.is_some());
148			let bytes = decoded.unwrap();
149			assert!(!bytes.is_empty());
150		}
151
152		#[rstest]
153		fn generate_totp_produces_six_digit_code() {
154			let code = generate_totp_code("JBSWY3DPEHPK3PXP").unwrap();
155			assert_eq!(code.len(), 6);
156			assert!(code.chars().all(|c| c.is_ascii_digit()));
157		}
158	}
159}