refactor(jwt): improve the implementation of jwt plugin and expose it to expressions (#534)

Author: dotansimha

Cargo.lock

@@ -44,6 +44,7 @@ retry-policies = { workspace = true}
 reqwest-retry = { workspace = true }
 reqwest-middleware = { workspace = true }
 vrl = { workspace = true }
+serde_json = { workspace = true }
 
 mimalloc = { version = "0.1.48", features = ["v3"] }
 moka = { version = "0.12.10", features = ["future"] }

bin/router/Cargo.toml

@@ -1,32 +1,44 @@
 use std::collections::HashMap;
 
-use hive_router_plan_executor::execution::jwt_forward::JwtAuthForwardingPlan;
+use hive_router_plan_executor::execution::jwt_forward::JwtForwardingError;
 use jsonwebtoken::TokenData;
 use serde::{Deserialize, Serialize};
-use sonic_rs::Value;
-
-use crate::jwt::errors::JwtForwardingError;
 
 pub type JwtTokenPayload = TokenData<JwtClaims>;
 
 #[derive(Debug, Clone)]
 pub struct JwtRequestContext {
-    // The payload extracted from the JWT token, and the extensions key to inject it into the request
-    pub token_payload: Option<(String, JwtTokenPayload)>,
+    pub token_prefix: Option<String>,
+    pub token_raw: String,
+    pub token_payload: JwtTokenPayload,
 }
 
-impl TryInto<Option<JwtAuthForwardingPlan>> for JwtRequestContext {
-    type Error = JwtForwardingError;
+impl JwtRequestContext {
+    pub fn get_claims_value(&self) -> Result<sonic_rs::Value, JwtForwardingError> {
+        Ok(sonic_rs::to_value(&self.token_payload.claims)?)
+    }
+
+    /// Extracts an optional "scope"/"scopes" field form the token's payload.
+    /// Supports both space-delimited and array formats.
+    pub fn extract_scopes(&self) -> Option<Vec<String>> {
+        let map = &self.token_payload.claims.additional_claims;
+        let maybe_scopes = map.get("scope").or_else(|| map.get("scopes"));
+
+        if let Some(serde_json::Value::String(scopes_str)) = maybe_scopes {
+            return Some(scopes_str.split(' ').map(String::from).collect());
+        }
 
-    fn try_into(self) -> Result<Option<JwtAuthForwardingPlan>, Self::Error> {
-        if let Some((extension_field_name, payload)) = &self.token_payload {
-            return Ok(Some(JwtAuthForwardingPlan {
-                extension_field_name: extension_field_name.clone(),
-                extension_field_value: sonic_rs::to_value(&payload.claims)?,
-            }));
+        if let Some(serde_json::Value::Array(scopes_arr)) = maybe_scopes {
+            return Some(
+                scopes_arr
+                    .iter()
+                    .filter_map(|s| s.as_str())
+                    .map(String::from)
+                    .collect::<Vec<_>>(),
+            );
         }
 
-        Ok(None)
+        None
     }
 }
 
@@ -61,6 +73,8 @@ pub struct JwtClaims {
     #[serde(skip_serializing_if = "Option::is_none")]
     pub jti: Option<String>,
 
+    // we are using serde to deserialize the additional claims
+    // because the jsonwebtoken crate is using `serde_json` internally, and the `sonic_rs::Value` is not recognized as valid type
     #[serde(flatten)]
-    pub additional_claims: HashMap<String, Value>,
+    pub additional_claims: HashMap<String, serde_json::Value>,
 }

bin/router/src/jwt/context.rs

@@ -94,11 +94,3 @@ impl From<&JwtError> for GraphQLError {
         }
     }
 }
-
-#[derive(Debug, thiserror::Error)]
-pub enum JwtForwardingError {
-    #[error("failed to serialized jwt claims")]
-    ClaimsSerializeError(#[from] sonic_rs::Error),
-    #[error("failed to parse  as valid header value")]
-    ValueIsNotValidHeader(#[from] http::header::InvalidHeaderValue),
-}

bin/router/src/jwt/errors.rs

@@ -51,7 +51,7 @@ impl JwtAuthRuntime {
         Ok(instance)
     }
 
-    fn lookup(&self, req: &HttpRequest) -> Result<String, LookupError> {
+    fn lookup(&self, req: &HttpRequest) -> Result<(Option<String>, String), LookupError> {
         for lookup_config in &self.config.lookup_locations {
             match lookup_config {
                 JwtAuthPluginLookupLocation::Header { name, prefix } => {
@@ -73,14 +73,17 @@ impl JwtAuthRuntime {
                                 .and_then(|s| s.strip_prefix(prefix))
                             {
                                 Some(stripped_value) => {
-                                    return Ok(stripped_value.trim().to_string());
+                                    return Ok((
+                                        Some(prefix.to_string()),
+                                        stripped_value.trim().to_string(),
+                                    ));
                                 }
                                 None => {
                                     return Err(LookupError::MismatchedPrefix);
                                 }
                             },
                             None => {
-                                return Ok(header_value.to_str().unwrap_or("").to_string());
+                                return Ok((None, header_value.to_str().unwrap_or("").to_string()));
                             }
                         }
                     }
@@ -101,7 +104,7 @@ impl JwtAuthRuntime {
                                     let (cookie_name, cookie_value) = v.name_value_trimmed();
 
                                     if cookie_name == name {
-                                        return Ok(cookie_value.to_string());
+                                        return Ok((None, cookie_value.to_string()));
                                     }
                                 }
                                 Err(e) => {
@@ -158,15 +161,15 @@ impl JwtAuthRuntime {
         &self,
         jwks: &Vec<Arc<JwkSet>>,
         req: &HttpRequest,
-    ) -> Result<(JwtTokenPayload, String), JwtError> {
+    ) -> Result<(JwtTokenPayload, Option<String>, String), JwtError> {
         match self.lookup(req) {
-            Ok(token) => {
+            Ok((maybe_prefix, token)) => {
                 // First, we need to decode the header to determine which provider to use.
                 let header = decode_header(&token).map_err(JwtError::InvalidJwtHeader)?;
                 let jwk = self.find_matching_jwks(&header, jwks)?;
 
                 self.decode_and_validate_token(&token, &jwk.keys)
-                    .map(|token_data| (token_data, token))
+                    .map(|token_data| (token_data, maybe_prefix, token))
             }
             Err(e) => {
                 warn!("jwt plugin failed to lookup token. error: {}", e);
@@ -266,25 +269,15 @@ impl JwtAuthRuntime {
         let valid_jwks = self.jwks.all();
 
         match self.authenticate(&valid_jwks, request) {
-            Ok((token_data, _token)) => {
-                let mut jwt_ctx = JwtRequestContext {
-                    token_payload: None,
-                };
-
-                if self.config.forward_claims_to_upstream_extensions.enabled {
-                    jwt_ctx.token_payload = Some((
-                        self.config
-                            .forward_claims_to_upstream_extensions
-                            .field_name
-                            .clone(),
-                        token_data,
-                    ));
-                }
-
-                request.extensions_mut().insert(jwt_ctx);
+            Ok((token_payload, maybe_token_prefix, token)) => {
+                request.extensions_mut().insert(JwtRequestContext {
+                    token_payload,
+                    token_raw: token,
+                    token_prefix: maybe_token_prefix,
+                });
             }
             Err(e) => {
-                warn!("jwt token error: {}", e);
+                warn!("jwt token error: {:?}", e);
 
                 if self.config.require_authentication.is_some_and(|v| v) {
                     return Err(e);

bin/router/src/jwt/mod.rs

@@ -106,9 +106,9 @@ pub async fn configure_app_from_config(
     router_config: HiveRouterConfig,
     bg_tasks_manager: &mut BackgroundTasksManager,
 ) -> Result<(Arc<RouterSharedState>, Arc<SchemaState>), Box<dyn std::error::Error>> {
-    let jwt_runtime = match &router_config.jwt {
-        Some(jwt_config) => Some(JwtAuthRuntime::init(bg_tasks_manager, jwt_config).await?),
-        None => None,
+    let jwt_runtime = match router_config.jwt.is_jwt_auth_enabled() {
+        true => Some(JwtAuthRuntime::init(bg_tasks_manager, &router_config.jwt).await?),
+        false => None,
     };
 
     let router_config_arc = Arc::new(router_config);

bin/router/src/lib.rs

@@ -2,7 +2,7 @@ use std::sync::Arc;
 
 use graphql_tools::validation::utils::ValidationError;
 use hive_router_plan_executor::{
-    execution::error::PlanExecutionError,
+    execution::{error::PlanExecutionError, jwt_forward::JwtForwardingError},
     response::graphql_error::{GraphQLError, GraphQLErrorExtensions},
 };
 use hive_router_query_planner::{
@@ -15,12 +15,9 @@ use ntex::{
 };
 use serde::{Deserialize, Serialize};
 
-use crate::{
-    jwt::errors::JwtForwardingError,
-    pipeline::{
-        header::{RequestAccepts, APPLICATION_GRAPHQL_RESPONSE_JSON_STR},
-        progressive_override::LabelEvaluationError,
-    },
+use crate::pipeline::{
+    header::{RequestAccepts, APPLICATION_GRAPHQL_RESPONSE_JSON_STR},
+    progressive_override::LabelEvaluationError,
 };
 
 #[derive(Debug)]

bin/router/src/pipeline/error.rs

@@ -1,17 +1,15 @@
 use std::collections::HashMap;
 use std::sync::Arc;
 
-use crate::jwt::context::JwtRequestContext;
 use crate::pipeline::coerce_variables::CoerceVariablesPayload;
 use crate::pipeline::error::{PipelineError, PipelineErrorFromAcceptHeader, PipelineErrorVariant};
 use crate::pipeline::normalize::GraphQLNormalizationPayload;
 use crate::schema_state::SupergraphData;
 use crate::shared_state::RouterSharedState;
 use hive_router_plan_executor::execute_query_plan;
+use hive_router_plan_executor::execution::client_request_details::ClientRequestDetails;
 use hive_router_plan_executor::execution::jwt_forward::JwtAuthForwardingPlan;
-use hive_router_plan_executor::execution::plan::{
-    ClientRequestDetails, PlanExecutionOutput, QueryPlanExecutionContext,
-};
+use hive_router_plan_executor::execution::plan::{PlanExecutionOutput, QueryPlanExecutionContext};
 use hive_router_plan_executor::introspection::resolve::IntrospectionContext;
 use hive_router_query_planner::planner::plan_nodes::QueryPlan;
 use http::HeaderName;
@@ -67,15 +65,23 @@ pub async fn execute_plan(
         metadata: &supergraph.metadata,
     };
 
-    let jwt_context = {
-        let req_extensions = req.extensions();
-        req_extensions.get::<JwtRequestContext>().cloned()
-    };
-    let jwt_forward_plan: Option<JwtAuthForwardingPlan> = match jwt_context {
-        Some(jwt_context) => jwt_context
-            .try_into()
-            .map_err(|e| req.new_pipeline_error(PipelineErrorVariant::JwtForwardingError(e)))?,
-        None => None,
+    let jwt_forward_plan: Option<JwtAuthForwardingPlan> = if app_state
+        .router_config
+        .jwt
+        .is_jwt_extensions_forwarding_enabled()
+    {
+        client_request_details
+            .jwt
+            .build_forwarding_plan(
+                &app_state
+                    .router_config
+                    .jwt
+                    .forward_claims_to_upstream_extensions
+                    .field_name,
+            )
+            .map_err(|e| req.new_pipeline_error(PipelineErrorVariant::JwtForwardingError(e)))?
+    } else {
+        None
     };
 
     execute_query_plan(QueryPlanExecutionContext {

bin/router/src/pipeline/execution.rs

@@ -1,7 +1,8 @@
 use std::sync::Arc;
 
-use hive_router_plan_executor::execution::plan::{
-    ClientRequestDetails, OperationDetails, PlanExecutionOutput,
+use hive_router_plan_executor::execution::{
+    client_request_details::{ClientRequestDetails, JwtRequestDetails, OperationDetails},
+    plan::PlanExecutionOutput,
 };
 use hive_router_query_planner::{
     state::supergraph_state::OperationKind, utils::cancellation::CancellationToken,
@@ -13,6 +14,7 @@ use ntex::{
 };
 
 use crate::{
+    jwt::context::JwtRequestContext,
     pipeline::{
         coerce_variables::coerce_request_variables,
         csrf_prevention::perform_csrf_prevention,
@@ -101,6 +103,7 @@ pub async fn graphql_request_handler(
 }
 
 #[inline]
+#[allow(clippy::await_holding_refcell_ref)]
 pub async fn execute_pipeline(
     req: &mut HttpRequest,
     body_bytes: Bytes,
@@ -129,6 +132,20 @@ pub async fn execute_pipeline(
     let query_plan_cancellation_token =
         CancellationToken::with_timeout(shared_state.router_config.query_planner.timeout);
 
+    let req_extensions = req.extensions();
+    let jwt_context = req_extensions.get::<JwtRequestContext>();
+    let jwt_request_details = match jwt_context {
+        Some(jwt_context) => JwtRequestDetails::Authenticated {
+            token: jwt_context.token_raw.as_str(),
+            prefix: jwt_context.token_prefix.as_deref(),
+            scopes: jwt_context.extract_scopes(),
+            claims: &jwt_context
+                .get_claims_value()
+                .map_err(|e| req.new_pipeline_error(PipelineErrorVariant::JwtForwardingError(e)))?,
+        },
+        None => JwtRequestDetails::Unauthenticated,
+    };
+
     let client_request_details = ClientRequestDetails {
         method: req.method(),
         url: req.uri(),
@@ -143,6 +160,7 @@ pub async fn execute_pipeline(
             },
             query: &execution_request.query,
         },
+        jwt: &jwt_request_details,
     };
 
     let progressive_override_ctx = request_override_context(

bin/router/src/pipeline/mod.rs

@@ -1,7 +1,7 @@
 use std::collections::{BTreeMap, HashMap, HashSet};
 
 use hive_router_config::override_labels::{LabelOverrideValue, OverrideLabelsConfig};
-use hive_router_plan_executor::execution::plan::ClientRequestDetails;
+use hive_router_plan_executor::execution::client_request_details::ClientRequestDetails;
 use hive_router_query_planner::{
     graph::{PlannerOverrideContext, PERCENTAGE_SCALE_FACTOR},
     state::supergraph_state::SupergraphState,