ceph/src/jaegertracing/thrift/lib/hs/src/Thrift/Protocol/Compact.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 {-# LANGUAGE CPP #-}
  21 {-# LANGUAGE ExistentialQuantification #-}
  22 {-# LANGUAGE OverloadedStrings #-}
  23 {-# LANGUAGE ScopedTypeVariables #-}
  24
  25 module Thrift.Protocol.Compact
  26     ( module Thrift.Protocol
  27     , CompactProtocol(..)
  28     , parseVarint
  29     , buildVarint
  30     ) where
  31
  32 import Control.Applicative
  33 import Control.Monad
  34 import Data.Attoparsec.ByteString as P
  35 import Data.Attoparsec.ByteString.Lazy as LP
  36 import Data.Bits
  37 import Data.ByteString.Lazy.Builder as B
  38 import Data.Int
  39 import Data.List as List
  40 import Data.Monoid
  41 import Data.Word
  42 import Data.Text.Lazy.Encoding ( decodeUtf8, encodeUtf8 )
  43
  44 import Thrift.Protocol
  45 import Thrift.Transport
  46 import Thrift.Types
  47
  48 import qualified Data.ByteString as BS
  49 import qualified Data.ByteString.Lazy as LBS
  50 import qualified Data.HashMap.Strict as Map
  51 import qualified Data.Text.Lazy as LT
  52
  53 -- | the Compact Protocol implements the standard Thrift 'TCompactProcotol'
  54 -- which is similar to the 'TBinaryProtocol', but takes less space on the wire.
  55 -- Integral types are encoded using as varints.
  56 data CompactProtocol a = CompactProtocol a
  57                          -- ^ Constuct a 'CompactProtocol' with a 'Transport'
  58
  59 protocolID, version, versionMask, typeMask, typeBits :: Word8
  60 protocolID  = 0x82 -- 1000 0010
  61 version     = 0x01
  62 versionMask = 0x1f -- 0001 1111
  63 typeMask    = 0xe0 -- 1110 0000
  64 typeBits    = 0x07 -- 0000 0111
  65 typeShiftAmount :: Int
  66 typeShiftAmount = 5
  67
  68 getTransport :: Transport t => CompactProtocol t -> t
  69 getTransport (CompactProtocol t) = t
  70
  71 instance Transport t => Protocol (CompactProtocol t) where
  72     readByte p = tReadAll (getTransport p) 1
  73     writeMessage p (n, t, s) f = do
  74       tWrite (getTransport p) messageBegin
  75       f
  76       tFlush $ getTransport p
  77       where
  78         messageBegin = toLazyByteString $
  79           B.word8 protocolID <>
  80           B.word8 ((version .&. versionMask) .|.
  81                   (((fromIntegral $ fromEnum t) `shiftL`
  82                   typeShiftAmount) .&. typeMask)) <>
  83           buildVarint (i32ToZigZag s) <>
  84           buildCompactValue (TString $ encodeUtf8 n)
  85
  86     readMessage p f = readMessageBegin >>= f
  87       where
  88         readMessageBegin = runParser p $ do
  89           pid <- fromIntegral <$> P.anyWord8
  90           when (pid /= protocolID) $ error "Bad Protocol ID"
  91           w <- fromIntegral <$> P.anyWord8
  92           let ver = w .&. versionMask
  93           when (ver /= version) $ error "Bad Protocol version"
  94           let typ = (w `shiftR` typeShiftAmount) .&. typeBits
  95           seqId <- parseVarint zigZagToI32
  96           TString name <- parseCompactValue T_STRING
  97           return (decodeUtf8 name, toEnum $ fromIntegral $ typ, seqId)
  98
  99     writeVal p = tWrite (getTransport p) . toLazyByteString . buildCompactValue
 100     readVal p ty = runParser p $ parseCompactValue ty
 101
 102 instance Transport t => StatelessProtocol (CompactProtocol t) where
 103     serializeVal _ = toLazyByteString . buildCompactValue
 104     deserializeVal _ ty bs =
 105       case LP.eitherResult $ LP.parse (parseCompactValue ty) bs of
 106         Left s -> error s
 107         Right val -> val
 108
 109 -- | Writing Functions
 110 buildCompactValue :: ThriftVal -> Builder
 111 buildCompactValue (TStruct fields) = buildCompactStruct fields
 112 buildCompactValue (TMap kt vt entries) =
 113   let len = fromIntegral $ length entries :: Word32 in
 114   if len == 0
 115   then B.word8 0x00
 116   else buildVarint len <>
 117        B.word8 (fromTType kt `shiftL` 4 .|. fromTType vt) <>
 118        buildCompactMap entries
 119 buildCompactValue (TList ty entries) =
 120   let len = length entries in
 121   (if len < 15
 122    then B.word8 $ (fromIntegral len `shiftL` 4) .|. fromTType ty
 123    else B.word8 (0xF0 .|. fromTType ty) <>
 124         buildVarint (fromIntegral len :: Word32)) <>
 125   buildCompactList entries
 126 buildCompactValue (TSet ty entries) = buildCompactValue (TList ty entries)
 127 buildCompactValue (TBool b) =
 128   B.word8 $ toEnum $ if b then 1 else 0
 129 buildCompactValue (TByte b) = int8 b
 130 buildCompactValue (TI16 i) = buildVarint $ i16ToZigZag i
 131 buildCompactValue (TI32 i) = buildVarint $ i32ToZigZag i
 132 buildCompactValue (TI64 i) = buildVarint $ i64ToZigZag i
 133 buildCompactValue (TDouble d) = doubleLE d
 134 buildCompactValue (TString s) = buildVarint len <> lazyByteString s
 135   where
 136     len = fromIntegral (LBS.length s) :: Word32
 137 buildCompactValue (TBinary s) = buildCompactValue (TString s)
 138
 139 buildCompactStruct :: Map.HashMap Int16 (LT.Text, ThriftVal) -> Builder
 140 buildCompactStruct = flip (loop 0) mempty . Map.toList
 141   where
 142     loop _ [] acc = acc <> B.word8 (fromTType T_STOP)
 143     loop lastId ((fid, (_,val)) : fields) acc = loop fid fields $ acc <>
 144       (if fid > lastId && fid - lastId <= 15
 145        then B.word8 $ fromIntegral ((fid - lastId) `shiftL` 4) .|. typeOf val
 146        else B.word8 (typeOf val) <> buildVarint (i16ToZigZag fid)) <>
 147       (if typeOf val > 0x02 -- Not a T_BOOL
 148        then buildCompactValue val
 149        else mempty) -- T_BOOLs are encoded in the type
 150 buildCompactMap :: [(ThriftVal, ThriftVal)] -> Builder
 151 buildCompactMap = foldl combine mempty
 152   where
 153     combine s (key, val) = buildCompactValue key <> buildCompactValue val <> s
 154
 155 buildCompactList :: [ThriftVal] -> Builder
 156 buildCompactList = foldr (mappend . buildCompactValue) mempty
 157
 158 -- | Reading Functions
 159 parseCompactValue :: ThriftType -> Parser ThriftVal
 160 parseCompactValue (T_STRUCT tmap) = TStruct <$> parseCompactStruct tmap
 161 parseCompactValue (T_MAP kt' vt') = do
 162   n <- parseVarint id
 163   if n == 0
 164     then return $ TMap kt' vt' []
 165     else do
 166     w <- P.anyWord8
 167     let kt = typeFrom $ w `shiftR` 4
 168         vt = typeFrom $ w .&. 0x0F
 169     TMap kt vt <$> parseCompactMap kt vt n
 170 parseCompactValue (T_LIST ty) = TList ty <$> parseCompactList
 171 parseCompactValue (T_SET ty) = TSet ty <$> parseCompactList
 172 parseCompactValue T_BOOL = TBool . (/=0) <$> P.anyWord8
 173 parseCompactValue T_BYTE = TByte . fromIntegral <$> P.anyWord8
 174 parseCompactValue T_I16 = TI16 <$> parseVarint zigZagToI16
 175 parseCompactValue T_I32 = TI32 <$> parseVarint zigZagToI32
 176 parseCompactValue T_I64 = TI64 <$> parseVarint zigZagToI64
 177 parseCompactValue T_DOUBLE = TDouble . bsToDoubleLE <$> P.take 8
 178 parseCompactValue T_STRING = parseCompactString TString
 179 parseCompactValue T_BINARY = parseCompactString TBinary
 180 parseCompactValue ty = error $ "Cannot read value of type " ++ show ty
 181
 182 parseCompactString ty = do
 183   len :: Word32 <- parseVarint id
 184   ty . LBS.fromStrict <$> P.take (fromIntegral len)
 185
 186 parseCompactStruct :: TypeMap -> Parser (Map.HashMap Int16 (LT.Text, ThriftVal))
 187 parseCompactStruct tmap = Map.fromList <$> parseFields 0
 188   where
 189     parseFields :: Int16 -> Parser [(Int16, (LT.Text, ThriftVal))]
 190     parseFields lastId = do
 191       w <- P.anyWord8
 192       if w == 0x00
 193         then return []
 194         else do
 195           let ty = typeFrom (w .&. 0x0F)
 196               modifier = (w .&. 0xF0) `shiftR` 4
 197           fid <- if modifier /= 0
 198                  then return (lastId + fromIntegral modifier)
 199                  else parseVarint zigZagToI16
 200           val <- if ty == T_BOOL
 201                  then return (TBool $ (w .&. 0x0F) == 0x01)
 202                  else case (ty, Map.lookup fid tmap) of
 203                         (T_STRING, Just (_, T_BINARY)) -> parseCompactValue T_BINARY
 204                         _ -> parseCompactValue ty
 205           ((fid, (LT.empty, val)) : ) <$> parseFields fid
 206
 207 parseCompactMap :: ThriftType -> ThriftType -> Int32 ->
 208                    Parser [(ThriftVal, ThriftVal)]
 209 parseCompactMap kt vt n | n <= 0 = return []
 210                         | otherwise = do
 211   k <- parseCompactValue kt
 212   v <- parseCompactValue vt
 213   ((k,v) :) <$> parseCompactMap kt vt (n-1)
 214
 215 parseCompactList :: Parser [ThriftVal]
 216 parseCompactList = do
 217   w <- P.anyWord8
 218   let ty = typeFrom $ w .&. 0x0F
 219       lsize = w `shiftR` 4
 220   size <- if lsize == 0xF
 221           then parseVarint id
 222           else return $ fromIntegral lsize
 223   loop ty size
 224   where
 225     loop :: ThriftType -> Int32 -> Parser [ThriftVal]
 226     loop ty n | n <= 0 = return []
 227               | otherwise = liftM2 (:) (parseCompactValue ty)
 228                             (loop ty (n-1))
 229
 230 -- Signed numbers must be converted to "Zig Zag" format before they can be
 231 -- serialized in the Varint format
 232 i16ToZigZag :: Int16 -> Word16
 233 i16ToZigZag n = fromIntegral $ (n `shiftL` 1) `xor` (n `shiftR` 15)
 234
 235 zigZagToI16 :: Word16 -> Int16
 236 zigZagToI16 n = fromIntegral $ (n `shiftR` 1) `xor` negate (n .&. 0x1)
 237
 238 i32ToZigZag :: Int32 -> Word32
 239 i32ToZigZag n = fromIntegral $ (n `shiftL` 1) `xor` (n `shiftR` 31)
 240
 241 zigZagToI32 :: Word32 -> Int32
 242 zigZagToI32 n = fromIntegral $ (n `shiftR` 1) `xor` negate (n .&. 0x1)
 243
 244 i64ToZigZag :: Int64 -> Word64
 245 i64ToZigZag n = fromIntegral $ (n `shiftL` 1) `xor` (n `shiftR` 63)
 246
 247 zigZagToI64 :: Word64 -> Int64
 248 zigZagToI64 n = fromIntegral $ (n `shiftR` 1) `xor` negate (n .&. 0x1)
 249
 250 buildVarint :: (Bits a, Integral a)  => a -> Builder
 251 buildVarint n | n .&. complement 0x7F == 0 = B.word8 $ fromIntegral n
 252               | otherwise = B.word8 (0x80 .|. (fromIntegral n .&. 0x7F)) <>
 253                             buildVarint (n `shiftR` 7)
 254
 255 parseVarint :: (Bits a, Integral a, Ord a) => (a -> b) -> Parser b
 256 parseVarint fromZigZag = do
 257   bytestemp <- BS.unpack <$> P.takeTill (not . flip testBit 7)
 258   lsb <- P.anyWord8
 259   let bytes = lsb : List.reverse bytestemp
 260   return $ fromZigZag $ List.foldl' combine 0x00 bytes
 261   where combine a b = (a `shiftL` 7) .|. (fromIntegral b .&. 0x7f)
 262
 263 -- | Compute the Compact Type
 264 fromTType :: ThriftType -> Word8
 265 fromTType ty = case ty of
 266   T_STOP -> 0x00
 267   T_BOOL -> 0x01
 268   T_BYTE -> 0x03
 269   T_I16 -> 0x04
 270   T_I32 -> 0x05
 271   T_I64 -> 0x06
 272   T_DOUBLE -> 0x07
 273   T_STRING -> 0x08
 274   T_BINARY -> 0x08
 275   T_LIST{} -> 0x09
 276   T_SET{} -> 0x0A
 277   T_MAP{} -> 0x0B
 278   T_STRUCT{} -> 0x0C
 279   T_VOID -> error "No Compact type for T_VOID"
 280
 281 typeOf :: ThriftVal -> Word8
 282 typeOf v = case v of
 283   TBool True -> 0x01
 284   TBool False -> 0x02
 285   TByte _ -> 0x03
 286   TI16 _ -> 0x04
 287   TI32 _ -> 0x05
 288   TI64 _ -> 0x06
 289   TDouble _ -> 0x07
 290   TString _ -> 0x08
 291   TBinary _ -> 0x08
 292   TList{} -> 0x09
 293   TSet{} -> 0x0A
 294   TMap{} -> 0x0B
 295   TStruct{} -> 0x0C
 296
 297 typeFrom :: Word8 -> ThriftType
 298 typeFrom w = case w of
 299   0x01 -> T_BOOL
 300   0x02 -> T_BOOL
 301   0x03 -> T_BYTE
 302   0x04 -> T_I16
 303   0x05 -> T_I32
 304   0x06 -> T_I64
 305   0x07 -> T_DOUBLE
 306   0x08 -> T_STRING
 307   0x09 -> T_LIST T_VOID
 308   0x0A -> T_SET T_VOID
 309   0x0B -> T_MAP T_VOID T_VOID
 310   0x0C -> T_STRUCT Map.empty
 311   n -> error $ "typeFrom: " ++ show n ++ " is not a compact type"