{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE NamedFieldPuns #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE TupleSections #-}

{- | B3 Propagation Requirements:
 https://github.com/openzipkin/b3-propagation
 https://github.com/open-telemetry/opentelemetry-specification/blob/main/specification/context/api-propagators.md#b3-requirements
-}
module OpenTelemetry.Propagator.B3 (
  b3TraceContextPropagator,
  b3MultiTraceContextPropagator,
) where

--------------------------------------------------------------------------------

import Control.Applicative ((<|>))
import Data.ByteString (ByteString)
import Data.List (intersperse)
import Data.Maybe (catMaybes, fromMaybe)
import qualified Data.Text.Encoding as Text
import Network.HTTP.Types (HeaderName, RequestHeaders, ResponseHeaders)
import OpenTelemetry.Common (TraceFlags (..))
import OpenTelemetry.Context (Context)
import qualified OpenTelemetry.Context as Context
import OpenTelemetry.Propagator (Propagator (..))
import OpenTelemetry.Propagator.B3.Internal
import qualified OpenTelemetry.Trace.Core as Core
import qualified OpenTelemetry.Trace.TraceState as TS
import Prelude


--------------------------------------------------------------------------------

b3TraceContextPropagator :: Propagator Context RequestHeaders ResponseHeaders
b3TraceContextPropagator :: Propagator Context RequestHeaders RequestHeaders
b3TraceContextPropagator =
  Propagator
    { propagatorNames :: [Text]
propagatorNames = [Text
"B3 Trace Context"]
    , extractor :: RequestHeaders -> Context -> IO Context
extractor = \RequestHeaders
hs Context
c ->
        case RequestHeaders -> Maybe SpanContext
b3Extractor RequestHeaders
hs of
          Maybe SpanContext
Nothing -> Context -> IO Context
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Context
c
          Just SpanContext
spanContext' -> Context -> IO Context
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Context -> IO Context) -> Context -> IO Context
forall a b. (a -> b) -> a -> b
$ Span -> Context -> Context
Context.insertSpan (SpanContext -> Span
Core.wrapSpanContext SpanContext
spanContext') Context
c
    , injector :: Context -> RequestHeaders -> IO RequestHeaders
injector = \Context
c RequestHeaders
hs ->
        case Context -> Maybe Span
Context.lookupSpan Context
c of
          Maybe Span
Nothing -> RequestHeaders -> IO RequestHeaders
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure RequestHeaders
hs
          Just Span
span' -> do
            Core.SpanContext {TraceId
traceId :: TraceId
traceId :: SpanContext -> TraceId
traceId, SpanId
spanId :: SpanId
spanId :: SpanContext -> SpanId
spanId, traceState :: SpanContext -> TraceState
traceState = TS.TraceState [(Key, Value)]
traceState} <- Span -> IO SpanContext
forall (m :: * -> *). MonadIO m => Span -> m SpanContext
Core.getSpanContext Span
span'
            let traceIdValue :: ByteString
traceIdValue = TraceId -> ByteString
encodeTraceId TraceId
traceId
                spanIdValue :: ByteString
spanIdValue = SpanId -> ByteString
encodeSpanId SpanId
spanId
                samplingStateValue :: Maybe Text
samplingStateValue = Key -> [(Key, Value)] -> Maybe Value
forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup (Text -> Key
TS.Key Text
"sampling-state") [(Key, Value)]
traceState Maybe Value
-> (Value -> Maybe SamplingState) -> Maybe SamplingState
forall a b. Maybe a -> (a -> Maybe b) -> Maybe b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Value -> Maybe SamplingState
samplingStateFromValue Maybe SamplingState -> (SamplingState -> Maybe Text) -> Maybe Text
forall a b. Maybe a -> (a -> Maybe b) -> Maybe b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= SamplingState -> Maybe Text
printSamplingStateSingle
                value :: ByteString
value = [ByteString] -> ByteString
forall a. Monoid a => [a] -> a
mconcat ([ByteString] -> ByteString) -> [ByteString] -> ByteString
forall a b. (a -> b) -> a -> b
$ ByteString -> [ByteString] -> [ByteString]
forall a. a -> [a] -> [a]
intersperse ByteString
"-" ([ByteString] -> [ByteString]) -> [ByteString] -> [ByteString]
forall a b. (a -> b) -> a -> b
$ [ByteString
traceIdValue, ByteString
spanIdValue] [ByteString] -> [ByteString] -> [ByteString]
forall a. Semigroup a => a -> a -> a
<> [Maybe ByteString] -> [ByteString]
forall a. [Maybe a] -> [a]
catMaybes [Text -> ByteString
Text.encodeUtf8 (Text -> ByteString) -> Maybe Text -> Maybe ByteString
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe Text
samplingStateValue]

            RequestHeaders -> IO RequestHeaders
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (RequestHeaders -> IO RequestHeaders)
-> RequestHeaders -> IO RequestHeaders
forall a b. (a -> b) -> a -> b
$ (HeaderName
b3Header, ByteString
value) (HeaderName, ByteString) -> RequestHeaders -> RequestHeaders
forall a. a -> [a] -> [a]
: RequestHeaders
hs
    }


b3MultiTraceContextPropagator :: Propagator Context RequestHeaders ResponseHeaders
b3MultiTraceContextPropagator :: Propagator Context RequestHeaders RequestHeaders
b3MultiTraceContextPropagator =
  Propagator
    { propagatorNames :: [Text]
propagatorNames = [Text
"B3 Multi Trace Context"]
    , extractor :: RequestHeaders -> Context -> IO Context
extractor = \RequestHeaders
hs Context
c -> do
        case RequestHeaders -> Maybe SpanContext
b3Extractor RequestHeaders
hs of
          Maybe SpanContext
Nothing -> Context -> IO Context
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Context
c
          Just SpanContext
spanContext' -> Context -> IO Context
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Context -> IO Context) -> Context -> IO Context
forall a b. (a -> b) -> a -> b
$ Span -> Context -> Context
Context.insertSpan (SpanContext -> Span
Core.wrapSpanContext SpanContext
spanContext') Context
c
    , injector :: Context -> RequestHeaders -> IO RequestHeaders
injector = \Context
c RequestHeaders
hs ->
        case Context -> Maybe Span
Context.lookupSpan Context
c of
          Maybe Span
Nothing -> RequestHeaders -> IO RequestHeaders
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure RequestHeaders
hs
          Just Span
span' -> do
            Core.SpanContext {TraceId
traceId :: SpanContext -> TraceId
traceId :: TraceId
traceId, SpanId
spanId :: SpanContext -> SpanId
spanId :: SpanId
spanId, traceState :: SpanContext -> TraceState
traceState = TS.TraceState [(Key, Value)]
traceState} <- Span -> IO SpanContext
forall (m :: * -> *). MonadIO m => Span -> m SpanContext
Core.getSpanContext Span
span'
            let traceIdValue :: ByteString
traceIdValue = TraceId -> ByteString
encodeTraceId TraceId
traceId
                spanIdValue :: ByteString
spanIdValue = SpanId -> ByteString
encodeSpanId SpanId
spanId
                samplingStateValue :: Maybe (HeaderName, Text)
samplingStateValue = Key -> [(Key, Value)] -> Maybe Value
forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup (Text -> Key
TS.Key Text
"sampling-state") [(Key, Value)]
traceState Maybe Value
-> (Value -> Maybe SamplingState) -> Maybe SamplingState
forall a b. Maybe a -> (a -> Maybe b) -> Maybe b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Value -> Maybe SamplingState
samplingStateFromValue Maybe SamplingState
-> (SamplingState -> Maybe (HeaderName, Text))
-> Maybe (HeaderName, Text)
forall a b. Maybe a -> (a -> Maybe b) -> Maybe b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= SamplingState -> Maybe (HeaderName, Text)
printSamplingStateMulti

            RequestHeaders -> IO RequestHeaders
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (RequestHeaders -> IO RequestHeaders)
-> RequestHeaders -> IO RequestHeaders
forall a b. (a -> b) -> a -> b
$
              (HeaderName
xb3TraceIdHeader, ByteString
traceIdValue)
                (HeaderName, ByteString) -> RequestHeaders -> RequestHeaders
forall a. a -> [a] -> [a]
: (HeaderName
xb3SpanIdHeader, ByteString
spanIdValue)
                (HeaderName, ByteString) -> RequestHeaders -> RequestHeaders
forall a. a -> [a] -> [a]
: RequestHeaders
hs
                RequestHeaders -> RequestHeaders -> RequestHeaders
forall a. [a] -> [a] -> [a]
++ [Maybe (HeaderName, ByteString)] -> RequestHeaders
forall a. [Maybe a] -> [a]
catMaybes [(Text -> ByteString)
-> (HeaderName, Text) -> (HeaderName, ByteString)
forall a b. (a -> b) -> (HeaderName, a) -> (HeaderName, b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Text -> ByteString
Text.encodeUtf8 ((HeaderName, Text) -> (HeaderName, ByteString))
-> Maybe (HeaderName, Text) -> Maybe (HeaderName, ByteString)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe (HeaderName, Text)
samplingStateValue]
    }


--------------------------------------------------------------------------------

{- | For both @B3@ and @B3 Multi@ formats, we must attempt single and
 multi header extraction:
 https://github.com/open-telemetry/opentelemetry-specification/blob/main/specification/context/api-propagators.md#configuration
-}
b3Extractor :: [(HeaderName, ByteString)] -> Maybe Core.SpanContext
b3Extractor :: RequestHeaders -> Maybe SpanContext
b3Extractor RequestHeaders
hs = RequestHeaders -> Maybe SpanContext
b3SingleExtractor RequestHeaders
hs Maybe SpanContext -> Maybe SpanContext -> Maybe SpanContext
forall a. Maybe a -> Maybe a -> Maybe a
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> RequestHeaders -> Maybe SpanContext
b3MultiExtractor RequestHeaders
hs


b3SingleExtractor :: [(HeaderName, ByteString)] -> Maybe Core.SpanContext
b3SingleExtractor :: RequestHeaders -> Maybe SpanContext
b3SingleExtractor RequestHeaders
hs = do
  B3SingleHeader {Maybe SpanId
SpanId
TraceId
SamplingState
traceId :: TraceId
spanId :: SpanId
samplingState :: SamplingState
parentSpanId :: Maybe SpanId
traceId :: B3SingleHeader -> TraceId
spanId :: B3SingleHeader -> SpanId
samplingState :: B3SingleHeader -> SamplingState
parentSpanId :: B3SingleHeader -> Maybe SpanId
..} <- ByteString -> Maybe B3SingleHeader
decodeB3SingleHeader (ByteString -> Maybe B3SingleHeader)
-> Maybe ByteString -> Maybe B3SingleHeader
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< HeaderName -> RequestHeaders -> Maybe ByteString
forall a b. Eq a => a -> [(a, b)] -> Maybe b
Prelude.lookup HeaderName
b3Header RequestHeaders
hs

  let traceFlags :: TraceFlags
traceFlags = if SamplingState
samplingState SamplingState -> SamplingState -> Bool
forall a. Eq a => a -> a -> Bool
== SamplingState
Accept Bool -> Bool -> Bool
|| SamplingState
samplingState SamplingState -> SamplingState -> Bool
forall a. Eq a => a -> a -> Bool
== SamplingState
Debug then Word8 -> TraceFlags
TraceFlags Word8
1 else Word8 -> TraceFlags
TraceFlags Word8
0

  SpanContext -> Maybe SpanContext
forall a. a -> Maybe a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (SpanContext -> Maybe SpanContext)
-> SpanContext -> Maybe SpanContext
forall a b. (a -> b) -> a -> b
$
    Core.SpanContext
      { traceId :: TraceId
traceId = TraceId
traceId
      , spanId :: SpanId
spanId = SpanId
spanId
      , isRemote :: Bool
isRemote = Bool
True
      , traceFlags :: TraceFlags
traceFlags = TraceFlags
traceFlags
      , traceState :: TraceState
traceState = [(Key, Value)] -> TraceState
TS.TraceState [(Text -> Key
TS.Key Text
"sampling-state", SamplingState -> Value
samplingStateToValue SamplingState
samplingState)]
      }


b3MultiExtractor :: [(HeaderName, ByteString)] -> Maybe Core.SpanContext
b3MultiExtractor :: RequestHeaders -> Maybe SpanContext
b3MultiExtractor RequestHeaders
hs = do
  TraceId
traceId <- ByteString -> Maybe TraceId
decodeXb3TraceIdHeader (ByteString -> Maybe TraceId) -> Maybe ByteString -> Maybe TraceId
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< HeaderName -> RequestHeaders -> Maybe ByteString
forall a b. Eq a => a -> [(a, b)] -> Maybe b
Prelude.lookup HeaderName
xb3TraceIdHeader RequestHeaders
hs
  SpanId
spanId <- ByteString -> Maybe SpanId
decodeXb3SpanIdHeader (ByteString -> Maybe SpanId) -> Maybe ByteString -> Maybe SpanId
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< HeaderName -> RequestHeaders -> Maybe ByteString
forall a b. Eq a => a -> [(a, b)] -> Maybe b
Prelude.lookup HeaderName
xb3SpanIdHeader RequestHeaders
hs

  let sampled :: Maybe SamplingState
sampled = ByteString -> Maybe SamplingState
decodeXb3SampledHeader (ByteString -> Maybe SamplingState)
-> Maybe ByteString -> Maybe SamplingState
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< HeaderName -> RequestHeaders -> Maybe ByteString
forall a b. Eq a => a -> [(a, b)] -> Maybe b
Prelude.lookup HeaderName
xb3SampledHeader RequestHeaders
hs
      debug :: Maybe SamplingState
debug = ByteString -> Maybe SamplingState
decodeXb3FlagsHeader (ByteString -> Maybe SamplingState)
-> Maybe ByteString -> Maybe SamplingState
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< HeaderName -> RequestHeaders -> Maybe ByteString
forall a b. Eq a => a -> [(a, b)] -> Maybe b
Prelude.lookup HeaderName
xb3FlagsHeader RequestHeaders
hs
      -- NOTE: Debug implies Accept (https://github.com/openzipkin/b3-propagation#debug-flag)
      samplingState :: SamplingState
samplingState = SamplingState -> Maybe SamplingState -> SamplingState
forall a. a -> Maybe a -> a
fromMaybe SamplingState
Defer (Maybe SamplingState -> SamplingState)
-> Maybe SamplingState -> SamplingState
forall a b. (a -> b) -> a -> b
$ Maybe SamplingState
sampled Maybe SamplingState -> Maybe SamplingState -> Maybe SamplingState
forall a. Maybe a -> Maybe a -> Maybe a
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> Maybe SamplingState
debug
  let traceFlags :: TraceFlags
traceFlags = if SamplingState
samplingState SamplingState -> SamplingState -> Bool
forall a. Eq a => a -> a -> Bool
== SamplingState
Accept Bool -> Bool -> Bool
|| SamplingState
samplingState SamplingState -> SamplingState -> Bool
forall a. Eq a => a -> a -> Bool
== SamplingState
Debug then Word8 -> TraceFlags
TraceFlags Word8
1 else Word8 -> TraceFlags
TraceFlags Word8
0

  SpanContext -> Maybe SpanContext
forall a. a -> Maybe a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (SpanContext -> Maybe SpanContext)
-> SpanContext -> Maybe SpanContext
forall a b. (a -> b) -> a -> b
$
    Core.SpanContext
      { traceId :: TraceId
traceId = TraceId
traceId
      , spanId :: SpanId
spanId = SpanId
spanId
      , isRemote :: Bool
isRemote = Bool
True
      , traceFlags :: TraceFlags
traceFlags = TraceFlags
traceFlags
      , traceState :: TraceState
traceState = [(Key, Value)] -> TraceState
TS.TraceState [(Text -> Key
TS.Key Text
"sampling-state", SamplingState -> Value
samplingStateToValue SamplingState
samplingState)]
      }