{-# LANGUAGE BlockArguments #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE NamedFieldPuns #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE TypeApplications #-}
module OTel.Instrumentation.Wai.Internal
  ( -- * Disclaimer
    -- $disclaimer
    buildMiddleware
  , middleware
  , spanSpecFromRequest
  , attrsFromRequest
  , includeReqLengthIfKnown
  , decodeBytes
  ) where

import Control.Monad (guard)
import Control.Monad.IO.Unlift (MonadUnliftIO(withRunInIO))
import Data.ByteString (ByteString)
import Data.Foldable (Foldable(..))
import Data.IP (fromHostAddress, fromHostAddress6)
import Data.Int (Int64)
import Data.Maybe (catMaybes)
import Data.Text (Text)
import Data.Text.Encoding (decodeUtf8With)
import Data.Text.Encoding.Error (lenientDecode)
import Data.Word (Word64)
import Network.HTTP.Types (Status(..))
import Network.Socket (SockAddr(..), PortNumber)
import Network.Wai
  ( Request
    ( httpVersion, rawPathInfo, rawQueryString, remoteHost, requestBodyLength, requestHeaderReferer
    , requestHeaderUserAgent, requestMethod
    )
  , RequestBodyLength(..), Middleware, requestHeaderHost, responseStatus
  )
import OTel.API.Common (AttrsFor(..), (.@), AttrsBuilder, Key)
import OTel.API.Trace (SpanName(..), SpanSpec, TracerProvider, TracingBackend)
import Prelude
import qualified Data.ByteString as ByteString
import qualified OTel.API.Common as OTel
import qualified OTel.API.Trace as OTel

buildMiddleware :: TracerProvider -> IO Middleware
buildMiddleware :: TracerProvider -> IO Middleware
buildMiddleware TracerProvider
tracerProvider = do
  TracingBackend
tracingBackend <- TracerProvider -> InstrumentationScope -> IO TracingBackend
forall (m :: * -> *).
MonadIO m =>
TracerProvider -> InstrumentationScope -> m TracingBackend
OTel.getTracingBackend TracerProvider
tracerProvider InstrumentationScope
"otel-instrumentation-wai"
    { OTel.instrumentationScopeVersion = Just "0.0.0" -- TODO: Automatically pull package version
    , OTel.instrumentationScopeSchemaURL = Just OTel.TRACE_SCHEMA_URL
    }
  Middleware -> IO Middleware
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Middleware -> IO Middleware) -> Middleware -> IO Middleware
forall a b. (a -> b) -> a -> b
$ TracingBackend -> Middleware
middleware TracingBackend
tracingBackend

middleware :: TracingBackend -> Middleware
middleware :: TracingBackend -> Middleware
middleware TracingBackend
tracingBackend Application
app Request
req Response -> IO ResponseReceived
sendResp = do
  (TracingT IO ResponseReceived
 -> TracingBackend -> IO ResponseReceived)
-> TracingBackend
-> TracingT IO ResponseReceived
-> IO ResponseReceived
forall a b c. (a -> b -> c) -> b -> a -> c
flip TracingT IO ResponseReceived
-> TracingBackend -> IO ResponseReceived
forall (m :: * -> *) a. TracingT m a -> TracingBackend -> m a
OTel.runTracingT TracingBackend
tracingBackend do
    -- TODO: Need to check if propagation is needed from request when propagator support is added
    SpanSpec
-> (MutableSpan -> TracingT IO ResponseReceived)
-> TracingT IO ResponseReceived
forall (m :: * -> *) a.
(MonadTracing m, HasCallStack) =>
SpanSpec -> (MutableSpan -> m a) -> m a
OTel.trace (Request -> SpanSpec
spanSpecFromRequest Request
req) \MutableSpan
mutableSpan -> do
      ((forall a. TracingT IO a -> IO a) -> IO ResponseReceived)
-> TracingT IO ResponseReceived
forall b.
((forall a. TracingT IO a -> IO a) -> IO b) -> TracingT IO b
forall (m :: * -> *) b.
MonadUnliftIO m =>
((forall a. m a -> IO a) -> IO b) -> m b
withRunInIO \forall a. TracingT IO a -> IO a
runInIO -> do
        Application
app Request
req \Response
resp -> do
          let Status { Int
statusCode :: Int
statusCode :: Status -> Int
statusCode, ByteString
statusMessage :: ByteString
statusMessage :: Status -> ByteString
statusMessage } = Response -> Status
responseStatus Response
resp
          TracingT IO () -> IO ()
