Fix tracing issues with authentication (hasura#70)

* Enable tracing for the whole request

* refactor webhook Error IntoResponse

* Add additional trace info

* add a comment

* lint

GitOrigin-RevId: 26f3932d753e3e510f2fb963dccc13390a518401

v3/dev.docker-compose.yaml

@@ -8,6 +8,24 @@ services:
     ports:
       - "3050:3050"
     build: ./hasura-authn-webhook/dev-auth-webhook
+  jaeger:
+    image: jaegertracing/all-in-one:1.37
+    restart: always
+    ports:
+      - 5775:5775/udp
+      - 6831:6831/udp
+      - 6832:6832/udp
+      - 5778:5778
+      - 4002:16686
+      - 14250:14250
+      - 14268:14268
+      - 14269:14269
+      - 4317:4317 # OTLP gRPC
+      - 4318:4318 # OTLP HTTP
+      - 9411:9411
+    environment:
+      COLLECTOR_OTLP_ENABLED: 'true'
+      COLLECTOR_ZIPKIN_HOST_PORT: '9411'
   source:
     build:
       context: .

v3/engine/bin/engine/main.rs

@@ -94,6 +94,12 @@ async fn main() {
                     state.clone(),
                     authentication_middleware,
                 ))
+                .layer(axum::middleware::from_fn(
+                    graphql_request_tracing_middleware,
+                ))
+                // *PLEASE DO NOT ADD ANY MIDDLEWARE
+                // BEFORE THE `graphql_request_tracing_middleware`*
+                // Refer to it for more details.
                 .layer(TraceLayer::new_for_http())
                 .with_state(state.clone());
 
@@ -117,6 +123,23 @@ async fn main() {
     global::shutdown_tracer_provider();
 }
 
+/// Middleware to start tracing of the `/graphql` request.
+/// This middleware must be active for the entire duration
+/// of the request i.e. this middleware should be the
+/// entry point and the exit point of the GraphQL request.
+async fn graphql_request_tracing_middleware<B>(
+    request: Request<B>,
+    next: Next<B>,
+) -> axum::response::Result<axum::response::Response> {
+    let tracer = global::tracer("engine");
+
+    Ok(tracer
+        .in_span("/graphql", |ctx| {
+            async { next.run(request).await }.with_context(ctx)
+        })
+        .await)
+}
+
 /// This middleware authenticates the incoming GraphQL request according to the
 /// authentication configuration present in the `auth_config` of `EngineState`. The
 /// result of the authentication is `hasura-authn-core::Identity`, which is then

v3/engine/src/execute.rs

