1use std::sync::Arc;
2
3use argon2::{Argon2, PasswordVerifier};
4use diesel::{
5 BoolExpressionMethods, ExpressionMethods, JoinOnDsl, NullableExpressionMethods, OptionalExtension, QueryDsl,
6 SelectableHelper,
7};
8use diesel_async::RunQueryDsl;
9use pkcs8::DecodePublicKey;
10use rand::RngCore;
11use sha2::Digest;
12use tonic::Code;
13use tonic_types::{ErrorDetails, StatusExt};
14
15use crate::CoreConfig;
16use crate::chrono_ext::ChronoDateTimeExt;
17use crate::id::Id;
18use crate::middleware::IpAddressInfo;
19use crate::models::{
20 MfaRecoveryCode, MfaWebauthnCredential, Organization, OrganizationId, User, UserEmail, UserId, UserSession,
21};
22use crate::schema::{
23 mfa_recovery_codes, mfa_totp_credentials, mfa_webauthn_auth_sessions, mfa_webauthn_credentials, organizations,
24 user_emails, user_sessions, users,
25};
26use crate::std_ext::{DisplayExt, OptionExt, ResultExt};
27
28pub(crate) fn generate_random_bytes() -> Result<[u8; 32], rand::Error> {
29 let mut token = [0u8; 32];
30 rand::rngs::OsRng.try_fill_bytes(&mut token)?;
31 Ok(token)
32}
33
34#[derive(Debug, thiserror::Error)]
35pub(crate) enum TxError {
36 #[error("diesel transaction error: {0}")]
37 Diesel(#[from] diesel::result::Error),
38 #[error("tonic status error: {0}")]
39 Status(#[from] tonic::Status),
40}
41
42impl From<TxError> for tonic::Status {
43 fn from(err: TxError) -> Self {
44 match err {
45 TxError::Diesel(e) => e.into_tonic_internal_err("transaction error"),
46 TxError::Status(s) => s,
47 }
48 }
49}
50
51pub(crate) fn encrypt_token(
52 algorithm: pb::scufflecloud::core::v1::DeviceAlgorithm,
53 token: &[u8],
54 pk_der_data: &[u8],
55) -> Result<Vec<u8>, tonic::Status> {
56 match algorithm {
57 pb::scufflecloud::core::v1::DeviceAlgorithm::RsaOaepSha256 => {
58 let pk = rsa::RsaPublicKey::from_public_key_der(pk_der_data)
59 .into_tonic_err_with_field_violation("public_key_data", "failed to parse public key")?;
60 let padding = rsa::Oaep::new::<sha2::Sha256>();
61 let enc_data = pk
62 .encrypt(&mut rsa::rand_core::OsRng, padding, token)
63 .into_tonic_internal_err("failed to encrypt token")?;
64 Ok(enc_data)
65 }
66 }
67}
68
69pub(crate) async fn get_user_by_id<G: CoreConfig>(global: &Arc<G>, user_id: UserId) -> Result<User, tonic::Status> {
70 global
71 .user_loader()
72 .load(user_id)
73 .await
74 .ok()
75 .into_tonic_internal_err("failed to query user")?
76 .into_tonic_not_found("user not found")
77}
78
79pub(crate) async fn get_user_by_id_in_tx(
80 db: &mut diesel_async::AsyncPgConnection,
81 user_id: UserId,
82) -> Result<User, tonic::Status> {
83 let user = users::dsl::users
84 .find(user_id)
85 .select(User::as_select())
86 .first::<User>(db)
87 .await
88 .optional()
89 .into_tonic_internal_err("failed to query user")?
90 .into_tonic_not_found("user not found")?;
91
92 Ok(user)
93}
94
95pub(crate) async fn get_user_by_email(db: &mut diesel_async::AsyncPgConnection, email: &str) -> Result<User, tonic::Status> {
96 let Some((user, _)) = users::dsl::users
97 .inner_join(user_emails::dsl::user_emails.on(users::dsl::primary_email.eq(user_emails::dsl::email.nullable())))
98 .filter(user_emails::dsl::email.eq(&email))
99 .select((User::as_select(), user_emails::dsl::email))
100 .first::<(User, String)>(db)
101 .await
102 .optional()
103 .into_tonic_internal_err("failed to query user by email")?
104 else {
105 return Err(tonic::Status::with_error_details(
106 tonic::Code::NotFound,
107 "user not found",
108 ErrorDetails::new(),
109 ));
110 };
111
112 Ok(user)
113}
114
115pub(crate) async fn get_organization_by_id(
116 db: &mut diesel_async::AsyncPgConnection,
117 organization_id: OrganizationId,
118) -> Result<Organization, tonic::Status> {
119 let organization = organizations::dsl::organizations
120 .find(organization_id)
121 .first::<Organization>(db)
122 .await
123 .optional()
124 .into_tonic_internal_err("failed to load organization")?
125 .ok_or_else(|| {
126 tonic::Status::with_error_details(tonic::Code::NotFound, "organization not found", ErrorDetails::new())
127 })?;
128
129 Ok(organization)
130}
131
132pub(crate) fn normalize_email(email: &str) -> String {
133 email.trim().to_ascii_lowercase()
134}
135
136pub(crate) async fn create_user(tx: &mut diesel_async::AsyncPgConnection, new_user: &User) -> Result<(), tonic::Status> {
137 diesel::insert_into(users::dsl::users)
138 .values(new_user)
139 .execute(tx)
140 .await
141 .into_tonic_internal_err("failed to insert user")?;
142
143 if let Some(email) = new_user.primary_email.as_ref() {
144 if user_emails::dsl::user_emails
146 .find(email)
147 .select(user_emails::dsl::email)
148 .first::<String>(tx)
149 .await
150 .optional()
151 .into_tonic_internal_err("failed to query user emails")?
152 .is_some()
153 {
154 return Err(tonic::Status::with_error_details(
155 Code::AlreadyExists,
156 "email is already registered",
157 ErrorDetails::new(),
158 ));
159 }
160
161 let user_email = UserEmail {
162 email: email.clone(),
163 user_id: new_user.id,
164 created_at: chrono::Utc::now(),
165 };
166
167 diesel::insert_into(user_emails::dsl::user_emails)
168 .values(&user_email)
169 .execute(tx)
170 .await
171 .into_tonic_internal_err("failed to insert user email")?;
172 }
173
174 Ok(())
175}
176
177pub(crate) async fn mfa_options(
178 tx: &mut diesel_async::AsyncPgConnection,
179 user_id: UserId,
180) -> Result<Vec<pb::scufflecloud::core::v1::MfaOption>, tonic::Status> {
181 let mut mfa_options = vec![];
182
183 if mfa_totp_credentials::dsl::mfa_totp_credentials
184 .filter(mfa_totp_credentials::dsl::user_id.eq(user_id))
185 .count()
186 .get_result::<i64>(tx)
187 .await
188 .into_tonic_internal_err("failed to query mfa factors")?
189 > 0
190 {
191 mfa_options.push(pb::scufflecloud::core::v1::MfaOption::Totp);
192 }
193
194 if mfa_webauthn_credentials::dsl::mfa_webauthn_credentials
195 .filter(mfa_webauthn_credentials::dsl::user_id.eq(user_id))
196 .count()
197 .get_result::<i64>(tx)
198 .await
199 .into_tonic_internal_err("failed to query mfa factors")?
200 > 0
201 {
202 mfa_options.push(pb::scufflecloud::core::v1::MfaOption::WebAuthn);
203 }
204
205 if mfa_recovery_codes::dsl::mfa_recovery_codes
206 .filter(mfa_recovery_codes::dsl::user_id.eq(user_id))
207 .count()
208 .get_result::<i64>(tx)
209 .await
210 .into_tonic_internal_err("failed to query mfa factors")?
211 > 0
212 {
213 mfa_options.push(pb::scufflecloud::core::v1::MfaOption::RecoveryCodes);
214 }
215
216 Ok(mfa_options)
217}
218
219pub(crate) async fn create_session<G: CoreConfig>(
220 global: &Arc<G>,
221 tx: &mut diesel_async::AsyncPgConnection,
222 user_id: UserId,
223 device: pb::scufflecloud::core::v1::Device,
224 ip_info: &IpAddressInfo,
225 check_mfa: bool,
226) -> Result<pb::scufflecloud::core::v1::NewUserSessionToken, tonic::Status> {
227 let mfa_options = if check_mfa { mfa_options(tx, user_id).await? } else { vec![] };
228
229 let device_fingerprint = sha2::Sha256::digest(&device.public_key_data).to_vec();
231
232 let session_expires_at = if !mfa_options.is_empty() {
233 chrono::Utc::now() + global.mfa_timeout()
234 } else {
235 chrono::Utc::now() + global.user_session_timeout()
236 };
237 let token_id = Id::new();
238 let token_expires_at = chrono::Utc::now() + global.user_session_token_timeout();
239
240 let token = generate_random_bytes().into_tonic_internal_err("failed to generate token")?;
241 let encrypted_token = encrypt_token(device.algorithm(), &token, &device.public_key_data)?;
242
243 let user_session = UserSession {
244 user_id,
245 device_fingerprint,
246 device_algorithm: device.algorithm().into(),
247 device_pk_data: device.public_key_data,
248 last_used_at: chrono::Utc::now(),
249 last_ip: ip_info.to_network(),
250 token_id: Some(token_id),
251 token: Some(token.to_vec()),
252 token_expires_at: Some(token_expires_at),
253 expires_at: session_expires_at,
254 mfa_pending: !mfa_options.is_empty(),
255 };
256
257 diesel::insert_into(user_sessions::dsl::user_sessions)
258 .values(&user_session)
259 .execute(tx)
260 .await
261 .into_tonic_internal_err("failed to insert user session")?;
262
263 let new_token = pb::scufflecloud::core::v1::NewUserSessionToken {
264 id: token_id.to_string(),
265 encrypted_token,
266 expires_at: Some(token_expires_at.to_prost_timestamp_utc()),
267 session_mfa_pending: user_session.mfa_pending,
268 mfa_options: mfa_options.into_iter().map(|o| o as i32).collect(),
269 };
270
271 Ok(new_token)
272}
273
274pub(crate) fn verify_password(password_hash: &str, password: &str) -> Result<(), tonic::Status> {
275 let password_hash = argon2::PasswordHash::new(password_hash).into_tonic_internal_err("failed to parse password hash")?;
276
277 match Argon2::default().verify_password(password.as_bytes(), &password_hash) {
278 Ok(_) => Ok(()),
279 Err(argon2::password_hash::Error::Password) => Err(tonic::Status::with_error_details(
280 tonic::Code::PermissionDenied,
281 "invalid password",
282 ErrorDetails::with_bad_request_violation("password", "invalid password"),
283 )),
284 Err(_) => Err(tonic::Status::with_error_details(
285 tonic::Code::Internal,
286 "failed to verify password",
287 ErrorDetails::new(),
288 )),
289 }
290}
291
292pub(crate) async fn finish_webauthn_authentication<G: CoreConfig>(
293 global: &Arc<G>,
294 tx: &mut diesel_async::AsyncPgConnection,
295 user_id: UserId,
296 reg: &webauthn_rs::prelude::PublicKeyCredential,
297) -> Result<(), tonic::Status> {
298 let state = diesel::delete(mfa_webauthn_auth_sessions::dsl::mfa_webauthn_auth_sessions)
299 .filter(
300 mfa_webauthn_auth_sessions::dsl::user_id
301 .eq(user_id)
302 .and(mfa_webauthn_auth_sessions::dsl::expires_at.gt(chrono::Utc::now())),
303 )
304 .returning(mfa_webauthn_auth_sessions::dsl::state)
305 .get_result::<serde_json::Value>(tx)
306 .await
307 .optional()
308 .into_tonic_internal_err("failed to query webauthn authentication session")?
309 .into_tonic_err(
310 tonic::Code::FailedPrecondition,
311 "no webauthn authentication session found",
312 ErrorDetails::new(),
313 )?;
314
315 let state: webauthn_rs::prelude::PasskeyAuthentication =
316 serde_json::from_value(state).into_tonic_internal_err("failed to deserialize webauthn state")?;
317
318 let result = global
319 .webauthn()
320 .finish_passkey_authentication(reg, &state)
321 .into_tonic_internal_err("failed to finish webauthn authentication")?;
322
323 let counter = result.counter() as i64;
324
325 let credential = mfa_webauthn_credentials::dsl::mfa_webauthn_credentials
326 .filter(mfa_webauthn_credentials::dsl::credential_id.eq(result.cred_id().as_ref()))
327 .select(MfaWebauthnCredential::as_select())
328 .first::<MfaWebauthnCredential>(tx)
329 .await
330 .into_tonic_internal_err("failed to find webauthn credential")?;
331
332 if counter == 0 || credential.counter.is_none_or(|c| c < counter) {
333 diesel::update(mfa_webauthn_credentials::dsl::mfa_webauthn_credentials)
334 .filter(mfa_webauthn_credentials::dsl::credential_id.eq(result.cred_id().as_ref()))
335 .set((
336 mfa_webauthn_credentials::dsl::counter.eq(counter),
337 mfa_webauthn_credentials::dsl::last_used_at.eq(chrono::Utc::now()),
338 ))
339 .execute(tx)
340 .await
341 .into_tonic_internal_err("failed to update webauthn credential")?;
342 } else {
343 diesel::delete(mfa_webauthn_credentials::dsl::mfa_webauthn_credentials)
345 .filter(mfa_webauthn_credentials::dsl::credential_id.eq(result.cred_id().as_ref()))
346 .execute(tx)
347 .await
348 .into_tonic_internal_err("failed to delete webauthn credential")?;
349
350 return Err(tonic::Status::with_error_details(
351 tonic::Code::FailedPrecondition,
352 "invalid webauthn credential",
353 ErrorDetails::new(),
354 ));
355 }
356
357 Ok(())
358}
359
360pub(crate) async fn process_recovery_code(
361 tx: &mut diesel_async::AsyncPgConnection,
362 user_id: UserId,
363 code: &str,
364) -> Result<(), tonic::Status> {
365 let codes = mfa_recovery_codes::dsl::mfa_recovery_codes
366 .filter(mfa_recovery_codes::dsl::user_id.eq(user_id))
367 .limit(20)
368 .load::<MfaRecoveryCode>(tx)
369 .await
370 .into_tonic_internal_err("failed to load MFA recovery codes")?;
371
372 let argon2 = Argon2::default();
373
374 for recovery_code in codes {
375 let hash = argon2::PasswordHash::new(&recovery_code.code_hash)
376 .into_tonic_internal_err("failed to parse recovery code hash")?;
377 match argon2.verify_password(code.as_bytes(), &hash) {
378 Ok(()) => {
379 diesel::delete(mfa_recovery_codes::dsl::mfa_recovery_codes)
380 .filter(mfa_recovery_codes::dsl::id.eq(recovery_code.id))
381 .execute(tx)
382 .await
383 .into_tonic_internal_err("failed to delete recovery code")?;
384
385 break;
386 }
387 Err(argon2::password_hash::Error::Password) => continue,
388 Err(e) => {
389 return Err(e.into_tonic_internal_err("failed to verify recovery code"));
390 }
391 }
392 }
393
394 Ok(())
395}