1 files changed, 99 insertions, 0 deletions
diff --git a/src/jaegertracing/thrift/lib/hs/src/Thrift/Transport/Framed.hs b/src/jaegertracing/thrift/lib/hs/src/Thrift/Transport/Framed.hs
new file mode 100644
index 000000000..42fc43f39
--- /dev/null
+++ b/src/jaegertracing/thrift/lib/hs/src/Thrift/Transport/Framed.hs
@@ -0,0 +1,99 @@
+{-# LANGUAGE FlexibleInstances #-}
+{-# LANGUAGE MultiParamTypeClasses #-}
+--
+-- Licensed to the Apache Software Foundation (ASF) under one
+-- or more contributor license agreements. See the NOTICE file
+-- distributed with this work for additional information
+-- regarding copyright ownership. The ASF licenses this file
+-- to you under the Apache License, Version 2.0 (the
+-- "License"); you may not use this file except in compliance
+-- with the License. You may obtain a copy of the License at
+--
+--   http://www.apache.org/licenses/LICENSE-2.0
+--
+-- Unless required by applicable law or agreed to in writing,
+-- software distributed under the License is distributed on an
+-- "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+-- KIND, either express or implied. See the License for the
+-- specific language governing permissions and limitations
+-- under the License.
+--
+
+module Thrift.Transport.Framed
+    ( module Thrift.Transport
+    , FramedTransport
+    , openFramedTransport
+    ) where
+
+import Thrift.Transport
+import Thrift.Transport.IOBuffer
+
+import Data.Int (Int32)
+import qualified Data.Binary as B
+import qualified Data.ByteString.Lazy as LBS
+
+
+-- | FramedTransport wraps a given transport in framed mode.
+data FramedTransport t = FramedTransport {
+      wrappedTrans :: t,           -- ^ Underlying transport.
+      writeBuffer  :: WriteBuffer, -- ^ Write buffer.
+      readBuffer   :: ReadBuffer   -- ^ Read buffer.
+    }
+
+-- | Create a new framed transport which wraps the given transport.
+openFramedTransport :: Transport t => t -> IO (FramedTransport t)
+openFramedTransport trans = do
+  wbuf <- newWriteBuffer
+  rbuf <- newReadBuffer
+  return FramedTransport{ wrappedTrans = trans, writeBuffer = wbuf, readBuffer = rbuf }
+
+instance Transport t => Transport (FramedTransport t) where
+
+    tClose = tClose . wrappedTrans
+
+    tRead trans n = do
+      -- First, check the read buffer for any data.
+      bs <- readBuf (readBuffer trans) n
+      if LBS.null bs
+         then
+         -- When the buffer is empty, read another frame from the
+         -- underlying transport.
+           do len <- readFrame trans
+              if len > 0
+                 then tRead trans n
+                 else return bs
+         else return bs
+    tPeek trans = do
+      mw <- peekBuf (readBuffer trans)
+      case mw of
+        Just _ -> return mw
+        Nothing -> do
+          len <- readFrame trans
+          if len > 0
+             then tPeek trans
+             else return Nothing
+
+    tWrite = writeBuf . writeBuffer
+
+    tFlush trans = do
+      bs <- flushBuf (writeBuffer trans)
+      let szBs = B.encode $ (fromIntegral $ LBS.length bs :: Int32)
+      tWrite (wrappedTrans trans) szBs
+      tWrite (wrappedTrans trans) bs
+      tFlush (wrappedTrans trans)
+
+    tIsOpen = tIsOpen . wrappedTrans
+
+readFrame :: Transport t => FramedTransport t -> IO Int
+readFrame trans = do
+  -- Read and decode the frame size.
+  szBs <- tRead (wrappedTrans trans) 4
+  let sz = fromIntegral (B.decode szBs :: Int32)
+
+  -- Read the frame and stuff it into the read buffer.
+  bs <- tRead (wrappedTrans trans) sz
+  fillBuf (readBuffer trans) bs
+
+  -- Return the frame size so that the caller knows whether to expect
+  -- something in the read buffer or not.
+  return sz