feat(jit): plan ahead auth requirements when required (#3139)

Co-authored-by: Tushar Mathur <[email protected]>

src/core/jit/exec_const.rs

@@ -6,14 +6,12 @@ use tailcall_valid::Validator;
 
 use super::context::Context;
 use super::exec::{Executor, IRExecutor};
-use super::{
-    transform, AnyResponse, BuildError, Error, OperationPlan, Pos, Positioned, Request, Response,
-    Result,
-};
+use super::graphql_error::GraphQLError;
+use super::{transform, AnyResponse, BuildError, Error, OperationPlan, Request, Response, Result};
 use crate::core::app_context::AppContext;
 use crate::core::http::RequestContext;
 use crate::core::ir::model::IR;
-use crate::core::ir::{self, EvalContext};
+use crate::core::ir::{self, EmptyResolverContext, EvalContext};
 use crate::core::jit::synth::Synth;
 use crate::core::jit::transform::InputResolver;
 use crate::core::json::{JsonLike, JsonLikeList};
@@ -42,6 +40,20 @@ impl ConstValueExecutor {
         req_ctx: &RequestContext,
         request: Request<ConstValue>,
     ) -> AnyResponse<Vec<u8>> {
+        // Run all the IRs in the before chain
+        if let Some(ir) = &self.plan.before {
+            let mut eval_context = EvalContext::new(req_ctx, &EmptyResolverContext {});
+            match ir.eval(&mut eval_context).await {
+                Ok(_) => (),
+                Err(err) => {
+                    let resp: Response<ConstValue> = Response::default();
+                    return resp
+                        .with_errors(vec![GraphQLError::new(err.to_string(), None)])
+                        .into();
+                }
+            }
+        }
+
         let is_introspection_query =
             req_ctx.server.get_enable_introspection() && self.plan.is_introspection_query;
         let variables = &request.variables;
@@ -54,7 +66,7 @@ impl ConstValueExecutor {
             let resp: Response<ConstValue> = Response::default();
             // this shouldn't actually ever happen
             return resp
-                .with_errors(vec![Positioned::new(Error::Unknown, Pos::default())])
+                .with_errors(vec![GraphQLError::new(Error::Unknown.to_string(), None)])
                 .into();
         };
 
@@ -69,9 +81,9 @@ impl ConstValueExecutor {
             Err(err) => {
                 let resp: Response<ConstValue> = Response::default();
                 return resp
-                    .with_errors(vec![Positioned::new(
-                        BuildError::from(err).into(),
-                        Pos::default(),
+                    .with_errors(vec![GraphQLError::new(
+                        BuildError::from(err).to_string(),
+                        None,
                     )])
                     .into();
             }
@@ -128,7 +140,7 @@ impl<'a> ConstValueExec<'a> {
     }
 }
 
-impl<'ctx> IRExecutor for ConstValueExec<'ctx> {
+impl IRExecutor for ConstValueExec<'_> {
     type Input = ConstValue;
     type Output = ConstValue;
     type Error = Error;

src/core/jit/model.rs

@@ -290,6 +290,9 @@ pub struct OperationPlan<Input> {
     pub is_protected: bool,
     pub min_cache_ttl: Option<NonZeroU64>,
     pub selection: Vec<Field<Input>>,
+
+    /// An IR that should be executed before the operation starts executing
+    pub before: Option<IR>,
 }
 
 impl<Input> std::fmt::Debug for OperationPlan<Input> {
@@ -312,15 +315,16 @@ impl<Input> OperationPlan<Input> {
         }
 
         Ok(OperationPlan {
+            selection,
             root_name: self.root_name,
             operation_type: self.operation_type,
-            selection,
             index: self.index,
             is_introspection_query: self.is_introspection_query,
             is_dedupe: self.is_dedupe,
             is_const: self.is_const,
-            min_cache_ttl: self.min_cache_ttl,
             is_protected: self.is_protected,
+            min_cache_ttl: self.min_cache_ttl,
+            before: self.before,
         })
     }
 }
@@ -345,8 +349,9 @@ impl<Input> OperationPlan<Input> {
             is_introspection_query,
             is_dedupe: false,
             is_const: false,
-            min_cache_ttl: None,
             is_protected: false,
+            min_cache_ttl: None,
+            before: Default::default(),
         }
     }
 

src/core/jit/request.rs

@@ -46,8 +46,9 @@ impl Request<ConstValue> {
         let plan = builder.build(self.operation_name.as_deref())?;
 
         transform::CheckConst::new()
-            .pipe(transform::CheckDedupe::new())
             .pipe(transform::CheckProtected::new())
+            .pipe(transform::AuthPlaner::new().when(blueprint.server.auth.is_some()))
+            .pipe(transform::CheckDedupe::new())
             .pipe(transform::CheckCache::new())
             .transform(plan)
             .to_result()

src/core/jit/transform/auth_planer.rs

@@ -0,0 +1,86 @@
+use std::convert::Infallible;
+
+use tailcall_valid::Valid;
+
+use crate::core::blueprint::DynamicValue;
+use crate::core::ir::model::IR;
+use crate::core::jit::{Field, OperationPlan};
+use crate::core::Transform;
+
+pub struct AuthPlaner<A> {
+    marker: std::marker::PhantomData<A>,
+}
+
+impl<A> AuthPlaner<A> {
+    pub fn new() -> Self {
+        Self { marker: std::marker::PhantomData }
+    }
+}
+
+impl<A> Transform for AuthPlaner<A> {
+    type Value = OperationPlan<A>;
+    type Error = Infallible;
+
+    fn transform(&self, mut plan: Self::Value) -> Valid<Self::Value, Self::Error> {
+        let mut before = Vec::new();
+
+        plan.selection = plan
+            .selection
+            .into_iter()
+            .map(|field| drop_protect_field(&mut before, field))
+            .collect();
+        let ir = before.into_iter();
+
+        plan.before = match plan.before {
+            Some(before) => Some(ir.fold(before, |a, b| a.pipe(b))),
+            None => ir.reduce(|a, b| a.pipe(b)),
+        };
+        Valid::succeed(plan)
+    }
+}
+
+/// Used to recursively update the field ands its selections to remove
+/// IR::Protected
+fn drop_protect_field<A>(before: &mut Vec<IR>, mut field: Field<A>) -> Field<A> {
+    if let Some(mut ir) = field.ir {
+        let is_protected = drop_protect_ir(&mut ir);
+
+        field.selection = field
+            .selection
+            .into_iter()
+            .map(|field| drop_protect_field(before, field))
+            .collect();
+
+        if is_protected {
+            let ir = IR::Protect(Box::new(IR::Dynamic(DynamicValue::Value(
+                Default::default(),
+            ))));
+            before.push(ir);
+        }
+
+        field.ir = Some(ir);
+    }
+    field
+}
+
+/// This function modifies an IR chain by detecting and removing any
+/// instances of IR::Protect from the chain. Returns `true` when it modifies the
+/// IR.
+pub fn drop_protect_ir(ir: &mut IR) -> bool {
+    match ir {
+        IR::Dynamic(_) => false,
+        IR::IO(_) => false,
+        IR::Cache(_) => false,
+        IR::Path(ir, _) => drop_protect_ir(ir),
+        IR::ContextPath(_) => false,
+        IR::Protect(inner_ir) => {
+            *ir = *inner_ir.clone();
+            true
+        }
+        IR::Map(_) => false,
+        IR::Pipe(ir1, ir2) => drop_protect_ir(ir1) || drop_protect_ir(ir2),
+        IR::Discriminate(_, ir) => drop_protect_ir(ir),
+        IR::Entity(_) => false,
+        IR::Service(_) => false,
+    }
+}

src/core/jit/transform/input_resolver.rs

@@ -78,6 +78,7 @@ where
             is_protected: self.plan.is_protected,
             min_cache_ttl: self.plan.min_cache_ttl,
             selection,
+            before: self.plan.before,
         })
     }
 

src/core/jit/transform/mod.rs

@@ -1,10 +1,12 @@
+mod auth_planer;
 mod check_cache;
 mod check_const;
 mod check_dedupe;
 mod check_protected;
 mod input_resolver;
 mod skip;
 
+pub use auth_planer::*;
 pub use check_cache::*;
 pub use check_const::*;
 pub use check_dedupe::*;

src/core/rest/partial_request.rs

@@ -15,7 +15,7 @@ pub struct PartialRequest<'a> {
     pub path: &'a Path,
 }
 
-impl<'a> PartialRequest<'a> {
+impl PartialRequest<'_> {
     pub async fn into_request(self, request: Request) -> Result<GraphQLRequest> {
         let mut variables = self.variables;
         if let Some(key) = self.body {

tests/core/snapshots/auth-basic.md_1.snap

@@ -1,6 +1,7 @@
 ---
 source: tests/core/spec.rs
 expression: response
+snapshot_kind: text
 ---
 {
   "status": 200,
@@ -11,13 +12,7 @@ expression: response
     "data": null,
     "errors": [
       {
-        "message": "Authentication Failure: Missing Authorization Header",
-        "locations": [
-          {
-            "line": 2,
-            "column": 3
-          }
-        ]
+        "message": "Authentication Failure: Missing Authorization Header"
       }
     ]
   }

tests/core/snapshots/auth-basic.md_2.snap

@@ -1,6 +1,7 @@
 ---
 source: tests/core/spec.rs
 expression: response
+snapshot_kind: text
 ---
 {
   "status": 200,
@@ -11,13 +12,7 @@ expression: response
     "data": null,
     "errors": [
       {
-        "message": "Authentication Failure: Invalid Authorization Header",
-        "locations": [
-          {
-            "line": 2,
-            "column": 3
-          }
-        ]
+        "message": "Authentication Failure: Invalid Authorization Header"
       }
     ]
   }

tests/core/snapshots/auth-basic.md_4.snap

@@ -1,6 +1,7 @@
 ---
 source: tests/core/spec.rs
 expression: response
+snapshot_kind: text
 ---
 {
   "status": 200,
@@ -11,13 +12,7 @@ expression: response
     "data": null,
     "errors": [
       {
-        "message": "Authentication Failure: Missing Authorization Header",
-        "locations": [
-          {
-            "line": 2,
-            "column": 3
-          }
-        ]
+        "message": "Authentication Failure: Missing Authorization Header"
       }
     ]
   }

tests/core/snapshots/auth-basic.md_6.snap

@@ -1,6 +1,7 @@
 ---
 source: tests/core/spec.rs
 expression: response
+snapshot_kind: text
 ---
 {
   "status": 200,
@@ -11,13 +12,7 @@ expression: response
     "data": null,
     "errors": [
       {
-        "message": "Authentication Failure: Invalid Authorization Header",
-        "locations": [
-          {
-            "line": 2,
-            "column": 3
-          }
-        ]
+        "message": "Authentication Failure: Invalid Authorization Header"
       }
     ]
   }

tests/core/snapshots/auth-jwt.md_1.snap

@@ -1,6 +1,7 @@
 ---
 source: tests/core/spec.rs
 expression: response
+snapshot_kind: text
 ---
 {
   "status": 200,
@@ -11,13 +12,7 @@ expression: response
     "data": null,
     "errors": [
       {
-        "message": "Authentication Failure: Missing Authorization Header",
-        "locations": [
-          {
-            "line": 2,
-            "column": 3
-          }
-        ]
+        "message": "Authentication Failure: Missing Authorization Header"
       }
     ]
   }

tests/core/snapshots/auth-jwt.md_3.snap

@@ -1,6 +1,7 @@
 ---
 source: tests/core/spec.rs
 expression: response
+snapshot_kind: text
 ---
 {
   "status": 200,
@@ -11,13 +12,7 @@ expression: response
     "data": null,
     "errors": [
       {
-        "message": "Authentication Failure: Missing Authorization Header",
-        "locations": [
-          {
-            "line": 2,
-            "column": 3
-          }
-        ]
+        "message": "Authentication Failure: Missing Authorization Header"
       }
     ]
   }

tests/core/snapshots/auth-jwt.md_5.snap

@@ -1,6 +1,7 @@
 ---
 source: tests/core/spec.rs
 expression: response
+snapshot_kind: text
 ---
 {
   "status": 200,
@@ -11,13 +12,7 @@ expression: response
     "data": null,
     "errors": [
       {
-        "message": "Authentication Failure: Invalid Authorization Header",
-        "locations": [
-          {
-            "line": 2,
-            "column": 3
-          }
-        ]
+        "message": "Authentication Failure: Invalid Authorization Header"
       }
     ]
   }