1use std::collections::HashMap;
5
6use reifydb_catalog::catalog::Catalog;
7use reifydb_core::interface::{
8 auth::{AuthStep, AuthenticationProvider},
9 catalog::{authentication::Authentication, identity::Identity},
10};
11use reifydb_transaction::transaction::{Transaction, query::QueryTransaction};
12use reifydb_value::{
13 error::Error,
14 reifydb_assertions,
15 value::{Value, identity::IdentityId},
16};
17use tracing::instrument;
18
19use super::{AuthResponse, AuthService, generate_session_token, solana::SOLANA_PUBLIC_KEY_ATTRIBUTE};
20use crate::error::AuthError;
21
22impl AuthService {
23 #[instrument(name = "auth::authenticate", level = "debug", skip(self, credentials))]
24 pub fn authenticate(&self, method: &str, credentials: HashMap<String, String>) -> Result<AuthResponse, Error> {
25 if let Some(challenge_id) = credentials.get("challenge_id").cloned() {
26 return self.authenticate_challenge_response(&challenge_id, credentials);
27 }
28 if method == "token" {
29 return self.authenticate_token(credentials);
30 }
31 if method == "github" {
32 return self.begin_github_login();
33 }
34 self.authenticate_with_provider(method, credentials)
35 }
36
37 fn authenticate_with_provider(
38 &self,
39 method: &str,
40 credentials: HashMap<String, String>,
41 ) -> Result<AuthResponse, Error> {
42 let identifier = credentials.get("identifier").map(|s| s.as_str()).unwrap_or("");
43 let mut txn = self.engine.begin_query()?;
44 let catalog = self.engine.catalog();
45
46 let Some(ident) = self.resolve_provider_identity(&mut txn, &catalog, method, identifier)? else {
47 drop(txn);
48 return self.handle_missing_identity(method, identifier, &credentials);
49 };
50 if !ident.enabled {
51 return Ok(AuthResponse::Failed {
52 reason: "identity is disabled".to_string(),
53 });
54 }
55
56 let Some(stored_auth) = self.load_stored_auth(&mut txn, &catalog, ident.id, method)? else {
57 return Ok(invalid_credentials());
58 };
59
60 self.run_provider_and_respond(&stored_auth, &credentials, ident.id, identifier, method)
61 }
62
63 #[inline]
64 fn resolve_provider_identity(
65 &self,
66 txn: &mut QueryTransaction,
67 catalog: &Catalog,
68 method: &str,
69 identifier: &str,
70 ) -> Result<Option<Identity>, Error> {
71 if let Some(u) = catalog.find_identity_by_name(&mut Transaction::Query(txn), identifier)? {
72 return Ok(Some(u));
73 }
74 if method == "solana" {
75 return catalog.find_identity_by_attribute_value(
76 &mut Transaction::Query(txn),
77 SOLANA_PUBLIC_KEY_ATTRIBUTE,
78 &Value::Utf8(identifier.to_string()),
79 );
80 }
81 Ok(None)
82 }
83
84 fn load_stored_auth(
85 &self,
86 txn: &mut QueryTransaction,
87 catalog: &Catalog,
88 identity: IdentityId,
89 method: &str,
90 ) -> Result<Option<Authentication>, Error> {
91 catalog.find_authentication_by_identity_and_method(&mut Transaction::Query(txn), identity, method)
92 }
93
94 #[inline]
95 fn run_provider_and_respond(
96 &self,
97 stored_auth: &Authentication,
98 credentials: &HashMap<String, String>,
99 identity: IdentityId,
100 identifier: &str,
101 method: &str,
102 ) -> Result<AuthResponse, Error> {
103 let provider = self.provider_for(method)?;
104 let step = provider.authenticate(&stored_auth.properties, credentials)?;
105 self.respond_to_initial_auth_step(step, identity, identifier, method)
106 }
107
108 #[inline]
109 fn handle_missing_identity(
110 &self,
111 method: &str,
112 identifier: &str,
113 credentials: &HashMap<String, String>,
114 ) -> Result<AuthResponse, Error> {
115 if method == "solana"
116 && let Some(public_key) = credentials.get("public_key").cloned()
117 {
118 return self.auto_provision_solana(identifier, &public_key, credentials);
119 }
120 Ok(invalid_credentials())
121 }
122
123 #[inline]
124 fn respond_to_initial_auth_step(
125 &self,
126 step: AuthStep,
127 identity: IdentityId,
128 identifier: &str,
129 method: &str,
130 ) -> Result<AuthResponse, Error> {
131 match step {
132 AuthStep::Authenticated => self.finalize_authentication(identity),
133 AuthStep::Failed => Ok(invalid_credentials()),
134 AuthStep::Challenge {
135 payload,
136 } => Ok(self.issue_challenge(identifier, method, payload)),
137 }
138 }
139
140 #[inline]
141 pub(super) fn finalize_authentication(&self, identity: IdentityId) -> Result<AuthResponse, Error> {
142 reifydb_assertions! {
143 assert!(
144 identity != IdentityId::default(),
145 "authentication finalized for the nil placeholder identity instead of a resolved one, so an unauthenticated principal would receive a valid session token and gain authorization (identity={:?})",
146 identity
147 );
148 }
149 let token = generate_session_token(&self.rng);
150 self.persist_token(&token, identity)?;
151 Ok(AuthResponse::Authenticated {
152 identity,
153 token,
154 })
155 }
156
157 #[inline]
158 fn issue_challenge(&self, identifier: &str, method: &str, payload: HashMap<String, String>) -> AuthResponse {
159 let challenge_id = self.challenges.create(
160 identifier.to_string(),
161 method.to_string(),
162 payload.clone(),
163 &self.clock,
164 &self.rng,
165 );
166 AuthResponse::Challenge {
167 challenge_id,
168 payload,
169 }
170 }
171
172 #[inline]
173 fn provider_for(&self, method: &str) -> Result<&dyn AuthenticationProvider, Error> {
174 self.auth_registry.get(method).ok_or_else(|| {
175 Error::from(AuthError::UnknownMethod {
176 method: method.to_string(),
177 })
178 })
179 }
180
181 fn authenticate_token(&self, credentials: HashMap<String, String>) -> Result<AuthResponse, Error> {
182 let token_value = match credentials.get("token") {
183 Some(t) if !t.is_empty() => t,
184 _ => return Ok(invalid_credentials()),
185 };
186
187 match self.validate_token(token_value) {
188 Some(token) => self.finalize_authentication(token.identity),
189 None => Ok(invalid_credentials()),
190 }
191 }
192
193 fn authenticate_challenge_response(
194 &self,
195 challenge_id: &str,
196 mut credentials: HashMap<String, String>,
197 ) -> Result<AuthResponse, Error> {
198 let Some(challenge) = self.challenges.consume(challenge_id) else {
199 return Ok(AuthResponse::Failed {
200 reason: "invalid or expired challenge".to_string(),
201 });
202 };
203
204 if challenge.method == "github" {
205 return self.complete_github_login(&challenge, &credentials);
206 }
207
208 merge_challenge_payload(&mut credentials, &challenge.payload);
209
210 let mut txn = self.engine.begin_query()?;
211 let catalog = self.engine.catalog();
212
213 let Some(ident) =
214 self.resolve_challenge_identity(&mut txn, &catalog, &challenge.identifier, &challenge.method)?
215 else {
216 return Ok(invalid_credentials());
217 };
218
219 let Some(stored_auth) = self.load_stored_auth(&mut txn, &catalog, ident.id, &challenge.method)? else {
220 return Ok(invalid_credentials());
221 };
222
223 self.run_challenge_provider_and_respond(&stored_auth, &credentials, ident.id, &challenge.method)
224 }
225
226 #[inline]
227 fn resolve_challenge_identity(
228 &self,
229 txn: &mut QueryTransaction,
230 catalog: &Catalog,
231 identifier: &str,
232 method: &str,
233 ) -> Result<Option<Identity>, Error> {
234 let resolved = match catalog.find_identity_by_name(&mut Transaction::Query(txn), identifier)? {
235 Some(u) if u.enabled => Some(u),
236 Some(_) => None,
237 None if method == "solana" => {
238 match catalog.find_identity_by_attribute_value(
239 &mut Transaction::Query(txn),
240 SOLANA_PUBLIC_KEY_ATTRIBUTE,
241 &Value::Utf8(identifier.to_string()),
242 )? {
243 Some(u) if u.enabled => Some(u),
244 _ => None,
245 }
246 }
247 None => None,
248 };
249 reifydb_assertions! {
250 if let Some(ref ident) = resolved {
251 assert!(
252 ident.enabled,
253 "challenge identity resolution returned a disabled identity (id={:?}, name={}); a disabled principal must never advance to provider authentication or it could obtain a session token",
254 ident.id,
255 ident.name
256 );
257 }
258 }
259 Ok(resolved)
260 }
261
262 #[inline]
263 fn run_challenge_provider_and_respond(
264 &self,
265 stored_auth: &Authentication,
266 credentials: &HashMap<String, String>,
267 identity: IdentityId,
268 method: &str,
269 ) -> Result<AuthResponse, Error> {
270 let provider = self.provider_for(method)?;
271 let step = provider.authenticate(&stored_auth.properties, credentials)?;
272 respond_to_challenge_step(step, identity, self)
273 }
274}
275
276#[inline]
277fn merge_challenge_payload(credentials: &mut HashMap<String, String>, payload: &HashMap<String, String>) {
278 for (k, v) in payload {
279 credentials.entry(k.clone()).or_insert_with(|| v.clone());
280 }
281 credentials.remove("challenge_id");
282}
283
284#[inline]
285fn respond_to_challenge_step(
286 step: AuthStep,
287 identity: IdentityId,
288 service: &AuthService,
289) -> Result<AuthResponse, Error> {
290 match step {
291 AuthStep::Authenticated => service.finalize_authentication(identity),
292 AuthStep::Failed => Ok(invalid_credentials()),
293 AuthStep::Challenge {
294 ..
295 } => Ok(AuthResponse::Failed {
296 reason: "nested challenges are not supported".to_string(),
297 }),
298 }
299}
300
301#[inline]
302fn invalid_credentials() -> AuthResponse {
303 AuthResponse::Failed {
304 reason: "invalid credentials".to_string(),
305 }
306}