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
21module 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
34import Thrift.Protocol
35import Thrift.Protocol.Binary
36import Thrift.Protocol.JSON
37import Thrift.Protocol.Compact
38import Thrift.Transport
39import Thrift.Transport.Header
40import Data.IORef
41import qualified Data.Map as Map
42
43data ProtocolWrap = forall a. (Protocol a) => ProtocolWrap(a)
44
45instance 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
52data HeaderProtocol i o = (Transport i, Transport o) => HeaderProtocol {
53    trans :: HeaderTransport i o,
54    wrappedProto :: IORef ProtocolWrap
55  }
56
57createProtocolWrap :: Transport t => ProtocolType -> t -> ProtocolWrap
58createProtocolWrap typ t =
59  case typ of
60    TBinary -> ProtocolWrap $ BinaryProtocol t
61    TCompact -> ProtocolWrap $ CompactProtocol t
62    TJSON -> ProtocolWrap $ JSONProtocol t
63
64createHeaderProtocol :: (Transport i, Transport o) => i -> o -> IO(HeaderProtocol i o)
65createHeaderProtocol 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
71createHeaderProtocol1 :: Transport t => t -> IO(HeaderProtocol t t)
72createHeaderProtocol1 t = createHeaderProtocol t t
73
74resetProtocol :: (Transport i, Transport o) => HeaderProtocol i o -> IO ()
75resetProtocol p = do
76  pid <- readIORef $ protocolType $ trans p
77  writeIORef (wrappedProto p) $ createProtocolWrap pid $ trans p
78
79getWrapped = readIORef . wrappedProto
80
81setTransport :: (Transport i, Transport o) => HeaderProtocol i o -> HeaderTransport i o -> HeaderProtocol i o
82setTransport p t = p { trans = t }
83
84updateTransport :: (Transport i, Transport o) => HeaderProtocol i o -> (HeaderTransport i o -> HeaderTransport i o)-> HeaderProtocol i o
85updateTransport p f = setTransport p (f $ trans p)
86
87type Headers = Map.Map String String
88
89-- TODO: we want to set headers without recreating client...
90setHeader :: (Transport i, Transport o) => HeaderProtocol i o -> String -> String -> HeaderProtocol i o
91setHeader p k v = updateTransport p $ \t -> t { writeHeaders = Map.insert k v $ writeHeaders t }
92
93setHeaders :: (Transport i, Transport o) => HeaderProtocol i o -> Headers -> HeaderProtocol i o
94setHeaders p h = updateTransport p $ \t -> t { writeHeaders = h }
95
96-- TODO: make it public once we have first transform implementation for Haskell
97setTransforms :: (Transport i, Transport o) => HeaderProtocol i o -> [TransformType] -> HeaderProtocol i o
98setTransforms p trs = updateTransport p $ \t -> t { writeTransforms = trs }
99
100setTransform :: (Transport i, Transport o) => HeaderProtocol i o -> TransformType -> HeaderProtocol i o
101setTransform p tr = updateTransport p $ \t -> t { writeTransforms = tr:(writeTransforms t) }
102
103getWriteHeaders :: (Transport i, Transport o) => HeaderProtocol i o -> Headers
104getWriteHeaders = writeHeaders . trans
105
106getHeaders :: (Transport i, Transport o) => HeaderProtocol i o -> IO [(String, String)]
107getHeaders = readIORef . headers . trans
108
109getProtocolType :: (Transport i, Transport o) => HeaderProtocol i o -> IO ProtocolType
110getProtocolType p = readIORef $ protocolType $ trans p
111
112setProtocolType :: (Transport i, Transport o) => HeaderProtocol i o -> ProtocolType -> IO ()
113setProtocolType 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
121instance (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
142