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