ceph/src/jaegertracing/thrift/lib/hs/src/Thrift/Protocol/Header.hs

   1 --
   2 -- Licensed to the Apache Software Foundation (ASF) under one
   3 -- or more contributor license agreements. See the NOTICE file
   4 -- distributed with this work for additional information
   5 -- regarding copyright ownership. The ASF licenses this file
   6 -- to you under the Apache License, Version 2.0 (the
   7 -- "License"); you may not use this file except in compliance
   8 -- with the License. You may obtain a copy of the License at
   9 --
  10 --   http://www.apache.org/licenses/LICENSE-2.0
  11 --
  12 -- Unless required by applicable law or agreed to in writing,
  13 -- software distributed under the License is distributed on an
  14 -- "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
  15 -- KIND, either express or implied. See the License for the
  16 -- specific language governing permissions and limitations
  17 -- under the License.
  18 --
  19
  20
  21 module Thrift.Protocol.Header
  22     ( module Thrift.Protocol
  23     , HeaderProtocol(..)
  24     , getProtocolType
  25     , setProtocolType
  26     , getHeaders
  27     , getWriteHeaders
  28     , setHeader
  29     , setHeaders
  30     , createHeaderProtocol
  31     , createHeaderProtocol1
  32     ) where
  33
  34 import Thrift.Protocol
  35 import Thrift.Protocol.Binary
  36 import Thrift.Protocol.JSON
  37 import Thrift.Protocol.Compact
  38 import Thrift.Transport
  39 import Thrift.Transport.Header
  40 import Data.IORef
  41 import qualified Data.Map as Map
  42
  43 data ProtocolWrap = forall a. (Protocol a) => ProtocolWrap(a)
  44
  45 instance Protocol ProtocolWrap where
  46   readByte (ProtocolWrap p) = readByte p
  47   readVal (ProtocolWrap p) = readVal p
  48   readMessage (ProtocolWrap p) = readMessage p
  49   writeVal (ProtocolWrap p) = writeVal p
  50   writeMessage (ProtocolWrap p) = writeMessage p
  51
  52 data HeaderProtocol i o = (Transport i, Transport o) => HeaderProtocol {
  53     trans :: HeaderTransport i o,
  54     wrappedProto :: IORef ProtocolWrap
  55   }
  56
  57 createProtocolWrap :: Transport t => ProtocolType -> t -> ProtocolWrap
  58 createProtocolWrap typ t =
  59   case typ of
  60     TBinary -> ProtocolWrap $ BinaryProtocol t
  61     TCompact -> ProtocolWrap $ CompactProtocol t
  62     TJSON -> ProtocolWrap $ JSONProtocol t
  63
  64 createHeaderProtocol :: (Transport i, Transport o) => i -> o -> IO(HeaderProtocol i o)
  65 createHeaderProtocol i o = do
  66   t <- openHeaderTransport i o
  67   pid <- readIORef $ protocolType t
  68   proto <- newIORef $ createProtocolWrap pid t
  69   return $ HeaderProtocol { trans = t, wrappedProto = proto }
  70
  71 createHeaderProtocol1 :: Transport t => t -> IO(HeaderProtocol t t)
  72 createHeaderProtocol1 t = createHeaderProtocol t t
  73
  74 resetProtocol :: (Transport i, Transport o) => HeaderProtocol i o -> IO ()
  75 resetProtocol p = do
  76   pid <- readIORef $ protocolType $ trans p
  77   writeIORef (wrappedProto p) $ createProtocolWrap pid $ trans p
  78
  79 getWrapped = readIORef . wrappedProto
  80
  81 setTransport :: (Transport i, Transport o) => HeaderProtocol i o -> HeaderTransport i o -> HeaderProtocol i o
  82 setTransport p t = p { trans = t }
  83
  84 updateTransport :: (Transport i, Transport o) => HeaderProtocol i o -> (HeaderTransport i o -> HeaderTransport i o)-> HeaderProtocol i o
  85 updateTransport p f = setTransport p (f $ trans p)
  86
  87 type Headers = Map.Map String String
  88
  89 -- TODO: we want to set headers without recreating client...
  90 setHeader :: (Transport i, Transport o) => HeaderProtocol i o -> String -> String -> HeaderProtocol i o
  91 setHeader p k v = updateTransport p $ \t -> t { writeHeaders = Map.insert k v $ writeHeaders t }
  92
  93 setHeaders :: (Transport i, Transport o) => HeaderProtocol i o -> Headers -> HeaderProtocol i o
  94 setHeaders p h = updateTransport p $ \t -> t { writeHeaders = h }
  95
  96 -- TODO: make it public once we have first transform implementation for Haskell
  97 setTransforms :: (Transport i, Transport o) => HeaderProtocol i o -> [TransformType] -> HeaderProtocol i o
  98 setTransforms p trs = updateTransport p $ \t -> t { writeTransforms = trs }
  99
 100 setTransform :: (Transport i, Transport o) => HeaderProtocol i o -> TransformType -> HeaderProtocol i o
 101 setTransform p tr = updateTransport p $ \t -> t { writeTransforms = tr:(writeTransforms t) }
 102
 103 getWriteHeaders :: (Transport i, Transport o) => HeaderProtocol i o -> Headers
 104 getWriteHeaders = writeHeaders . trans
 105
 106 getHeaders :: (Transport i, Transport o) => HeaderProtocol i o -> IO [(String, String)]
 107 getHeaders = readIORef . headers . trans
 108
 109 getProtocolType :: (Transport i, Transport o) => HeaderProtocol i o -> IO ProtocolType
 110 getProtocolType p = readIORef $ protocolType $ trans p
 111
 112 setProtocolType :: (Transport i, Transport o) => HeaderProtocol i o -> ProtocolType -> IO ()
 113 setProtocolType p typ = do
 114   typ0 <- getProtocolType p
 115   if typ == typ0
 116     then return ()
 117     else do
 118       tSetProtocol (trans p) typ
 119       resetProtocol p
 120
 121 instance (Transport i, Transport o) => Protocol (HeaderProtocol i o) where
 122   readByte p = tReadAll (trans p) 1
 123
 124   readVal p tp = do
 125     proto <- getWrapped p
 126     readVal proto tp
 127
 128   readMessage p f = do
 129     tResetProtocol (trans p)
 130     resetProtocol p
 131     proto <- getWrapped p
 132     readMessage proto f
 133
 134   writeVal p v = do
 135     proto <- getWrapped p
 136     writeVal proto v
 137
 138   writeMessage p x f = do
 139     proto <- getWrapped p
 140     writeMessage proto x f
 141