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{-# LANGUAGE CPP #-}
21{-# LANGUAGE ExistentialQuantification #-}
22{-# LANGUAGE OverloadedStrings #-}
23{-# LANGUAGE ScopedTypeVariables #-}
24
25module Thrift.Protocol.Binary
26    ( module Thrift.Protocol
27    , BinaryProtocol(..)
28    ) where
29
30import Control.Exception ( throw )
31import Control.Monad
32import Data.Bits
33import Data.ByteString.Builder
34#if __GLASGOW_HASKELL__ < 710
35import Data.Functor
36#endif
37import Data.Int
38#if __GLASGOW_HASKELL__ < 804
39import Data.Monoid
40#endif
41import Data.Word
42import Data.Text.Lazy.Encoding ( decodeUtf8, encodeUtf8 )
43
44import Thrift.Protocol
45import Thrift.Transport
46import Thrift.Types
47
48import qualified Data.Attoparsec.ByteString as P
49import qualified Data.Attoparsec.ByteString.Lazy as LP
50import qualified Data.Binary as Binary
51import qualified Data.ByteString.Lazy as LBS
52import qualified Data.HashMap.Strict as Map
53import qualified Data.Text.Lazy as LT
54
55-- | The Binary Protocol uses the standard Thrift 'TBinaryProtocol'
56data BinaryProtocol a = BinaryProtocol a
57                        -- ^ Construct a 'BinaryProtocol' with a 'Transport'
58
59versionMask :: Int32
60versionMask = fromIntegral (0xffff0000 :: Word32)
61
62version1 :: Int32
63version1 = fromIntegral (0x80010000 :: Word32)
64
65-- NOTE: Reading and Writing functions rely on Builders and Data.Binary to
66-- encode and decode data.  Data.Binary assumes that the binary values it is
67-- encoding to and decoding from are in BIG ENDIAN format, and converts the
68-- endianness as necessary to match the local machine.
69instance Protocol BinaryProtocol where
70    mkProtocol = BinaryProtocol
71    getTransport (BinaryProtocol t) = t
72
73    writeMessage p (n, t, s) = (writeMessageBegin >>)
74      where
75        writeMessageBegin = tWrite (getTransport p) $ toLazyByteString $
76          buildBinaryValue (TI32 (version1 .|. fromIntegral (fromEnum t))) <>
77          buildBinaryValue (TString $ encodeUtf8 n) <>
78          buildBinaryValue (TI32 s)
79
80    readMessage p = (readMessageBegin >>=)
81      where
82        readMessageBegin =runParser p $ do
83          TI32 ver <- parseBinaryValue T_I32
84          if ver .&. versionMask /= version1
85            then throw $ ProtocolExn PE_BAD_VERSION "Missing version identifier"
86            else do
87              TString s <- parseBinaryValue T_STRING
88              TI32 sz <- parseBinaryValue T_I32
89              return (decodeUtf8 s, toEnum $ fromIntegral $ ver .&. 0xFF, sz)
90
91    serializeVal _ = toLazyByteString . buildBinaryValue
92    deserializeVal _ ty bs =
93      case LP.eitherResult $ LP.parse (parseBinaryValue ty) bs of
94        Left s -> error s
95        Right val -> val
96
97    readVal p = runParser p . parseBinaryValue
98
99-- | Writing Functions
100buildBinaryValue :: ThriftVal -> Builder
101buildBinaryValue (TStruct fields) = buildBinaryStruct fields <> buildType T_STOP
102buildBinaryValue (TMap ky vt entries) =
103  buildType ky <>
104  buildType vt <>
105  int32BE (fromIntegral (length entries)) <>
106  buildBinaryMap entries
107buildBinaryValue (TList ty entries) =
108  buildType ty <>
109  int32BE (fromIntegral (length entries)) <>
110  buildBinaryList entries
111buildBinaryValue (TSet ty entries) =
112  buildType ty <>
113  int32BE (fromIntegral (length entries)) <>
114  buildBinaryList entries
115buildBinaryValue (TBool b) =
116  word8 $ toEnum $ if b then 1 else 0
117buildBinaryValue (TByte b) = int8 b
118buildBinaryValue (TI16 i) = int16BE i
119buildBinaryValue (TI32 i) = int32BE i
120buildBinaryValue (TI64 i) = int64BE i
121buildBinaryValue (TFloat f) = floatBE f
122buildBinaryValue (TDouble d) = doubleBE d
123buildBinaryValue (TString s) = int32BE len <> lazyByteString s
124  where
125    len :: Int32 = fromIntegral (LBS.length s)
126
127buildBinaryStruct :: Map.HashMap Int16 (LT.Text, ThriftVal) -> Builder
128buildBinaryStruct = Map.foldrWithKey combine mempty
129  where
130    combine fid (_,val) s =
131      buildTypeOf val <> int16BE fid <> buildBinaryValue val <> s
132
133buildBinaryMap :: [(ThriftVal, ThriftVal)] -> Builder
134buildBinaryMap = foldl combine mempty
135  where
136    combine s (key, val) = s <> buildBinaryValue key <> buildBinaryValue val
137
138buildBinaryList :: [ThriftVal] -> Builder
139buildBinaryList = foldr (mappend . buildBinaryValue) mempty
140
141-- | Reading Functions
142parseBinaryValue :: ThriftType -> P.Parser ThriftVal
143parseBinaryValue (T_STRUCT _) = TStruct <$> parseBinaryStruct
144parseBinaryValue (T_MAP _ _) = do
145  kt <- parseType
146  vt <- parseType
147  n <- Binary.decode . LBS.fromStrict <$> P.take 4
148  TMap kt vt <$> parseBinaryMap kt vt n
149parseBinaryValue (T_LIST _) = do
150  t <- parseType
151  n <- Binary.decode . LBS.fromStrict <$> P.take 4
152  TList t <$> parseBinaryList t n
153parseBinaryValue (T_SET _) = do
154  t <- parseType
155  n <- Binary.decode . LBS.fromStrict <$> P.take 4
156  TSet t <$> parseBinaryList t n
157parseBinaryValue T_BOOL = TBool . (/=0) <$> P.anyWord8
158parseBinaryValue T_BYTE = TByte . Binary.decode . LBS.fromStrict <$> P.take 1
159parseBinaryValue T_I16 = TI16 . Binary.decode . LBS.fromStrict <$> P.take 2
160parseBinaryValue T_I32 = TI32 . Binary.decode . LBS.fromStrict <$> P.take 4
161parseBinaryValue T_I64 = TI64 . Binary.decode . LBS.fromStrict <$> P.take 8
162parseBinaryValue T_FLOAT = TFloat . bsToFloating byteSwap32 <$> P.take 4
163parseBinaryValue T_DOUBLE = TDouble . bsToFloating byteSwap64 <$> P.take 8
164parseBinaryValue T_STRING = do
165  i :: Int32  <- Binary.decode . LBS.fromStrict <$> P.take 4
166  TString . LBS.fromStrict <$> P.take (fromIntegral i)
167parseBinaryValue ty = error $ "Cannot read value of type " ++ show ty
168
169parseBinaryStruct :: P.Parser (Map.HashMap Int16 (LT.Text, ThriftVal))
170parseBinaryStruct = Map.fromList <$> P.manyTill parseField (matchType T_STOP)
171  where
172    parseField = do
173      t <- parseType
174      n <- Binary.decode . LBS.fromStrict <$> P.take 2
175      v <- parseBinaryValue t
176      return (n, ("", v))
177
178parseBinaryMap :: ThriftType -> ThriftType -> Int32 -> P.Parser [(ThriftVal, ThriftVal)]
179parseBinaryMap kt vt n | n <= 0 = return []
180                       | otherwise = do
181  k <- parseBinaryValue kt
182  v <- parseBinaryValue vt
183  ((k,v) :) <$> parseBinaryMap kt vt (n-1)
184
185parseBinaryList :: ThriftType -> Int32 -> P.Parser [ThriftVal]
186parseBinaryList ty n | n <= 0 = return []
187                     | otherwise = liftM2 (:) (parseBinaryValue ty)
188                                   (parseBinaryList ty (n-1))
189
190
191
192-- | Write a type as a byte
193buildType :: ThriftType -> Builder
194buildType t = word8 $ fromIntegral $ fromEnum t
195
196-- | Write type of a ThriftVal as a byte
197buildTypeOf :: ThriftVal -> Builder
198buildTypeOf v = buildType $ case v of
199  TStruct{} -> T_STRUCT Map.empty
200  TMap{} -> T_MAP T_VOID T_VOID
201  TList{} -> T_LIST T_VOID
202  TSet{} -> T_SET T_VOID
203  TBool{} -> T_BOOL
204  TByte{} -> T_BYTE
205  TI16{} -> T_I16
206  TI32{} -> T_I32
207  TI64{} -> T_I64
208  TString{} -> T_STRING
209  TFloat{} -> T_FLOAT
210  TDouble{} -> T_DOUBLE
211
212-- | Read a byte as though it were a ThriftType
213parseType :: P.Parser ThriftType
214parseType = toEnum . fromIntegral <$> P.anyWord8
215
216matchType :: ThriftType -> P.Parser ThriftType
217matchType t = t <$ P.word8 (fromIntegral $ fromEnum t)
218