From ff79a5227bb395cf52fe44f3d3533c16c5b95de0 Mon Sep 17 00:00:00 2001
From: Paolo Capriotti <paolo@capriotti.io>
Date: Fri, 3 May 2024 11:13:32 +0200
Subject: [PATCH 1/2] Generalise catchErrors middleware

---
 .../src/Network/Wai/Utilities/Server.hs       | 38 ++++++++++++-------
 services/federator/src/Federator/Response.hs  |  5 ++-
 2 files changed, 29 insertions(+), 14 deletions(-)

diff --git a/libs/wai-utilities/src/Network/Wai/Utilities/Server.hs b/libs/wai-utilities/src/Network/Wai/Utilities/Server.hs
index 80ee4329c13..14fa566c12a 100644
--- a/libs/wai-utilities/src/Network/Wai/Utilities/Server.hs
+++ b/libs/wai-utilities/src/Network/Wai/Utilities/Server.hs
@@ -31,6 +31,7 @@ module Network.Wai.Utilities.Server
 
     -- * Middlewares
     catchErrors,
+    catchErrorsWithRequestId,
     OnErrorMetrics,
     heavyDebugLogging,
     rethrow5xx,
@@ -184,6 +185,9 @@ route rt rq k = Route.routeWith (Route.Config $ errorRs' noEndpoint) rt rq (lift
 --------------------------------------------------------------------------------
 -- Middlewares
 
+catchErrors :: Logger -> OnErrorMetrics -> Middleware
+catchErrors l m = catchErrorsWithRequestId lookupRequestId l m
+
 -- | Create a middleware that catches exceptions and turns
 -- them into appropriate 'Error' responses, thereby logging
 -- as well as counting server errors (i.e. exceptions that
@@ -191,14 +195,21 @@ route rt rq k = Route.routeWith (Route.Config $ errorRs' noEndpoint) rt rq (lift
 --
 -- This does not log any 'Response' values with error status.
 -- See 'catchErrors'.
-catchErrors :: Logger -> OnErrorMetrics -> Middleware
-catchErrors l m app req k =
-  rethrow5xx l app req k `catch` errorResponse
+catchErrorsWithRequestId ::
+  (Request -> Maybe ByteString) ->
+  Logger ->
+  OnErrorMetrics ->
+  Middleware
+catchErrorsWithRequestId getRequestId l m app req k =
+  rethrow5xx getRequestId l app req k `catch` errorResponse
   where
+    mReqId = getRequestId req
+
     errorResponse :: SomeException -> IO ResponseReceived
     errorResponse ex = do
       er <- runHandlers ex errorHandlers
-      onError l m req k er
+      onError l mReqId m req k er
+
 {-# INLINEABLE catchErrors #-}
 
 -- | Standard handlers for turning exceptions into appropriate
@@ -298,15 +309,15 @@ emitLByteString lbs = do
 
 -- | Run the 'Application'; check the response status; if >=500, throw a 'Wai.Error' with
 -- label @"server-error"@ and the body as the error message.
-rethrow5xx :: Logger -> Middleware
-rethrow5xx logger app req k = app req k'
+rethrow5xx :: (Request -> Maybe ByteString) -> Logger -> Middleware
+rethrow5xx getRequestId logger app req k = app req k'
   where
     k' resp@WaiInt.ResponseRaw {} = do
       -- See Note [Raw Response]
       let logMsg =
             field "canoncalpath" (show $ pathInfo req)
               . field "rawpath" (rawPathInfo req)
-              . field "request" (fromMaybe "N/A" $ lookupRequestId req)
+              . field "request" (fromMaybe "N/A" $ getRequestId req)
               . msg (val "ResponseRaw - cannot collect metrics or log info on errors")
       Log.log logger Log.Debug logMsg
       k resp
@@ -349,15 +360,16 @@ type OnErrorMetrics = [Either Prm.Counter Metrics]
 onError ::
   MonadIO m =>
   Logger ->
+  Maybe ByteString ->
   OnErrorMetrics ->
   Request ->
   Continue IO ->
   Either Wai.Error JSONResponse ->
   m ResponseReceived
-onError g m r k e = liftIO $ do
+onError g mReqId m r k e = liftIO $ do
   case e of
-    Left we -> logError g (Just r) we
-    Right jr -> logJSONResponse g (Just r) jr
+    Left we -> logError' g mReqId we
+    Right jr -> logJSONResponse g mReqId jr
   let resp = either waiErrorToJSONResponse id e
   let code = statusCode (resp.status)
   when (code >= 500) $
@@ -379,9 +391,9 @@ logError' g mr e = liftIO $ doLog g (logErrorMsgWithRequest mr e)
       | statusCode (Error.code e) >= 500 = Log.err
       | otherwise = Log.debug
 
-logJSONResponse :: (MonadIO m, HasRequest r) => Logger -> Maybe r -> JSONResponse -> m ()
-logJSONResponse g mr e = do
-  let r = fromMaybe "N/A" (mr >>= lookupRequestId)
+logJSONResponse :: MonadIO m => Logger -> Maybe ByteString -> JSONResponse -> m ()
+logJSONResponse g mReqId e = do
+  let r = fromMaybe "N/A" mReqId
   liftIO $
     doLog g $
       field "request" r
diff --git a/services/federator/src/Federator/Response.hs b/services/federator/src/Federator/Response.hs
index 6f70df5a390..f4082f93c1a 100644
--- a/services/federator/src/Federator/Response.hs
+++ b/services/federator/src/Federator/Response.hs
@@ -119,7 +119,7 @@ serveServant ::
   IO ()
 serveServant middleware server env port =
   Warp.run port
-    . Wai.catchErrors (view applog env) []
+    . Wai.catchErrorsWithRequestId getRequestId (view applog env) []
     . middleware
     $ app
   where
@@ -127,6 +127,9 @@ serveServant middleware server env port =
     app =
       genericServe server
 
+    getRequestId :: Wai.Request -> Maybe ByteString
+    getRequestId = lookup "Wire-Origin-Request-Id" . Wai.requestHeaders
+
 type AllEffects =
   '[ Metrics,
      Remote,

From f20c5b3f8d5fb6112e13944b510a84b18848c5bc Mon Sep 17 00:00:00 2001
From: Paolo Capriotti <paolo@capriotti.io>
Date: Fri, 3 May 2024 11:30:28 +0200
Subject: [PATCH 2/2] Add CHANGELOG

---
 changelog.d/5-internal/federator-request-id | 1 +
 1 file changed, 1 insertion(+)
 create mode 100644 changelog.d/5-internal/federator-request-id

diff --git a/changelog.d/5-internal/federator-request-id b/changelog.d/5-internal/federator-request-id
new file mode 100644
index 00000000000..4f8c042bfa8
--- /dev/null
+++ b/changelog.d/5-internal/federator-request-id
@@ -0,0 +1 @@
+Log federator request ID on exceptions