ceph/src/jaegertracing/thrift/lib/hs/src/Thrift/Transport/Framed.hs

   1 {-# LANGUAGE FlexibleInstances #-}
   2 {-# LANGUAGE MultiParamTypeClasses #-}
   3 --
   4 -- Licensed to the Apache Software Foundation (ASF) under one
   5 -- or more contributor license agreements. See the NOTICE file
   6 -- distributed with this work for additional information
   7 -- regarding copyright ownership. The ASF licenses this file
   8 -- to you under the Apache License, Version 2.0 (the
   9 -- "License"); you may not use this file except in compliance
  10 -- with the License. You may obtain a copy of the License at
  11 --
  12 --   http://www.apache.org/licenses/LICENSE-2.0
  13 --
  14 -- Unless required by applicable law or agreed to in writing,
  15 -- software distributed under the License is distributed on an
  16 -- "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
  17 -- KIND, either express or implied. See the License for the
  18 -- specific language governing permissions and limitations
  19 -- under the License.
  20 --
  21
  22 module Thrift.Transport.Framed
  23     ( module Thrift.Transport
  24     , FramedTransport
  25     , openFramedTransport
  26     ) where
  27
  28 import Thrift.Transport
  29 import Thrift.Transport.IOBuffer
  30
  31 import Data.Int (Int32)
  32 import qualified Data.Binary as B
  33 import qualified Data.ByteString.Lazy as LBS
  34
  35
  36 -- | FramedTransport wraps a given transport in framed mode.
  37 data FramedTransport t = FramedTransport {
  38       wrappedTrans :: t,           -- ^ Underlying transport.
  39       writeBuffer  :: WriteBuffer, -- ^ Write buffer.
  40       readBuffer   :: ReadBuffer   -- ^ Read buffer.
  41     }
  42
  43 -- | Create a new framed transport which wraps the given transport.
  44 openFramedTransport :: Transport t => t -> IO (FramedTransport t)
  45 openFramedTransport trans = do
  46   wbuf <- newWriteBuffer
  47   rbuf <- newReadBuffer
  48   return FramedTransport{ wrappedTrans = trans, writeBuffer = wbuf, readBuffer = rbuf }
  49
  50 instance Transport t => Transport (FramedTransport t) where
  51
  52     tClose = tClose . wrappedTrans
  53
  54     tRead trans n = do
  55       -- First, check the read buffer for any data.
  56       bs <- readBuf (readBuffer trans) n
  57       if LBS.null bs
  58          then
  59          -- When the buffer is empty, read another frame from the
  60          -- underlying transport.
  61            do len <- readFrame trans
  62               if len > 0
  63                  then tRead trans n
  64                  else return bs
  65          else return bs
  66     tPeek trans = do
  67       mw <- peekBuf (readBuffer trans)
  68       case mw of
  69         Just _ -> return mw
  70         Nothing -> do
  71           len <- readFrame trans
  72           if len > 0
  73              then tPeek trans
  74              else return Nothing
  75
  76     tWrite = writeBuf . writeBuffer
  77
  78     tFlush trans = do
  79       bs <- flushBuf (writeBuffer trans)
  80       let szBs = B.encode $ (fromIntegral $ LBS.length bs :: Int32)
  81       tWrite (wrappedTrans trans) szBs
  82       tWrite (wrappedTrans trans) bs
  83       tFlush (wrappedTrans trans)
  84
  85     tIsOpen = tIsOpen . wrappedTrans
  86
  87 readFrame :: Transport t => FramedTransport t -> IO Int
  88 readFrame trans = do
  89   -- Read and decode the frame size.
  90   szBs <- tRead (wrappedTrans trans) 4
  91   let sz = fromIntegral (B.decode szBs :: Int32)
  92
  93   -- Read the frame and stuff it into the read buffer.
  94   bs <- tRead (wrappedTrans trans) sz
  95   fillBuf (readBuffer trans) bs
  96
  97   -- Return the frame size so that the caller knows whether to expect
  98   -- something in the read buffer or not.
  99   return sz