the-guild-org · dotansimha · Feb 12, 2024 · Feb 12, 2024 · Feb 12, 2024
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/libs/common/src/lib.rs b/libs/common/src/lib.rs
@@ -4,5 +4,6 @@ pub mod http;
 pub mod json;
 pub mod plugin;
 pub mod serde_utils;
+pub mod vrl_functions;
 pub mod vrl_utils;
 pub use graphql_parser::query::{Definition, Document, OperationDefinition, ParseError};
diff --git a/plugins/vrl/src/vrl_functions.rs → libs/common/src/vrl_functions.rs b/plugins/vrl/src/vrl_functions.rs → libs/common/src/vrl_functions.rs
@@ -1,4 +1,4 @@
-use conductor_common::http::Bytes;
+use bytes::Bytes;
 use vrl::{
   compiler::{
     state, value::kind, Context, Expression, Function, FunctionExpression, Parameter, Resolved,

diff --git a/libs/common/src/vrl_utils.rs b/libs/common/src/vrl_utils.rs
@@ -1,10 +1,28 @@
+use crate::{
+  graphql::GraphQLRequest,
+  http::{ConductorHttpRequest, ConductorHttpResponse},
+  serde_utils::LocalFileReference,
+  vrl_functions::vrl_fns,
+};
 use anyhow::{anyhow, Result};
 use schemars::JsonSchema;
 use serde::{Deserialize, Serialize};
+use std::collections::BTreeMap;
 use std::str::FromStr;
-use vrl::{prelude::NotNan, value::Value};
-
-use crate::serde_utils::LocalFileReference;
+use vrl::{
+  compiler::{state::RuntimeState, TargetValue},
+  value,
+};
+use vrl::{
+  compiler::{Context, TimeZone},
+  value::KeyString,
+};
+use vrl::{
+  compiler::{Program, Resolved},
+  diagnostic::DiagnosticList,
+  prelude::NotNan,
+  value::Value,
+};
 
 #[derive(Serialize, Deserialize, Debug, Clone, JsonSchema)]
 #[serde(tag = "from")]
@@ -19,13 +37,61 @@ pub enum VrlConfigReference {
   File { path: LocalFileReference },
 }
 
+#[derive(Debug)]
+pub struct VrlProgramProxy {
+  program: Program,
+}
+
+impl VrlProgramProxy {
+  pub fn new(program: Program) -> Self {
+    Self { program }
+  }
+
+  pub fn resolve_with_state(
+    &self,
+    write_value: Value,
+    read_value: Value,
+    state: &mut RuntimeState,
+  ) -> Resolved {
+    let mut target = TargetValue {
+      metadata: read_value,
+      value: write_value,
+      secrets: Default::default(),
+    };
+    self
+      .program
+      .resolve(&mut Context::new(&mut target, state, &TimeZone::default()))
+  }
+
+  pub fn resolve(&self, write_value: Value, read_value: Value) -> Resolved {
+    self.resolve_with_state(write_value, read_value, &mut RuntimeState::default())
+  }
+}
+
 impl VrlConfigReference {
   pub fn contents(&self) -> &String {
     match self {
       VrlConfigReference::Inline { content } => content,
       VrlConfigReference::File { path } => &path.contents,
     }
   }
+
+  pub fn program(&self) -> Result<VrlProgramProxy, DiagnosticList> {
+    let contents = self.contents();
+
+    match vrl::compiler::compile(contents, &vrl_fns()) {
+      Err(err) => Err(err),
+      Ok(result) => {
+        if result.warnings.len() > 0 {
+          tracing::warn!("vrl compiler warning: {:?}", result.warnings);
+        }
+
+        let r = VrlProgramProxy::new(result.program);
+
+        Ok(r)
+      }
+    }
+  }
 }
 
 pub fn vrl_value_to_serde_value(value: &Value) -> Result<serde_json::Value> {
@@ -87,3 +153,69 @@ pub fn serde_value_to_vrl_value(value: &serde_json::Value) -> Result<Value> {
     ),
   })
 }
