{-# OPTIONS_GHC -Wno-orphans #-}

module Network.GRPC.Spec.Serialization.TraceContext (
    buildTraceContext
  , parseTraceContext
  ) where

import Control.Applicative (many)
import Control.Monad.Except
import Data.Binary (Binary(..))
import Data.Binary qualified as Binary
import Data.Binary.Get qualified as Get
import Data.Binary.Put qualified as Put
import Data.ByteString qualified as Strict (ByteString)
import Data.ByteString.Lazy qualified as BS.Lazy
import Data.Default
import Data.Maybe (maybeToList)
import Data.Word

import Network.GRPC.Spec

{-------------------------------------------------------------------------------
  Serialization
-------------------------------------------------------------------------------}

-- | Serialize t'TraceContext'
buildTraceContext :: TraceContext -> Strict.ByteString
buildTraceContext :: TraceContext -> ByteString
buildTraceContext = LazyByteString -> ByteString
BS.Lazy.toStrict (LazyByteString -> ByteString)
-> (TraceContext -> LazyByteString) -> TraceContext -> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TraceContext -> LazyByteString
forall a. Binary a => a -> LazyByteString
Binary.encode

-- | Parse t'TraceContext'
parseTraceContext :: MonadError String m => Strict.ByteString -> m TraceContext
parseTraceContext :: forall (m :: * -> *).
MonadError String m =>
ByteString -> m TraceContext
parseTraceContext ByteString
bs =
    case LazyByteString
-> Either
     (LazyByteString, ByteOffset, String)
     (LazyByteString, ByteOffset, TraceContext)
forall a.
Binary a =>
LazyByteString
-> Either
     (LazyByteString, ByteOffset, String)
     (LazyByteString, ByteOffset, a)
Binary.decodeOrFail (ByteString -> LazyByteString
BS.Lazy.fromStrict ByteString
bs) of
      Right (LazyByteString
_, ByteOffset
_, TraceContext
ctxt) -> TraceContext -> m TraceContext
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return TraceContext
ctxt
      Left  (LazyByteString
_, ByteOffset
_, String
err)  -> String -> m TraceContext
forall a. String -> m a
forall e (m :: * -> *) a. MonadError e m => e -> m a
throwError String
err

{-------------------------------------------------------------------------------
  Internal auxiliary: parsing

  <https://github.com/census-instrumentation/opencensus-specs/blob/master/encodings/BinaryEncoding.md>
-------------------------------------------------------------------------------}

instance Binary TraceId where
  put :: TraceId -> Put
put = ByteString -> Put
Put.putByteString (ByteString -> Put) -> (TraceId -> ByteString) -> TraceId -> Put
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TraceId -> ByteString
getTraceId
  get :: Get TraceId
get = ByteString -> TraceId
TraceId (ByteString -> TraceId) -> Get ByteString -> Get TraceId
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> Get ByteString
Get.getByteString Int
16

instance Binary SpanId where
  put :: SpanId -> Put
put = ByteString -> Put
Put.putByteString (ByteString -> Put) -> (SpanId -> ByteString) -> SpanId -> Put
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SpanId -> ByteString
getSpanId
  get :: Get SpanId
get = ByteString -> SpanId
SpanId (ByteString -> SpanId) -> Get ByteString -> Get SpanId
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> Get ByteString
Get.getByteString Int
8

instance Binary TraceOptions where
  put :: TraceOptions -> Put
put = Word8 -> Put
Put.putWord8 (Word8 -> Put) -> (TraceOptions -> Word8) -> TraceOptions -> Put
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TraceOptions -> Word8
traceOptionsToWord8
  get :: Get TraceOptions
get = Word8 -> Get TraceOptions
forall (m :: * -> *). MonadFail m => Word8 -> m TraceOptions
traceOptionsFromWord8 (Word8 -> Get TraceOptions) -> Get Word8 -> Get TraceOptions
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Get Word8
Get.getWord8

instance Binary Field where
  put :: Field -> Put
put (FieldTraceId TraceId
tid)  = Word8 -> Put
Put.putWord8 Word8
0 Put -> Put -> Put
forall a. Semigroup a => a -> a -> a
<> TraceId -> Put
forall t. Binary t => t -> Put
put TraceId
tid
  put (FieldSpanId  SpanId
sid)  = Word8 -> Put
Put.putWord8 Word8
1 Put -> Put -> Put
forall a. Semigroup a => a -> a -> a
<> SpanId -> Put
forall t. Binary t => t -> Put
put SpanId
sid
  put (FieldOptions TraceOptions
opts) = Word8 -> Put
Put.putWord8 Word8
2 Put -> Put -> Put
forall a. Semigroup a => a -> a -> a
<> TraceOptions -> Put
forall t. Binary t => t -> Put
put TraceOptions
opts

  get :: Get Field
get = do
      fieldId <- Get Word8
Get.getWord8
      case fieldId of
        Word8
0 -> TraceId -> Field
FieldTraceId (TraceId -> Field) -> Get TraceId -> Get Field
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Get TraceId
forall t. Binary t => Get t
get
        Word8
1 -> SpanId -> Field
FieldSpanId  (SpanId -> Field) -> Get SpanId -> Get Field
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Get SpanId
forall t. Binary t => Get t
get
        Word8
2 -> TraceOptions -> Field
FieldOptions (TraceOptions -> Field) -> Get TraceOptions -> Get Field
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Get TraceOptions
forall t. Binary t => Get t
get
        Word8
_ -> String -> Get Field
forall a. String -> Get a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail (String -> Get Field) -> String -> Get Field
forall a b. (a -> b) -> a -> b
$ String
"Invalid fieldId " String -> String -> String
forall a. [a] -> [a] -> [a]
++ Word8 -> String
forall a. Show a => a -> String
show Word8
fieldId

instance Binary TraceContext where
  put :: TraceContext -> Put
put TraceContext
ctxt = [Put] -> Put
forall a. Monoid a => [a] -> a
mconcat [
        Word8 -> Put
Put.putWord8 Word8
0 -- Version 0
      , (Field -> Put) -> [Field] -> Put
forall m a. Monoid m => (a -> m) -> [a] -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap Field -> Put
forall t. Binary t => t -> Put
put (TraceContext -> [Field]
traceContextToFields TraceContext
ctxt)
      ]

  get :: Get TraceContext
get = do
      version <- Get Word8
Get.getWord8
      case version of
        Word8
0 -> [Field] -> Get TraceContext
forall (m :: * -> *). MonadFail m => [Field] -> m TraceContext
traceContextFromFields ([Field] -> Get TraceContext) -> Get [Field] -> Get TraceContext
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Get Field -> Get [Field]
forall a. Get a -> Get [a]
forall (f :: * -> *) a. Alternative f => f a -> f [a]
many Get Field
forall t. Binary t => Get t
get
        Word8
_ -> String -> Get TraceContext
forall a. String -> Get a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail (String -> Get TraceContext) -> String -> Get TraceContext
forall a b. (a -> b) -> a -> b
$ String
"Invalid version " String -> String -> String
forall a. [a] -> [a] -> [a]
++ Word8 -> String
forall a. Show a => a -> String
show Word8
version

{-------------------------------------------------------------------------------
  Internal: fields
-------------------------------------------------------------------------------}

data Field =
    FieldTraceId TraceId
  | FieldSpanId SpanId
  | FieldOptions TraceOptions

traceContextToFields :: TraceContext -> [Field]
traceContextToFields :: TraceContext -> [Field]
traceContextToFields (TraceContext Maybe TraceId
tid Maybe SpanId
sid Maybe TraceOptions
opts) = [[Field]] -> [Field]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [
      TraceId -> Field
FieldTraceId (TraceId -> Field) -> [TraceId] -> [Field]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe TraceId -> [TraceId]
forall a. Maybe a -> [a]
maybeToList Maybe TraceId
tid
    , SpanId -> Field
FieldSpanId  (SpanId -> Field) -> [SpanId] -> [Field]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe SpanId -> [SpanId]
forall a. Maybe a -> [a]
maybeToList Maybe SpanId
sid
    , TraceOptions -> Field
FieldOptions (TraceOptions -> Field) -> [TraceOptions] -> [Field]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe TraceOptions -> [TraceOptions]
forall a. Maybe a -> [a]
maybeToList Maybe TraceOptions
opts
    ]

traceContextFromFields :: forall m. MonadFail m => [Field] -> m TraceContext
traceContextFromFields :: forall (m :: * -> *). MonadFail m => [Field] -> m TraceContext
traceContextFromFields = ([Field] -> TraceContext -> m TraceContext)
-> TraceContext -> [Field] -> m TraceContext
forall a b c. (a -> b -> c) -> b -> a -> c
flip [Field] -> TraceContext -> m TraceContext
go TraceContext
forall a. Default a => a
def
  where
    go :: [Field] -> TraceContext -> m TraceContext
    go :: [Field] -> TraceContext -> m TraceContext
go []     TraceContext
acc = TraceContext -> m TraceContext
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return TraceContext
acc
    go (Field
f:[Field]
fs) TraceContext
acc =
        case Field
f of
          FieldTraceId TraceId
tid ->
            case TraceContext -> Maybe TraceId
traceContextTraceId TraceContext
acc of
              Maybe TraceId
Nothing -> [Field] -> TraceContext -> m TraceContext
go [Field]
fs (TraceContext -> m TraceContext) -> TraceContext -> m TraceContext
forall a b. (a -> b) -> a -> b
$ TraceContext
acc{traceContextTraceId = Just tid}
              Just TraceId
_  -> String -> m TraceContext
forall a. String -> m a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"Multiple TraceId fields"
          FieldSpanId SpanId
sid ->
            case TraceContext -> Maybe SpanId
traceContextSpanId TraceContext
acc of
              Maybe SpanId
Nothing -> [Field] -> TraceContext -> m TraceContext
go [Field]
fs (TraceContext -> m TraceContext) -> TraceContext -> m TraceContext
forall a b. (a -> b) -> a -> b
$ TraceContext
acc{traceContextSpanId = Just sid}
              Just SpanId
_  -> String -> m TraceContext
forall a. String -> m a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"Multiple SpanId fields"
          FieldOptions TraceOptions
opts ->
            case TraceContext -> Maybe TraceOptions
traceContextOptions TraceContext
acc of
              Maybe TraceOptions
Nothing -> [Field] -> TraceContext -> m TraceContext
go [Field]
fs (TraceContext -> m TraceContext) -> TraceContext -> m TraceContext
forall a b. (a -> b) -> a -> b
$ TraceContext
acc{traceContextOptions = Just opts}
              Just TraceOptions
_  -> String -> m TraceContext
forall a. String -> m a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"Multiple TraceOptions fields"

{-------------------------------------------------------------------------------
  Internal: dealing with 'TraceOptions'

  We take advantage of the fact that currently only a single option is defined.
  Once we have more than one, this code will be a bit more complicated.
-------------------------------------------------------------------------------}

traceOptionsToWord8 :: TraceOptions -> Word8
traceOptionsToWord8 :: TraceOptions -> Word8
traceOptionsToWord8 (TraceOptions Bool
False) = Word8
0
traceOptionsToWord8 (TraceOptions Bool
True)  = Word8
1

traceOptionsFromWord8 :: MonadFail m => Word8 -> m TraceOptions
traceOptionsFromWord8 :: forall (m :: * -> *). MonadFail m => Word8 -> m TraceOptions
traceOptionsFromWord8 Word8
0 = TraceOptions -> m TraceOptions
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (TraceOptions -> m TraceOptions) -> TraceOptions -> m TraceOptions
forall a b. (a -> b) -> a -> b
$ Bool -> TraceOptions
TraceOptions Bool
False
traceOptionsFromWord8 Word8
1 = TraceOptions -> m TraceOptions
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (TraceOptions -> m TraceOptions) -> TraceOptions -> m TraceOptions
forall a b. (a -> b) -> a -> b
$ Bool -> TraceOptions
TraceOptions Bool
True
traceOptionsFromWord8 Word8
n = String -> m TraceOptions
forall a. String -> m a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail (String -> m TraceOptions) -> String -> m TraceOptions
forall a b. (a -> b) -> a -> b
$ String
"Invalid TraceOptions " String -> String -> String
forall a. [a] -> [a] -> [a]
++ Word8 -> String
forall a. Show a => a -> String
show Word8
n