{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RankNTypes #-}

-- This is an approximate implementation of https://www.w3.org/TR/trace-context

module OpenTelemetry.Propagation where

import Control.Applicative
import qualified Data.ByteString as BS
import qualified Data.ByteString.Char8 as BS8
import Data.Char (ord)
import Data.List (find)
import Data.String
import Data.Word
import OpenTelemetry.SpanContext
import Text.Printf

data PropagationFormat = PropagationFormat
  { PropagationFormat
-> forall key.
   (Semigroup key, IsString key, Eq key) =>
   [(key, ByteString)] -> Maybe SpanContext
propagateFromHeaders :: forall key. (Semigroup key, IsString key, Eq key) => [(key, BS.ByteString)] -> Maybe SpanContext,
    PropagationFormat
-> forall key.
   (Semigroup key, IsString key, Eq key) =>
   SpanContext -> [(key, ByteString)]
propagateToHeaders :: forall key. (Semigroup key, IsString key, Eq key) => SpanContext -> [(key, BS.ByteString)]
  }

-- | (p1 <> p2) parses like p1, then p2 as a fallback. (p1 <> p2) injects like p1.
instance Semigroup PropagationFormat where
  PropagationFormat forall key.
(Semigroup key, IsString key, Eq key) =>
[(key, ByteString)] -> Maybe SpanContext
from1 forall key.
(Semigroup key, IsString key, Eq key) =>
SpanContext -> [(key, ByteString)]
to1 <> :: PropagationFormat -> PropagationFormat -> PropagationFormat
<> PropagationFormat forall key.
(Semigroup key, IsString key, Eq key) =>
[(key, ByteString)] -> Maybe SpanContext
from2 forall key.
(Semigroup key, IsString key, Eq key) =>
SpanContext -> [(key, ByteString)]
_to2 =
    let from :: [(key, ByteString)] -> Maybe SpanContext
from [(key, ByteString)]
headers = forall key.
(Semigroup key, IsString key, Eq key) =>
[(key, ByteString)] -> Maybe SpanContext
from1 [(key, ByteString)]
headers forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> forall key.
(Semigroup key, IsString key, Eq key) =>
[(key, ByteString)] -> Maybe SpanContext
from2 [(key, ByteString)]
headers
        to :: SpanContext -> [(key, ByteString)]
to SpanContext
context = forall key.
(Semigroup key, IsString key, Eq key) =>
SpanContext -> [(key, ByteString)]
to1 SpanContext
context
     in (forall key.
 (Semigroup key, IsString key, Eq key) =>
 [(key, ByteString)] -> Maybe SpanContext)
-> (forall key.
    (Semigroup key, IsString key, Eq key) =>
    SpanContext -> [(key, ByteString)])
-> PropagationFormat
PropagationFormat forall key.
(Semigroup key, IsString key, Eq key) =>
[(key, ByteString)] -> Maybe SpanContext
from forall key.
(Semigroup key, IsString key, Eq key) =>
SpanContext -> [(key, ByteString)]
to

w3cTraceContext :: PropagationFormat
w3cTraceContext :: PropagationFormat
w3cTraceContext = (forall key.
 (Semigroup key, IsString key, Eq key) =>
 [(key, ByteString)] -> Maybe SpanContext)
-> (forall key.
    (Semigroup key, IsString key, Eq key) =>
    SpanContext -> [(key, ByteString)])
-> PropagationFormat
PropagationFormat forall {t :: * -> *} {b}.
(Foldable t, Eq b, IsString b) =>
t (b, ByteString) -> Maybe SpanContext
from forall {a}. IsString a => SpanContext -> [(a, ByteString)]
to
  where
    from :: t (b, ByteString) -> Maybe SpanContext
from t (b, ByteString)
headers =
      case forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
find ((forall a. Eq a => a -> a -> Bool
== b
"traceparent") forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> a
fst) t (b, ByteString)
headers of
        Just (b
_, (ByteString -> Maybe SpanContext
parseSpanContext -> Maybe SpanContext
mctx)) -> Maybe SpanContext
mctx
        Maybe (b, ByteString)
_ -> forall a. Maybe a
Nothing
    to :: SpanContext -> [(a, ByteString)]
to (SpanContext (SId Word64
sid) (TId Word64
tid)) =
      [(a
"traceparent", String -> ByteString
BS8.pack forall a b. (a -> b) -> a -> b
$ forall r. PrintfType r => String -> r
printf String
"00-%x-%x-00" Word64
tid Word64
sid)]

b3 :: PropagationFormat
b3 :: PropagationFormat
b3 = String -> PropagationFormat
typical_opentracing_format_with_prefix String
"x-b3-"

otTracer :: PropagationFormat
otTracer :: PropagationFormat
otTracer = String -> PropagationFormat
typical_opentracing_format_with_prefix String
"ot-tracer-"

typical_opentracing_format_with_prefix :: String -> PropagationFormat
typical_opentracing_format_with_prefix :: String -> PropagationFormat
typical_opentracing_format_with_prefix String
prefix = (forall key.
 (Semigroup key, IsString key, Eq key) =>
 [(key, ByteString)] -> Maybe SpanContext)
-> (forall key.
    (Semigroup key, IsString key, Eq key) =>
    SpanContext -> [(key, ByteString)])
-> PropagationFormat
PropagationFormat forall key.
(Semigroup key, IsString key, Eq key) =>
[(key, ByteString)] -> Maybe SpanContext
from forall {a}.
(Semigroup a, IsString a) =>
SpanContext -> [(a, ByteString)]
to
  where
    to :: SpanContext -> [(a, ByteString)]
to (SpanContext (SId Word64
sid) (TId Word64
tid)) =
      [ (forall a. IsString a => String -> a
fromString String
prefix forall a. Semigroup a => a -> a -> a
<> a
"traceid", Word64 -> ByteString
encode_u64 Word64
tid),
        (forall a. IsString a => String -> a
fromString String
prefix forall a. Semigroup a => a -> a -> a
<> a
"spanid", Word64 -> ByteString
encode_u64 Word64
sid),
        (forall a. IsString a => String -> a
fromString String
prefix forall a. Semigroup a => a -> a -> a
<> a
"sampled", ByteString
"true")
      ]
    from :: [(a, ByteString)] -> Maybe SpanContext
from [(a, ByteString)]
headers =
      let traceidKey :: a
traceidKey = forall a. IsString a => String -> a
fromString String
prefix forall a. Semigroup a => a -> a -> a
<> a
"traceid"
          spanidKey :: a
spanidKey = forall a. IsString a => String -> a
fromString String
prefix forall a. Semigroup a => a -> a -> a
<> a
"spanid"
          go :: [(a, ByteString)]
-> (Maybe Word64, Maybe Word64) -> Maybe (Word64, Word64)
go [(a, ByteString)]
_ (Just Word64
tid, Just Word64
sid) = forall a. a -> Maybe a
Just (Word64
tid, Word64
sid)
          go [] (Maybe Word64, Maybe Word64)
_ = forall a. Maybe a
Nothing
          go ((a
k, ByteString
v) : [(a, ByteString)]
xs) (Maybe Word64
tid, Maybe Word64
sid)
            | a
k forall a. Eq a => a -> a -> Bool
== a
traceidKey = [(a, ByteString)]
-> (Maybe Word64, Maybe Word64) -> Maybe (Word64, Word64)
go [(a, ByteString)]
xs (ByteString -> Maybe Word64
decode_u64 ByteString
v, Maybe Word64
sid)
            | a
k forall a. Eq a => a -> a -> Bool
== a
spanidKey = [(a, ByteString)]
-> (Maybe Word64, Maybe Word64) -> Maybe (Word64, Word64)
go [(a, ByteString)]
xs (Maybe Word64
tid, ByteString -> Maybe Word64
decode_u64 ByteString
v)
            | Bool
otherwise = [(a, ByteString)]
-> (Maybe Word64, Maybe Word64) -> Maybe (Word64, Word64)
go [(a, ByteString)]
xs (Maybe Word64
tid, Maybe Word64
sid)
       in (\(Word64
t, Word64
s) -> SpanId -> TraceId -> SpanContext
SpanContext (Word64 -> SpanId
SId Word64
s) (Word64 -> TraceId
TId Word64
t)) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [(a, ByteString)]
-> (Maybe Word64, Maybe Word64) -> Maybe (Word64, Word64)
go [(a, ByteString)]
headers (forall a. Maybe a
Nothing, forall a. Maybe a
Nothing)

parseSpanContext :: BS.ByteString -> Maybe SpanContext
parseSpanContext :: ByteString -> Maybe SpanContext
parseSpanContext ByteString
input =
  case Char -> ByteString -> [ByteString]
BS8.split Char
'-' ByteString
input of
    [ByteString
"00", (ByteString -> Maybe Word64
fromHex -> Just Word64
tid), (ByteString -> Maybe Word64
fromHex -> Just Word64
sid), ByteString
_] ->
      forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ SpanId -> TraceId -> SpanContext
SpanContext (Word64 -> SpanId
SId Word64
sid) (Word64 -> TraceId
TId Word64
tid)
    [ByteString]
_ -> forall a. Maybe a
Nothing

isLowerHexDigit :: Char -> Bool
isLowerHexDigit :: Char -> Bool
isLowerHexDigit (Char -> Int
ord -> Int
w) = (Int
w forall a. Ord a => a -> a -> Bool
>= Int
48 Bool -> Bool -> Bool
&& Int
w forall a. Ord a => a -> a -> Bool
<= Int
57) Bool -> Bool -> Bool
|| (Int
w forall a. Ord a => a -> a -> Bool
>= Int
97 Bool -> Bool -> Bool
&& Int
w forall a. Ord a => a -> a -> Bool
<= Int
102)

fromHex :: BS.ByteString -> Maybe Word64
fromHex :: ByteString -> Maybe Word64
fromHex ByteString
bytes = forall a. (a -> Char -> a) -> a -> ByteString -> a
BS8.foldl' forall {a}. Num a => Maybe a -> Char -> Maybe a
go (forall a. a -> Maybe a
Just Word64
0) ByteString
bytes
  where
    go :: Maybe a -> Char -> Maybe a
go Maybe a
Nothing Char
_ = forall a. Maybe a
Nothing
    go (Just !a
result) (Char -> Int
ord -> Int
d) | Int
d forall a. Ord a => a -> a -> Bool
>= Int
48 Bool -> Bool -> Bool
&& Int
d forall a. Ord a => a -> a -> Bool
< Int
58 = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ a
result forall a. Num a => a -> a -> a
* a
16 forall a. Num a => a -> a -> a
+ forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
d forall a. Num a => a -> a -> a
- a
48
    go (Just a
result) (Char -> Int
ord -> Int
d) | Int
d forall a. Ord a => a -> a -> Bool
>= Int
97 Bool -> Bool -> Bool
&& Int
d forall a. Ord a => a -> a -> Bool
< Int
124 = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ a
result forall a. Num a => a -> a -> a
* a
16 forall a. Num a => a -> a -> a
+ forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
d forall a. Num a => a -> a -> a
- a
87
    go Maybe a
_ Char
_ = forall a. Maybe a
Nothing

encode_u64 :: Word64 -> BS.ByteString
encode_u64 :: Word64 -> ByteString
encode_u64 Word64
x = String -> ByteString
BS8.pack (forall r. PrintfType r => String -> r
printf String
"%016x" Word64
x)

decode_u64 :: BS.ByteString -> Maybe Word64
decode_u64 :: ByteString -> Maybe Word64
decode_u64 ByteString
bytes | ByteString -> Int
BS.length ByteString
bytes forall a. Ord a => a -> a -> Bool
> Int
16 = forall a. Maybe a
Nothing
decode_u64 ByteString
bytes = forall a. (a -> Word8 -> a) -> a -> ByteString -> a
BS.foldl' forall {a} {a}. (Num a, Integral a) => Maybe a -> a -> Maybe a
go (forall a. a -> Maybe a
Just Word64
0) ByteString
bytes
  where
    go :: Maybe a -> a -> Maybe a
go Maybe a
Nothing a
_ = forall a. Maybe a
Nothing
    go (Just !a
result) a
d | a
d forall a. Ord a => a -> a -> Bool
>= a
48 Bool -> Bool -> Bool
&& a
d forall a. Ord a => a -> a -> Bool
< a
58 = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ a
result forall a. Num a => a -> a -> a
* a
16 forall a. Num a => a -> a -> a
+ forall a b. (Integral a, Num b) => a -> b
fromIntegral a
d forall a. Num a => a -> a -> a
- a
48
    go (Just a
result) a
d | a
d forall a. Ord a => a -> a -> Bool
>= a
97 Bool -> Bool -> Bool
&& a
d forall a. Ord a => a -> a -> Bool
< a
124 = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ a
result forall a. Num a => a -> a -> a
* a
16 forall a. Num a => a -> a -> a
+ forall a b. (Integral a, Num b) => a -> b
fromIntegral a
d forall a. Num a => a -> a -> a
- a
87
    go Maybe a
_ a
_ = forall a. Maybe a
Nothing