+
+pub fn conductor_response_to_value(res: &ConductorHttpResponse) -> Value {
+  let body = res.body.clone();
+  let status = res.status.as_u16();
+  let mut headers_map: BTreeMap<KeyString, Value> = BTreeMap::new();
+
+  for (k, v) in res.headers.iter() {
+    headers_map.insert(k.to_string().into(), v.as_bytes().into());
+  }
+
+  let headers = Value::Object(headers_map);
+
+  value!({
+      body: body,
+      status: status,
+      headers: headers,
+  })
+}
+
+pub fn conductor_graphql_request_to_value(gql_req: &GraphQLRequest) -> Result<Value> {
+  let operation = gql_req.operation.as_bytes();
+  let operation_name = gql_req.operation_name.as_ref().map(|v| v.as_bytes());
+  let variables = match gql_req.variables.as_ref() {
+    Some(v) => Some(serde_value_to_vrl_value(&serde_json::Value::Object(
+      v.clone(),
+    ))?),
+    None => None,
+  };
+
+  let extensions = match gql_req.extensions.as_ref() {
+    Some(v) => Some(serde_value_to_vrl_value(&serde_json::Value::Object(
+      v.clone(),
+    ))?),
+    None => None,
+  };
+
+  Ok(value!({
+      operation: operation,
+      operation_name: operation_name,
+      variables: variables,
+      extensions: extensions,
+  }))
+}
+
+pub fn conductor_request_to_value(req: &ConductorHttpRequest) -> Value {
+  let body = req.body.clone();
+  let uri = req.uri.as_bytes();
+  let query_string = req.query_string.as_bytes();
+  let method_str = req.method.to_string();
+  let method = method_str.as_bytes();
+  let mut headers_map: BTreeMap<KeyString, Value> = BTreeMap::new();
+
+  for (k, v) in req.headers.iter() {
+    headers_map.insert(k.to_string().into(), v.as_bytes().into());
+  }
+
+  let headers = Value::Object(headers_map);
+
+  value!({
+      body: body,
+      uri: uri,
+      query_string: query_string,
+      method: method,
+      headers: headers,
+  })
+}
diff --git a/plugins/disable_introspection/Cargo.toml b/plugins/disable_introspection/Cargo.toml
@@ -8,6 +8,7 @@ path = "src/lib.rs"
 
 [dependencies]
 tracing = { workspace = true }
+anyhow = { workspace = true }
 serde = { workspace = true }
 serde_json = { workspace = true }
 async-trait = { workspace = true }

diff --git a/plugins/disable_introspection/src/plugin.rs b/plugins/disable_introspection/src/plugin.rs
@@ -3,52 +3,36 @@ use conductor_common::{
   graphql::GraphQLResponse,
   http::StatusCode,
   plugin::{CreatablePlugin, Plugin, PluginError},
+  vrl_utils::{conductor_request_to_value, VrlProgramProxy},
 };
-use tracing::{error, warn};
-use vrl::{
-  compiler::{Context, Program, TargetValue, TimeZone},
-  value,
-  value::Secrets,
-};
+use tracing::error;
+use vrl::value;
 
 use conductor_common::execute::RequestExecutionContext;
 
