{-# 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 = [(key, ByteString)] -> Maybe SpanContext
forall key.
(Semigroup key, IsString key, Eq key) =>
[(key, ByteString)] -> Maybe SpanContext
from1 [(key, ByteString)]
headers Maybe SpanContext -> Maybe SpanContext -> Maybe SpanContext
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> [(key, ByteString)] -> Maybe SpanContext
forall key.
(Semigroup key, IsString key, Eq key) =>
[(key, ByteString)] -> Maybe SpanContext
from2 [(key, ByteString)]
headers
        to :: SpanContext -> [(key, ByteString)]
to SpanContext
context = SpanContext -> [(key, ByteString)]
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 key.
(Semigroup key, IsString key, Eq key) =>
[(key, ByteString)] -> Maybe SpanContext
forall (t :: * -> *) a.
(Foldable t, Eq a, IsString a) =>
t (a, ByteString) -> Maybe SpanContext
from forall a. IsString a => SpanContext -> [(a, ByteString)]
forall key.
(Semigroup key, IsString key, Eq key) =>
SpanContext -> [(key, ByteString)]
to
  where
    from :: t (a, ByteString) -> Maybe SpanContext
from t (a, ByteString)
headers =
      case ((a, ByteString) -> Bool)
-> t (a, ByteString) -> Maybe (a, ByteString)
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
find ((a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
"traceparent") (a -> Bool) -> ((a, ByteString) -> a) -> (a, ByteString) -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a, ByteString) -> a
forall a b. (a, b) -> a
fst) t (a, ByteString)
headers of
        Just (a
_, (ByteString -> Maybe SpanContext
parseSpanContext -> Maybe SpanContext
mctx)) -> Maybe SpanContext
mctx
        Maybe (a, ByteString)
_ -> Maybe SpanContext
forall a. Maybe a
Nothing
    to :: SpanContext -> [(a, ByteString)]
to (SpanContext (SId Word64
sid) (TId Word64
tid)) =
      [(a
"traceparent", String -> ByteString
BS8.pack (String -> ByteString) -> String -> ByteString
forall a b. (a -> b) -> a -> b
$ String -> Word64 -> Word64 -> String
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 key.
(Semigroup key, IsString key, Eq key) =>
SpanContext -> [(key, ByteString)]
forall a.
(Semigroup a, IsString a) =>
SpanContext -> [(a, ByteString)]
to
  where
    to :: SpanContext -> [(a, ByteString)]
to (SpanContext (SId Word64
sid) (TId Word64
tid)) =
      [ (String -> a
forall a. IsString a => String -> a
fromString String
prefix a -> a -> a
forall a. Semigroup a => a -> a -> a
<> a
"traceid", Word64 -> ByteString
encode_u64 Word64
tid),
        (String -> a
forall a. IsString a => String -> a
fromString String
prefix a -> a -> a
forall a. Semigroup a => a -> a -> a
<> a
"spanid", Word64 -> ByteString
encode_u64 Word64
sid),
        (String -> a
forall a. IsString a => String -> a
fromString String
prefix a -> a -> a
forall a. Semigroup a => a -> a -> a
<> a
"sampled", ByteString
"true")
      ]
    from :: [(a, ByteString)] -> Maybe SpanContext
from [(a, ByteString)]
headers =
      let traceidKey :: a
traceidKey = String -> a
forall a. IsString a => String -> a
fromString String
prefix a -> a -> a
forall a. Semigroup a => a -> a -> a
<> a
"traceid"
          spanidKey :: a
spanidKey = String -> a
forall a. IsString a => String -> a
fromString String
prefix a -> a -> a
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) = (Word64, Word64) -> Maybe (Word64, Word64)
forall a. a -> Maybe a
Just (Word64
tid, Word64
sid)
          go [] (Maybe Word64, Maybe Word64)
_ = Maybe (Word64, Word64)
forall a. Maybe a
Nothing
          go ((a
k, ByteString
v) : [(a, ByteString)]
xs) (Maybe Word64
tid, Maybe Word64
sid)
            | a
k a -> a -> Bool
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 a -> a -> Bool
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)) ((Word64, Word64) -> SpanContext)
-> Maybe (Word64, Word64) -> Maybe SpanContext
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 (Maybe Word64
forall a. Maybe a
Nothing, Maybe Word64
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
_] ->
      SpanContext -> Maybe SpanContext
forall a. a -> Maybe a
Just (SpanContext -> Maybe SpanContext)
-> SpanContext -> Maybe SpanContext
forall a b. (a -> b) -> a -> b
$ SpanId -> TraceId -> SpanContext
SpanContext (Word64 -> SpanId
SId Word64
sid) (Word64 -> TraceId
TId Word64
tid)
    [ByteString]
_ -> Maybe SpanContext
forall a. Maybe a
Nothing

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

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

encode_u64 :: Word64 -> BS.ByteString
encode_u64 :: Word64 -> ByteString
encode_u64 Word64
x = String -> ByteString
BS8.pack (String -> Word64 -> String
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 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
16 = Maybe Word64
forall a. Maybe a
Nothing
decode_u64 ByteString
bytes = (Maybe Word64 -> Word8 -> Maybe Word64)
-> Maybe Word64 -> ByteString -> Maybe Word64
forall a. (a -> Word8 -> a) -> a -> ByteString -> a
BS.foldl' Maybe Word64 -> Word8 -> Maybe Word64
forall a a. (Num a, Integral a) => Maybe a -> a -> Maybe a
go (Word64 -> Maybe Word64
forall a. a -> Maybe a
Just Word64
0) ByteString
bytes
  where
    go :: Maybe a -> a -> Maybe a
go Maybe a
Nothing a
_ = Maybe a
forall a. Maybe a
Nothing
    go (Just !a
result) a
d | a
d a -> a -> Bool
forall a. Ord a => a -> a -> Bool
>= a
48 Bool -> Bool -> Bool
&& a
d a -> a -> Bool
forall a. Ord a => a -> a -> Bool
< a
58 = a -> Maybe a
forall a. a -> Maybe a
Just (a -> Maybe a) -> a -> Maybe a
forall a b. (a -> b) -> a -> b
$ a
result a -> a -> a
forall a. Num a => a -> a -> a
* a
16 a -> a -> a
forall a. Num a => a -> a -> a
+ a -> a
forall a b. (Integral a, Num b) => a -> b
fromIntegral a
d a -> a -> a
forall a. Num a => a -> a -> a
- a
48
    go (Just a
result) a
d | a
d a -> a -> Bool
forall a. Ord a => a -> a -> Bool
>= a
97 Bool -> Bool -> Bool
&& a
d a -> a -> Bool
forall a. Ord a => a -> a -> Bool
< a
124 = a -> Maybe a
forall a. a -> Maybe a
Just (a -> Maybe a) -> a -> Maybe a
forall a b. (a -> b) -> a -> b
$ a
result a -> a -> a
forall a. Num a => a -> a -> a
* a
16 a -> a -> a
forall a. Num a => a -> a -> a
+ a -> a
forall a b. (Integral a, Num b) => a -> b
fromIntegral a
d a -> a -> a
forall a. Num a => a -> a -> a
- a
87
    go Maybe a
_ a
_ = Maybe a
forall a. Maybe a
Nothing