forall a. TracingT IO a -> IO a
runInIO (TracingT IO () -> IO ()) -> TracingT IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ MutableSpan -> UpdateSpanSpec -> TracingT IO ()
forall (m :: * -> *).
MonadTracing m =>
MutableSpan -> UpdateSpanSpec -> m ()
OTel.updateSpan MutableSpan
mutableSpan UpdateSpanSpec
OTel.defaultUpdateSpanSpec
            { OTel.updateSpanSpecAttrs =
                Just $ OTel.HTTP_STATUS_CODE .@ statusCode
            , OTel.updateSpanSpecStatus = do
                guard $ statusCode >= 500
                pure $ OTel.SpanStatusError $ decodeBytes statusMessage
            }
          Response -> IO ResponseReceived
sendResp Response
resp

spanSpecFromRequest :: Request -> SpanSpec
spanSpecFromRequest :: Request -> SpanSpec
spanSpecFromRequest Request
req =
  SpanSpec
OTel.defaultSpanSpec
    { OTel.spanSpecName = SpanName $ decodeBytes $ rawPathInfo req
    , OTel.spanSpecKind = OTel.SpanKindServer
    , OTel.spanSpecAttrs = attrsFromRequest req
    }

attrsFromRequest :: Request -> AttrsBuilder 'AttrsForSpan
attrsFromRequest :: Request -> AttrsBuilder 'AttrsForSpan
attrsFromRequest Request
req =
  [AttrsBuilder 'AttrsForSpan] -> AttrsBuilder 'AttrsForSpan
forall m. Monoid m => [m] -> m
forall (t :: * -> *) m. (Foldable t, Monoid m) => t m -> m
fold ([AttrsBuilder 'AttrsForSpan] -> AttrsBuilder 'AttrsForSpan)
-> [AttrsBuilder 'AttrsForSpan] -> AttrsBuilder 'AttrsForSpan
forall a b. (a -> b) -> a -> b
$ [Maybe (AttrsBuilder 'AttrsForSpan)]
-> [AttrsBuilder 'AttrsForSpan]
forall a. [Maybe a] -> [a]
catMaybes
    [ AttrsBuilder 'AttrsForSpan -> Maybe (AttrsBuilder 'AttrsForSpan)
forall a. a -> Maybe a
Just (AttrsBuilder 'AttrsForSpan -> Maybe (AttrsBuilder 'AttrsForSpan))
-> AttrsBuilder 'AttrsForSpan -> Maybe (AttrsBuilder 'AttrsForSpan)
forall a b. (a -> b) -> a -> b
$ Key Text
OTel.HTTP_METHOD Key Text -> Text -> AttrsBuilder 'AttrsForSpan
forall kv from to.
(KV kv, KVConstraints kv from to) =>
Key to -> from -> kv
forall from to.
KVConstraints (AttrsBuilder 'AttrsForSpan) from to =>
Key to -> from -> AttrsBuilder 'AttrsForSpan
.@ ByteString -> Text
decodeBytes (Request -> ByteString
requestMethod Request
req)
    , AttrsBuilder 'AttrsForSpan -> Maybe (AttrsBuilder 'AttrsForSpan)
forall a. a -> Maybe a
Just (AttrsBuilder 'AttrsForSpan -> Maybe (AttrsBuilder 'AttrsForSpan))
-> AttrsBuilder 'AttrsForSpan -> Maybe (AttrsBuilder 'AttrsForSpan)
forall a b. (a -> b) -> a -> b
$ Key Text
OTel.HTTP_FLAVOR Key Text -> String -> AttrsBuilder 'AttrsForSpan
forall kv from to.
(KV kv, KVConstraints kv from to) =>
Key to -> from -> kv
forall from to.
KVConstraints (AttrsBuilder 'AttrsForSpan) from to =>
Key to -> from -> AttrsBuilder 'AttrsForSpan
.@ HttpVersion -> String
forall a. Show a => a -> String
show (Request -> HttpVersion
httpVersion Request
req)
    , case Request -> SockAddr
remoteHost Request
req of
        SockAddrUnix String
_path ->
          Maybe (AttrsBuilder 'AttrsForSpan)
forall a. Maybe a
Nothing
        SockAddrInet PortNumber
port HostAddress
hostAddress ->
          AttrsBuilder 'AttrsForSpan -> Maybe (AttrsBuilder 'AttrsForSpan)
forall a. a -> Maybe a
Just
            (AttrsBuilder 'AttrsForSpan -> Maybe (AttrsBuilder 'AttrsForSpan))
-> AttrsBuilder 'AttrsForSpan -> Maybe (AttrsBuilder 'AttrsForSpan)
forall a b. (a -> b) -> a -> b
$ Key Text
OTel.NET_HOST_IP Key Text -> String -> AttrsBuilder 'AttrsForSpan
forall kv from to.
(KV kv, KVConstraints kv from to) =>
Key to -> from -> kv
forall from to.
KVConstraints (AttrsBuilder 'AttrsForSpan) from to =>
Key to -> from -> AttrsBuilder 'AttrsForSpan
.@ IPv4 -> String
forall a. Show a => a -> String
show (HostAddress -> IPv4
fromHostAddress HostAddress
hostAddress)
                AttrsBuilder 'AttrsForSpan
-> AttrsBuilder 'AttrsForSpan -> AttrsBuilder 'AttrsForSpan
forall a. Semigroup a => a -> a -> a
<> Key Int64
OTel.NET_HOST_PORT Key Int64 -> Int64 -> AttrsBuilder 'AttrsForSpan
forall kv from to.
(KV kv, KVConstraints kv from to) =>
Key to -> from -> kv
forall from to.
KVConstraints (AttrsBuilder 'AttrsForSpan) from to =>
Key to -> from -> AttrsBuilder 'AttrsForSpan
.@ forall a b. (Integral a, Num b) => a -> b
fromIntegral @PortNumber @Int64  PortNumber
port
        SockAddrInet6 PortNumber
port HostAddress
_flowInfo HostAddress6
hostAddress6 HostAddress
_scopeID ->
          AttrsBuilder 'AttrsForSpan -> Maybe (AttrsBuilder 'AttrsForSpan)
forall a. a -> Maybe a
Just
            (AttrsBuilder 'AttrsForSpan -> Maybe (AttrsBuilder 'AttrsForSpan))
-> AttrsBuilder 'AttrsForSpan -> Maybe (AttrsBuilder 'AttrsForSpan)
forall a b. (a -> b) -> a -> b
$ Key Text
OTel.NET_HOST_IP Key Text -> String -> AttrsBuilder 'AttrsForSpan
forall kv from to.
(KV kv, KVConstraints kv from to) =>
Key to -> from -> kv
forall from to.
KVConstraints (AttrsBuilder 'AttrsForSpan) from to =>
Key to -> from -> AttrsBuilder 'AttrsForSpan
.@ IPv6 -> String
forall a. Show a => a -> String
show (HostAddress6 -> IPv6
fromHostAddress6 HostAddress6
hostAddress6)
                AttrsBuilder 'AttrsForSpan
-> AttrsBuilder 'AttrsForSpan -> AttrsBuilder 'AttrsForSpan
forall a. Semigroup a => a -> a -> a
<> Key Int64
OTel.NET_HOST_PORT Key Int64 -> Int64 -> AttrsBuilder 'AttrsForSpan
forall kv from to.
(KV kv, KVConstraints kv from to) =>
Key to -> from -> kv
forall from to.
KVConstraints (AttrsBuilder 'AttrsForSpan) from to =>
Key to -> from -> AttrsBuilder 'AttrsForSpan
.@ forall a b. (Integral a, Num b) => a -> b
fromIntegral @PortNumber @Int64  PortNumber
port
    , Request
-> Key Text
-> (Request -> Maybe ByteString)
-> Maybe (AttrsBuilder 'AttrsForSpan)
includeIfNotNull Request
req Key Text
OTel.HTTP_HOST Request -> Maybe ByteString
requestHeaderHost
    , Request
-> Key Text
-> (Request -> Maybe ByteString)
-> Maybe (AttrsBuilder 'AttrsForSpan)
includeIfNotNull Request
req Key Text
OTel.HTTP_USER_AGENT Request -> Maybe ByteString
requestHeaderUserAgent
    , Request
-> Key Text
-> (Request -> Maybe ByteString)
-> Maybe (AttrsBuilder 'AttrsForSpan)
includeIfNotNull Request
req Key Text
OTel.HTTP_TARGET \Request
r ->
        ByteString -> Maybe ByteString
forall a. a -> Maybe a
Just (ByteString -> Maybe ByteString) -> ByteString -> Maybe ByteString
forall a b. (a -> b) -> a -> b
$ Request -> ByteString
rawPathInfo Request
r ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> Request -> ByteString
rawQueryString Request
r
    , Request
-> Key Text
-> (Request -> Maybe ByteString)
-> Maybe (AttrsBuilder 'AttrsForSpan)
includeIfNotNull Request
req Key Text
OTel.NET_PEER_NAME Request -> Maybe ByteString
requestHeaderReferer
    , Request -> Maybe (AttrsBuilder 'AttrsForSpan)
includeReqLengthIfKnown Request
req
    ]

includeIfNotNull
  :: Request
  -> Key Text
  -> (Request -> Maybe ByteString)
  -> Maybe (AttrsBuilder 'AttrsForSpan)
includeIfNotNull :: Request
-> Key Text
-> (Request -> Maybe ByteString)
-> Maybe (AttrsBuilder 'AttrsForSpan)
includeIfNotNull Request
req Key Text
key Request -> Maybe ByteString
selector = do
  ByteString
val <- Request -> Maybe ByteString
selector Request
req
  Bool -> Maybe ()
forall (f :: * -> *). Alternative f => Bool -> f ()
guard (Bool -> Maybe ()) -> Bool -> Maybe ()
forall a b. (a -> b) -> a -> b
$ Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ ByteString -> Bool
ByteString.null ByteString
val
  AttrsBuilder 'AttrsForSpan -> Maybe (AttrsBuilder 'AttrsForSpan)
forall a. a -> Maybe a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (AttrsBuilder 'AttrsForSpan -> Maybe (AttrsBuilder 'AttrsForSpan))
-> AttrsBuilder 'AttrsForSpan -> Maybe (AttrsBuilder 'AttrsForSpan)
forall a b. (a -> b) -> a -> b
$ Key Text
key Key Text -> Text -> AttrsBuilder 'AttrsForSpan
forall kv from to.
(KV kv, KVConstraints kv from to) =>
Key to -> from -> kv
forall from to.
KVConstraints (AttrsBuilder 'AttrsForSpan) from to =>
Key to -> from -> AttrsBuilder 'AttrsForSpan
.@ ByteString -> Text
decodeBytes ByteString
val

includeReqLengthIfKnown :: Request -> Maybe (AttrsBuilder 'AttrsForSpan)
includeReqLengthIfKnown :: Request -> Maybe (AttrsBuilder 'AttrsForSpan)
includeReqLengthIfKnown Request
req = do
  KnownLength Word64
len <- RequestBodyLength -> Maybe RequestBodyLength
forall a. a -> Maybe a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (RequestBodyLength -> Maybe RequestBodyLength)
-> RequestBodyLength -> Maybe RequestBodyLength
forall a b. (a -> b) -> a -> b
$ Request -> RequestBodyLength
requestBodyLength Request
req
  Bool -> Maybe ()
forall (f :: * -> *). Alternative f => Bool -> f ()
guard (Bool -> Maybe ()) -> Bool -> Maybe ()
forall a b. (a -> b) -> a -> b
$ Word64
len Word64 -> Word64 -> Bool
forall a. Ord a => a -> a -> Bool
> Word64
0 Bool -> Bool -> Bool
&& Word64
len Word64 -> Word64 -> Bool
forall a. Ord a => a -> a -> Bool
<= Int64 -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral (forall a. Bounded a => a
maxBound @Int64)
  AttrsBuilder 'AttrsForSpan -> Maybe (AttrsBuilder 'AttrsForSpan)
forall a. a -> Maybe a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (AttrsBuilder 'AttrsForSpan -> Maybe (AttrsBuilder 'AttrsForSpan))
-> AttrsBuilder 'AttrsForSpan -> Maybe (AttrsBuilder 'AttrsForSpan)
forall a b. (a -> b) -> a -> b
$ Key Int64
OTel.HTTP_REQUEST_CONTENT_LENGTH Key Int64 -> Int64 -> AttrsBuilder 'AttrsForSpan
forall kv from to.
(KV kv, KVConstraints kv from to) =>
Key to -> from -> kv
forall from to.
KVConstraints (AttrsBuilder 'AttrsForSpan) from to =>
Key to -> from -> AttrsBuilder 'AttrsForSpan
.@ forall a b. (Integral a, Num b) => a -> b
fromIntegral @Word64 @Int64 Word64
len

decodeBytes :: ByteString -> Text
decodeBytes :: ByteString -> Text
decodeBytes = OnDecodeError -> ByteString -> Text
decodeUtf8With OnDecodeError
lenientDecode

-- $disclaimer
--
-- In general, changes to this module will not be reflected in the library's
-- version updates. Direct use of this module should be done with utmost care,
-- otherwise invariants will easily be violated.