Skip to main content

alibi_plugins/oauth_provider/
extension_claims.rs

1use 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    /// The issuance grant, or `None` when re-deriving opaque access claims.
22    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
77/// Resource policy wins over extension/per-issuance/operator claims, but no source may
78/// supply authority owned by the authorization server. Shared by JWT minting
79/// and live opaque validation (including reserved claims absent from the token).
80pub(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            // The operator overrides additive/per-request claims. The final
220            // resource merge and authority stripping remain at the callers,
221            // using resource policy captured before any contributor executes.
222            // JS stripReservedClaims assigns into {}, omitting __proto__.
223            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}