alibi_plugins/oauth_provider/
extension_claims.rs1use super::{OAuthProviderConfig, OAuthProviderPlugin, OAuthTokenIssueParams};
2use alibi_core::session::SessionRequest;
3use alibi_core::store::OAuthClient;
4use alibi_core::wire::UserView;
5use alibi_core::{AuthContext, AuthResult, AuthSchema, CallbackContext};
6use async_trait::async_trait;
7use serde_json::{Map, Value};
8use std::sync::Arc;
9
10#[derive(Debug, Clone, Default)]
11pub struct OAuthExtensionClaims {
12 pub access_token: Option<Arc<dyn OAuthTokenClaimsContributor>>,
13 pub id_token: Option<Arc<dyn OAuthTokenClaimsContributor>>,
14 pub user_info: Option<Arc<dyn OAuthUserInfoClaimsContributor>>,
15}
16
17pub struct OAuthTokenClaimsInput<'a> {
18 pub context: CallbackContext,
19 pub options: &'a OAuthProviderConfig,
20 pub params: &'a OAuthTokenIssueParams,
21 pub grant_type: Option<&'a str>,
23 pub metadata: Option<Value>,
24}
25
26#[derive(Clone)]
27pub struct OAuthUserInfoClaimsInput<'a> {
28 pub context: CallbackContext,
29 pub options: &'a OAuthProviderConfig,
30 pub user: &'a UserView,
31 pub client: Option<&'a OAuthClient>,
32 pub scopes: &'a [String],
33 pub payload: &'a Value,
34 pub requested_claims: &'a [String],
35}
36
37#[async_trait]
38pub trait OAuthTokenClaimsContributor: Send + Sync + std::fmt::Debug {
39 async fn contribute(&self, input: OAuthTokenClaimsInput<'_>) -> AuthResult<Map<String, Value>>;
40}
41
42#[async_trait]
43pub trait OAuthUserInfoClaimsContributor: Send + Sync + std::fmt::Debug {
44 async fn contribute(
45 &self,
46 input: OAuthUserInfoClaimsInput<'_>,
47 ) -> AuthResult<Map<String, Value>>;
48}
49
50pub(super) enum TokenClaimKind {
51 Access,
52 Id,
53}
54
55pub(super) fn merge_additive(base: &mut Map<String, Value>, claims: Map<String, Value>) {
56 for (key, value) in claims {
57 if !matches!(
58 key.as_str(),
59 "__proto__"
60 | "constructor"
61 | "toString"
62 | "toLocaleString"
63 | "valueOf"
64 | "hasOwnProperty"
65 | "isPrototypeOf"
66 | "propertyIsEnumerable"
67 | "__defineGetter__"
68 | "__defineSetter__"
69 | "__lookupGetter__"
70 | "__lookupSetter__"
71 ) {
72 _ = base.entry(key).or_insert(value);
73 }
74 }
75}
76
77pub(super) fn compose_access_claims(
81 mut claims: Map<String, Value>,
82 resource_claims: Map<String, Value>,
83) -> Map<String, Value> {
84 claims.extend(resource_claims);
85 for key in [
86 "sub",
87 "aud",
88 "client_id",
89 "azp",
90 "scope",
91 "sid",
92 "iss",
93 "iat",
94 "exp",
95 "jti",
96 "cnf",
97 "auth_time",
98 "acr",
99 "amr",
100 ] {
101 _ = claims.remove(key);
102 }
103 claims
104}
105
106pub(super) fn callback_context<S: AuthSchema>(
107 ctx: &AuthContext<S>,
108 req: &impl SessionRequest,
109) -> CallbackContext {
110 let mut callback = CallbackContext::new(ctx, req.physical_request());
111 callback.endpoint = req.endpoint_call().cloned().or(callback.endpoint);
112 callback
113}
114
115impl OAuthProviderPlugin {
116 pub(super) async fn extend_token_response_fields<S: AuthSchema>(
117 &self,
118 ctx: &AuthContext<S>,
119 req: &impl SessionRequest,
120 params: &mut OAuthTokenIssueParams,
121 grant_type: &str,
122 ) -> Result<(), alibi_core::endpoint::EndpointError> {
123 let Some(contributor) = &self.config.custom_token_response_fields else {
124 return Ok(());
125 };
126 let mut fields = contributor
127 .contribute(OAuthTokenClaimsInput {
128 context: callback_context(ctx, req),
129 options: &self.config,
130 params,
131 grant_type: Some(grant_type),
132 metadata: alibi_core::store::oauth_clients::metadata_to_columns(
133 params.client.metadata.clone(),
134 )?
135 .get("metadata")
136 .filter(|value| !value.is_null())
137 .cloned(),
138 })
139 .await
140 .map_err(|error| {
141 if alibi_core::endpoint::is_endpoint_api_error(&error) {
142 error
143 } else {
144 crate::helpers::callback_failure(error)
145 }
146 })?;
147 if let Some(per_request) = params.token_response.as_ref().and_then(Value::as_object) {
148 fields.extend(per_request.clone());
149 }
150 params.token_response = Some(Value::Object(fields));
151 Ok(())
152 }
153
154 pub(super) async fn collect_token_claims(
155 &self,
156 context: CallbackContext,
157 params: &OAuthTokenIssueParams,
158 grant_type: Option<&str>,
159 kind: TokenClaimKind,
160 ) -> AuthResult<Map<String, Value>> {
161 let mut claims = Map::new();
162 for extension in &self.config.extensions {
163 let contributor = match kind {
164 TokenClaimKind::Access => extension.claims.access_token.as_ref(),
165 TokenClaimKind::Id => extension.claims.id_token.as_ref(),
166 };
167 if let Some(contributor) = contributor {
168 let contribution = contributor
169 .contribute(OAuthTokenClaimsInput {
170 context: context.clone(),
171 options: &self.config,
172 params,
173 grant_type,
174 metadata: alibi_core::store::oauth_clients::metadata_to_columns(
175 params.client.metadata.clone(),
176 )?
177 .get("metadata")
178 .filter(|value| !value.is_null())
179 .cloned(),
180 })
181 .await?;
182 merge_additive(&mut claims, contribution);
183 }
184 }
185 Ok(claims)
186 }
187
188 pub(super) async fn resolve_access_token_claims(
189 &self,
190 context: CallbackContext,
191 params: &OAuthTokenIssueParams,
192 grant_type: Option<&str>,
193 ) -> AuthResult<Map<String, Value>> {
194 let mut claims = self
195 .collect_token_claims(context.clone(), params, grant_type, TokenClaimKind::Access)
196 .await?;
197 if let Some(per_request) = params
198 .access_token_claims
199 .as_ref()
200 .and_then(Value::as_object)
201 {
202 claims.extend(per_request.clone());
203 }
204 if let Some(contributor) = &self.config.custom_access_token_claims {
205 let operator = contributor
206 .contribute(OAuthTokenClaimsInput {
207 context,
208 options: &self.config,
209 params,
210 grant_type,
211 metadata: alibi_core::store::oauth_clients::metadata_to_columns(
212 params.client.metadata.clone(),
213 )?
214 .get("metadata")
215 .filter(|value| !value.is_null())
216 .cloned(),
217 })
218 .await?;
219 claims.extend(operator.into_iter().filter(|(key, _)| key != "__proto__"));
224 }
225 Ok(claims)
226 }
227
228 pub(super) async fn extend_access_token_claims<S: AuthSchema>(
229 &self,
230 ctx: &AuthContext<S>,
231 req: &impl SessionRequest,
232 params: &mut OAuthTokenIssueParams,
233 grant_type: &str,
234 ) -> AuthResult<()> {
235 let claims = self
236 .resolve_access_token_claims(callback_context(ctx, req), params, Some(grant_type))
237 .await?;
238 params.access_token_claims = Some(Value::Object(claims));
239 Ok(())
240 }
241
242 pub(super) async fn extend_id_token_claims<S: AuthSchema>(
243 &self,
244 ctx: &AuthContext<S>,
245 req: &impl SessionRequest,
246 params: &mut OAuthTokenIssueParams,
247 grant_type: &str,
248 ) -> AuthResult<()> {
249 if params.user.is_some() && params.scopes.iter().any(|scope| scope == "openid") {
250 let mut claims = self
251 .collect_token_claims(
252 callback_context(ctx, req),
253 params,
254 Some(grant_type),
255 TokenClaimKind::Id,
256 )
257 .await?;
258 if let Some(per_request) = params.id_token_claims.as_ref().and_then(Value::as_object) {
259 claims.extend(per_request.clone());
260 }
261 params.id_token_claims = Some(Value::Object(claims));
262 }
263 Ok(())
264 }
265
266 pub(super) async fn collect_user_info_claims(
267 &self,
268 input: OAuthUserInfoClaimsInput<'_>,
269 ) -> AuthResult<Map<String, Value>> {
270 let mut claims = Map::new();
271 for extension in &self.config.extensions {
272 if let Some(contributor) = &extension.claims.user_info {
273 let contribution = contributor.contribute(input.clone()).await?;
274 merge_additive(&mut claims, contribution);
275 }
276 }
277 Ok(claims)
278 }
279}