StreamingSpec.hs

{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE CPP #-}
module Data.Store.StreamingSpec where

import           Control.Concurrent (threadDelay)
import           Control.Concurrent.Async (race, concurrently)
import           Control.Concurrent.MVar
import           Control.Exception (try)
import           Control.Monad (void, (<=<), forM_)
import           Control.Monad.Trans.Free (runFreeT, FreeF(..))
import           Control.Monad.Trans.Free.Church (fromFT)
import           Control.Monad.Trans.Resource
import qualified Data.ByteString as BS
import           Data.Conduit ((=$=), ($$))
import qualified Data.Conduit.List as C
import           Data.List (unfoldr)
import           Data.Monoid
import           Data.Store.Internal
import           Data.Store.Streaming
import           Data.Store.Streaming.Internal
import           Data.Streaming.Network (runTCPServer, runTCPClient, clientSettingsTCP, serverSettingsTCP, setAfterBind)
import           Data.Streaming.Network.Internal (AppData(..))
import           Data.Void (absurd, Void)
import           Network.Socket (Socket(..), socketPort)
import           Network.Socket.ByteString (send)
import qualified System.IO.ByteBuffer as BB
import           System.Posix.Types (Fd(..))
import           Test.Hspec
import           Test.Hspec.SmallCheck
import           Test.SmallCheck

spec :: Spec
spec = do
  describe "conduitEncode and conduitDecode" $ do
    it "Roundtrips ([Int])." $ property roundtrip
    it "Roundtrips ([Int]), with chunked transfer." $ property roundtripChunked
    it "Throws an Exception on incomplete messages." conduitIncomplete
    it "Throws an Exception on excess input." $ property conduitExcess
  describe "peekMessage" $ do
    describe "ByteString" $ do
      it "demands more input when needed." $ property (askMoreBS (headerLength + 1))
      it "demands more input on incomplete message magic." $ property (askMoreBS 1)
      it "demands more input on incomplete SizeTag." $ property (askMoreBS (headerLength - 1))
      it "successfully decodes valid input." $ property canPeekBS
  describe "decodeMessage" $ do
    describe "ByteString" $ do
      it "Throws an Exception on incomplete messages." decodeIncomplete
      it "Throws an Exception on messages that are shorter than indicated." decodeTooShort
#ifndef mingw32_HOST_OS
    describe "Socket" $ do
      it "Decodes data trickling through a socket." $ property decodeTricklingMessageFd
#endif

roundtrip :: [Int] -> Property IO
roundtrip xs = monadic $ do
  xs' <- runResourceT $ C.sourceList xs
    =$= C.map Message
    =$= conduitEncode
    =$= conduitDecode Nothing
    =$= C.map fromMessage
    $$ C.consume
  return $ xs' == xs

roundtripChunked :: [Int] -> Property IO
roundtripChunked input = monadic $ do
  let (xs, chunkLengths) = splitAt (length input `div` 2) input
  bs <- C.sourceList xs
    =$= C.map Message
    =$= conduitEncode
    $$ C.fold (<>) mempty
  let chunks = unfoldr takeChunk (bs, chunkLengths)
  xs' <- runResourceT $ C.sourceList chunks
    =$= conduitDecode (Just 10)
    =$= C.map fromMessage
    $$ C.consume
  return $ xs' == xs
  where
    takeChunk (x, _) | BS.null x = Nothing
    takeChunk (x, []) = Just (x, (BS.empty, []))
    takeChunk (x, l:ls) =
        let (chunk, rest) = BS.splitAt l x
        in Just (chunk, (rest, ls))

conduitIncomplete :: Expectation
conduitIncomplete =
    (runResourceT (C.sourceList [incompleteInput]
                  =$= conduitDecode (Just 10)
                  $$ C.consume)
    :: IO [Message Integer]) `shouldThrow` \PeekException{} -> True

conduitExcess :: [Int] -> Property IO
conduitExcess xs = monadic $ do
  bs <- C.sourceList xs
    =$= C.map Message
    =$= conduitEncode
    $$ C.fold (<>) mempty
  res <- try (runResourceT (C.sourceList [bs `BS.append` "excess bytes"]
                            =$= conduitDecode (Just 10)
                            $$ C.consume) :: IO [Message Int])
  case res of
      Right _ -> return False
      Left (PeekException _ _) -> return True

-- splits an encoded message after n bytes.  Feeds the first part to
-- peekResult, expecting it to require more input.  Then, feeds the
-- second part and checks if the decoded result is the original
-- message.
askMoreBS :: Int -> Integer -> Property IO
askMoreBS n x = monadic $ BB.with (Just 10) $ \ bb -> do
  let bs = encodeMessage (Message x)
      (start, end) = BS.splitAt n $ bs
  BB.copyByteString bb start
  peekResult <- runFreeT (fromFT (peekMessageBS bb))
  case peekResult of
    Free cont ->
      runFreeT (cont end) >>= \case
        Pure (Message x') -> return $ x' == x
        Free _ -> return False
    Pure _ -> return False

canPeekBS :: Integer -> Property IO
canPeekBS x = monadic $ BB.with (Just 10) $ \ bb -> do
  let bs = encodeMessage (Message x)
  BB.copyByteString bb bs
  peekResult <- runFreeT (fromFT (peekMessageBS bb))
  case peekResult of
    Free _ -> return False
    Pure (Message x') -> return $ x' == x

#ifndef mingw32_HOST_OS

socketFd :: Socket -> Fd
socketFd (MkSocket fd _ _ _ _) = Fd fd

withServer :: (Socket -> Socket -> IO a) -> IO a
withServer cont = do
  sock1Var :: MVar Socket <- newEmptyMVar
  sock2Var :: MVar Socket <- newEmptyMVar
  portVar :: MVar Int <- newEmptyMVar
  doneVar :: MVar Void <- newEmptyMVar
  let adSocket ad = case appRawSocket' ad of
        Nothing -> error "withServer.adSocket: no raw socket in AppData"
        Just sock -> sock
  let ss = setAfterBind
        (putMVar portVar . fromIntegral <=< socketPort)
        (serverSettingsTCP 0 "127.0.0.1")
  x <- fmap (either (either absurd absurd) id) $ race
    (race
      (runTCPServer ss $ \ad -> do
        putMVar sock1Var (adSocket ad)
        void (readMVar doneVar))
      (do port <- takeMVar portVar
          runTCPClient (clientSettingsTCP port "127.0.0.1") $ \ad -> do
            putMVar sock2Var (adSocket ad)
            readMVar doneVar))
    (do sock1 <- takeMVar sock1Var
        sock2 <- takeMVar sock2Var
        cont sock1 sock2)
  putMVar doneVar (error "withServer: impossible: read from doneVar")
  return x

decodeTricklingMessageFd :: Integer -> Property IO
decodeTricklingMessageFd v = monadic $ do
  let bs = encodeMessage (Message v)
  BB.with Nothing $ \bb ->
    withServer $ \sock1 sock2 -> do
      let generateChunks :: [Int] -> BS.ByteString -> [BS.ByteString]
          generateChunks xs0 bs_ = case xs0 of
            [] -> generateChunks [1,3,10] bs_
            x : xs -> if BS.null bs_
              then []
              else BS.take x bs_ : generateChunks xs (BS.drop x bs_)
      let chunks = generateChunks [] bs
      ((), Message v') <- concurrently
        (forM_ chunks $ \chunk -> do
          void (send sock1 chunk)
          threadDelay (10 * 1000))
        (decodeMessageFd bb (socketFd sock2))
      return (v == v')

#endif

decodeIncomplete :: IO ()
decodeIncomplete = BB.with (Just 0) $ \ bb -> do
  BB.copyByteString bb (BS.take 1 incompleteInput)
  (decodeMessageBS bb (return Nothing) :: IO (Maybe (Message Integer)))
    `shouldThrow` \PeekException{} -> True

incompleteInput :: BS.ByteString
incompleteInput =
  let bs = encodeMessage (Message (42 :: Integer))
  in BS.take (BS.length bs - 1) bs

decodeTooShort :: IO ()
decodeTooShort = BB.with Nothing $ \bb -> do
    BB.copyByteString bb (encodeMessageTooShort . Message $ (1 :: Int))
    (decodeMessageBS bb (return Nothing) :: IO (Maybe (Message Int)))
        `shouldThrow` \PeekException{} -> True

encodeMessageTooShort :: Message Int -> BS.ByteString
encodeMessageTooShort msg =
    BS.take (BS.length encoded - (getSize (0 :: Int))) encoded
  where
    encoded = encodeMessage msg