@@ -53,7 +53,8 @@ pub async fn execute_query_internal(
                     schema,
                     session,
                     &normalized_request.selection_set,
-                )?;
+                )
+                .with_traced_errors()?;
                 tracer
                     .in_span("execute", |ctx| {
                         async {

v3/hasura-authn-webhook/Cargo.toml

@@ -18,6 +18,8 @@ open-dds = { path = "../open-dds" }
 hasura-authn-core = { path = "../hasura-authn-core" }
 lang-graphql = { path = "../lang-graphql" }
 schemars = { version = "0.8.12", features = ["smol_str", "url"] }
+opentelemetry = { version = "0.20.0" }
+tracing-util = { path = "../tracing-util" }
 
 [dev-dependencies]
 pretty_assertions = "1.3.0"

v3/hasura-authn-webhook/src/webhook.rs

@@ -10,82 +10,54 @@ use axum::{
 };
 use lang_graphql::http::Response;
 use lazy_static::lazy_static;
+use opentelemetry::{
+    global,
+    trace::{FutureExt, Tracer},
+};
 use reqwest::{header::ToStrError, Url};
 use serde::{de::Error as SerdeDeError, Deserialize, Deserializer, Serialize, Serializer};
 
 use hasura_authn_core as auth_base;
 use open_dds::data_specification::session_variables;
 use schemars::JsonSchema;
+use thiserror::Error;
+use tracing_util::{FutureTracing, Tracing};
 
-#[derive(Debug)]
+#[derive(Error, Debug)]
 pub enum Error {
+    #[error("Error in converting the header value corresponding to the {header_name} to a String - {error}")]
     ErrorInConvertingHeaderValueToString {
         header_name: HeaderName,
         error: ToStrError,
     },
+    #[error("internal: Error while making the authentication HTTP request to the webhook - {0}")]
     ErrorWhileMakingHTTPRequestToTheAuthHook(reqwest::Error),
+    #[error("The Authentication hook has denied to execute the request.")]
     AuthenticationFailed,
+    #[error("The authentication hook has returned the status {0}. Only 200 and 401 response status are recognized.")]
     AuthHookUnexpectedStatus(reqwest::StatusCode),
+    #[error("internal: Error while parsing the response from the webhook - {0}")]
     ErrorInParsingResponse(serde_json::Error),
+    #[error("internal reqwest error: {0}")]
     InternalReqwestError(reqwest::Error),
+    #[error("'x-hasura-role' session variable not found in the webhook response.")]
     RoleSessionVariableNotFound,
 }
 
 impl IntoResponse for Error {
     fn into_response(self) -> axum::response::Response {
-        match self {
-            Error::ErrorInConvertingHeaderValueToString {
-                header_name,
-                error,
-            } => Response::error_message_with_status(
-                StatusCode::BAD_REQUEST,
-                format!(
-                    "Error in converting the header value corresponding to the {header_name} to a String - {error}"
-                ),
-            ),
-            Error::ErrorWhileMakingHTTPRequestToTheAuthHook(e) => {
-                Response::error_message_with_status(
-                    StatusCode::INTERNAL_SERVER_ERROR,
-                    format!(
-                        "internal: Error while making the authentication HTTP request to the webhook - {e}"
-                    ),
-                )
-            },
-            Error::ErrorInParsingResponse(e) => {
-                Response::error_message_with_status(
-                    StatusCode::INTERNAL_SERVER_ERROR,
-                    format!(
-                        "internal: Error while parsing the response from the webhook - {e}"
-                    ),
-                )
-            },
-            Error::AuthenticationFailed => {
-                Response::error_message_with_status(
-                    StatusCode::FORBIDDEN,
-                    "The Authentication hook has denied to execute the request.".to_string()
-                )
-            },
-            Error::AuthHookUnexpectedStatus(s) => {
-                Response::error_message_with_status(
-                    StatusCode::INTERNAL_SERVER_ERROR,
-                    format!("The authentication hook has returned the status {}. Only 200 and 401 response status are recognized.", s)
-                )
-
-            }
-            Error::InternalReqwestError(e) => {
-                Response::error_message_with_status(
-                    StatusCode::INTERNAL_SERVER_ERROR,
-                    format!("internal reqwest error: {e}")
-                )
-
-            }
-            Error::RoleSessionVariableNotFound => {
-                Response::error_message_with_status(
-                    StatusCode::INTERNAL_SERVER_ERROR,
-                    "'x-hasura-role' session variable not found in the webhook response.".into()
-                )
+        let response_status = match &self {
+            Error::ErrorInConvertingHeaderValueToString { .. } => StatusCode::BAD_REQUEST,
+            Error::ErrorWhileMakingHTTPRequestToTheAuthHook(_e) => {
+                StatusCode::INTERNAL_SERVER_ERROR
             }
-        }.into_response()
+            Error::ErrorInParsingResponse(_e) => StatusCode::INTERNAL_SERVER_ERROR,
+            Error::AuthenticationFailed => StatusCode::FORBIDDEN,
+            Error::AuthHookUnexpectedStatus(_s) => StatusCode::INTERNAL_SERVER_ERROR,
+            Error::InternalReqwestError(_e) => StatusCode::INTERNAL_SERVER_ERROR,
+            Error::RoleSessionVariableNotFound => StatusCode::INTERNAL_SERVER_ERROR,
+        };
+        Response::error_message_with_status(response_status, self.to_string()).into_response()
     }
 }
 
@@ -159,6 +131,7 @@ async fn make_auth_hook_request(
     client_headers: &HeaderMap,
     allow_role_emulation_for: Option<Role>,
 ) -> Result<auth_base::Identity, Error> {
+    let tracer = global::tracer("engine");
     let http_request_builder = match auth_hook_config.mode {
         AuthHookMode::GET => {
             let mut auth_hook_headers = HeaderMap::new();
@@ -198,10 +171,18 @@ async fn make_auth_hook_request(
         .build()
         .map_err(Error::InternalReqwestError)?;
 
-    let response = http_client
-        .execute(req)
-        .await
-        .map_err(Error::ErrorWhileMakingHTTPRequestToTheAuthHook)?;
+    let response = tracer
+        .in_span("request_to_webhook", |ctx| {
+            async {
+                http_client
+                    .execute(req)
+                    .await
+                    .map_err(Error::ErrorWhileMakingHTTPRequestToTheAuthHook)
+            }
+            .with_traced_errors()
+            .with_context(ctx)
+        })
+        .await?;
 
     match response.status() {
         reqwest::StatusCode::UNAUTHORIZED => Err(Error::AuthenticationFailed),
@@ -266,14 +247,23 @@ pub async fn authenticate_request(
     client_headers: &HeaderMap,
     allow_role_emulation_for: Option<Role>,
 ) -> axum::response::Result<auth_base::Identity> {
-    Ok(make_auth_hook_request(
-        http_client,
-        auth_hook_config,
-        client_headers,
-        allow_role_emulation_for,
-    )
-    .await
-    .map_err(|e| e.into_response())?)
+    let tracer = global::tracer("engine");
+    tracer
+        .in_span("webhook_authenticate_request", |ctx| {
+            async {
+                Ok(make_auth_hook_request(
+                    http_client,
+                    auth_hook_config,
+                    client_headers,
+                    allow_role_emulation_for,
+                )
+                .await
+                .with_traced_errors()
+                .map_err(|e| e.into_response())?)
+            }
+            .with_context(ctx)
+        })
+        .await
 }
 
 #[cfg(test)]