-use vrl_plugin::{utils::conductor_request_to_value, vrl_functions::vrl_fns};
-
 #[derive(Debug)]
 pub struct DisableIntrospectionPlugin {
-  condition: Option<Program>,
+  condition: Option<VrlProgramProxy>,
 }
 
 #[async_trait::async_trait(?Send)]
 impl CreatablePlugin for DisableIntrospectionPlugin {
   type Config = DisableIntrospectionPluginConfig;
 
   async fn create(config: Self::Config) -> Result<Box<Self>, PluginError> {
-    let instance = match &config.condition {
-      Some(condition) => match vrl::compiler::compile(condition.contents(), &vrl_fns()) {
-        Err(err) => {
-          error!("vrl compiler error: {:?}", err);
-          // @expected: we need to exit the process if our provided VRL condition has incorrect syntax.
-          panic!("failed to compile vrl program for disable_introspection plugin");
-        }
-        Ok(result) => {
-          if result.warnings.len() > 0 {
-            warn!(
-              "vrl compiler warning for disable_introspection plugin: {:?}",
-              result.warnings
-            );
-          }
-
-          Self {
-            condition: Some(result.program),
-          }
+    let condition = match &config.condition {
+      Some(condition) => match condition.program() {
+        Ok(program) => Some(program),
+        Err(e) => {
+          return Err(PluginError::InitError {
+            source: anyhow::anyhow!("vrl compiler error: {:?}", e),
+          })
         }
       },
-      None => Self { condition: None },
+      None => None,
     };
 
-    Ok(Box::new(instance))
+    Ok(Box::new(Self { condition }))
   }
 }
 
@@ -60,19 +44,14 @@ impl Plugin for DisableIntrospectionPlugin {
         let should_disable = match &self.condition {
           Some(program) => {
             let downstream_http_req = conductor_request_to_value(&ctx.downstream_http_request);
-            let mut target = TargetValue {
-              value: value!({}),
-              metadata: value!({
+
+            match program.resolve_with_state(
+              value::Value::Null,
+              value!({
                 downstream_http_req: downstream_http_req,
               }),
-              secrets: Secrets::default(),
-            };
-
-            match program.resolve(&mut Context::new(
-              &mut target,
               ctx.vrl_shared_state(),
-              &TimeZone::default(),
-            )) {
+            ) {
               Ok(ret) => match ret {
                 vrl::value::Value::Boolean(b) => b,
                 _ => {

diff --git a/plugins/vrl/src/downstream_graphql_request.rs b/plugins/vrl/src/downstream_graphql_request.rs
@@ -1,7 +1,8 @@
 use conductor_common::{
   graphql::{GraphQLRequest, GraphQLResponse, ParsedGraphQLRequest},
   http::StatusCode,
-  vrl_utils::vrl_value_to_serde_value,
+  vrl_functions::ShortCircuitFn,
+  vrl_utils::{conductor_graphql_request_to_value, vrl_value_to_serde_value},
 };
 use tracing::error;
 use vrl::{
@@ -12,8 +13,6 @@ use vrl::{
 
 use conductor_common::execute::RequestExecutionContext;
 
-use super::{utils::conductor_graphql_request_to_value, vrl_functions::ShortCircuitFn};
-
 static METADATA_GRAPHQL_OPERATION_INFO: &str = "downstream_graphql_req";
 static TARGET_GRAPHQL_OPERATION_KEY: &str = "graphql.operation";
 static TARGET_GRAPHQL_OPERATION_NAME: &str = "graphql.operation_name";

diff --git a/plugins/vrl/src/downstream_http_request.rs b/plugins/vrl/src/downstream_http_request.rs
@@ -1,7 +1,8 @@
 use conductor_common::{
   graphql::{GraphQLRequest, GraphQLResponse, ParsedGraphQLRequest},
   http::StatusCode,
-  vrl_utils::vrl_value_to_serde_value,
+  vrl_functions::ShortCircuitFn,
+  vrl_utils::{conductor_request_to_value, vrl_value_to_serde_value},
 };
 use tracing::error;
 use vrl::{
@@ -12,8 +13,6 @@ use vrl::{
 
 use conductor_common::execute::RequestExecutionContext;
 
-use super::{utils::conductor_request_to_value, vrl_functions::ShortCircuitFn};
-
 static TARGET_GRAPHQL_OPERATION_KEY: &str = "graphql.operation";
 static TARGET_GRAPHQL_OPERATION_NAME: &str = "graphql.operation_name";
 static TARGET_GRAPHQL_OPERATION_VARIABLES: &str = "graphql.variables";

diff --git a/plugins/vrl/src/downstream_http_response.rs b/plugins/vrl/src/downstream_http_response.rs
@@ -3,6 +3,8 @@ use std::str::FromStr;
 use conductor_common::{
   graphql::GraphQLResponse,
   http::{ConductorHttpResponse, HeaderName, HeaderValue, StatusCode},
+  vrl_functions::ShortCircuitFn,
+  vrl_utils::conductor_response_to_value,
 };
 use tracing::error;
 use vrl::{
@@ -13,8 +15,6 @@ use vrl::{
 
 use conductor_common::execute::RequestExecutionContext;
 
-use super::{utils::conductor_response_to_value, vrl_functions::ShortCircuitFn};
-
 static METADATA_DOWNSTREAM_HTTP_RES: &str = "downstream_http_res";
 static TARGET_DOWNSTREAM_HTTP_RES_VALUE_HEADERS: &str = "downstream_http_res.headers";
 static TARGET_DOWNSTREAM_HTTP_RES_VALUE_STATUS: &str = "downstream_http_res.status";

diff --git a/plugins/vrl/src/lib.rs b/plugins/vrl/src/lib.rs
@@ -4,8 +4,6 @@ mod downstream_http_request;
 mod downstream_http_response;
 mod plugin;
 mod upstream_http_request;
-pub mod utils;
-pub mod vrl_functions;
 
 pub use config::VrlPluginConfig as Config;
 pub use plugin::VrlPlugin as Plugin;
diff --git a/plugins/vrl/src/plugin.rs b/plugins/vrl/src/plugin.rs
@@ -1,5 +1,6 @@
 use conductor_common::http::{ConductorHttpRequest, ConductorHttpResponse};
 use conductor_common::plugin::{CreatablePlugin, Plugin, PluginError};
+use conductor_common::vrl_functions::vrl_fns;
 use tracing::{error, warn};
 use vrl::compiler::{Function, Program, TypeState};
 
@@ -11,7 +12,6 @@ use super::downstream_graphql_request::vrl_downstream_graphql_request;
 use super::downstream_http_request::vrl_downstream_http_request;
 use super::downstream_http_response::vrl_downstream_http_response;
 use super::upstream_http_request::vrl_upstream_http_request;
-use super::vrl_functions::vrl_fns;
 
 #[derive(Debug)]
 pub struct VrlPlugin {

diff --git a/plugins/vrl/src/upstream_http_request.rs b/plugins/vrl/src/upstream_http_request.rs
@@ -3,6 +3,8 @@ use std::str::FromStr;
 use conductor_common::{
   graphql::GraphQLResponse,
   http::{ConductorHttpRequest, HeaderName, HeaderValue, Method, StatusCode},
+  vrl_functions::ShortCircuitFn,
+  vrl_utils::conductor_request_to_value,
 };
 use tracing::error;
 use vrl::{
@@ -13,8 +15,6 @@ use vrl::{
 
 use conductor_common::execute::RequestExecutionContext;
 
-use super::{utils::conductor_request_to_value, vrl_functions::ShortCircuitFn};
-
 static METADATA_UPSTREAM_HTTP_REQ: &str = "upstream_http_req";
 static TARGET_UPSTREAM_HTTP_REQ_VALUE_HEADERS: &str = "upstream_http_req.headers";
 static TARGET_UPSTREAM_HTTP_REQ_VALUE_METHOD: &str = "upstream_http_req.method";