{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE RecordWildCards #-}
module Database.Bolt.Connection.Pipe where

import qualified Database.Bolt.Connection.Connection as C (close, connect, recv, send, sendMany)
import           Database.Bolt.Connection.Instances
import           Database.Bolt.Connection.Type
import           Database.Bolt.Value.Helpers
import           Database.Bolt.Value.Instances
import           Database.Bolt.Value.Type            (BoltValue (pack, unpackT),
                                                      FromStructure (fromStructure),
                                                      ToStructure (toStructure), unpackAction)

import           Control.Exception                   (throwIO)
import           Control.Monad                       (forM_, unless, void, when)
import           Control.Monad.Except                (ExceptT, MonadError (..), runExceptT)
import           Control.Monad.Trans                 (MonadIO (..))
import           Data.Binary                         (decode)
import           Data.Binary.Put                     (runPut)
import           Data.ByteString                     (ByteString)
import qualified Data.ByteString                     as B (concat, length)
import qualified Data.ByteString.Lazy                as BSL
import qualified Data.ByteString.Lazy.Internal       as BSL
import           Data.Int                            (Int64)
import           Data.Word                           (Word16)
import           GHC.Stack                           (HasCallStack)

type MonadPipe m = (MonadIO m, MonadError BoltError m)

-- |Creates new 'Pipe' instance to use all requests through
connect :: MonadIO m => HasCallStack => BoltCfg -> m Pipe
connect :: BoltCfg -> m Pipe
connect = (BoltCfg -> ExceptT BoltError m Pipe) -> BoltCfg -> m Pipe
forall (m :: * -> *) a b.
(MonadIO m, HasCallStack) =>
(a -> ExceptT BoltError m b) -> a -> m b
makeIO BoltCfg -> ExceptT BoltError m Pipe
forall (m :: * -> *). MonadPipe m => BoltCfg -> m Pipe
connect'
  where
    connect' :: MonadPipe m => BoltCfg -> m Pipe
    connect' :: BoltCfg -> m Pipe
connect' BoltCfg
bcfg = do ConnectionWithTimeout
conn <- Bool -> String -> PortNumber -> Int -> m ConnectionWithTimeout
forall (m :: * -> *).
(MonadIO m, HasCallStack) =>
Bool -> String -> PortNumber -> Int -> m ConnectionWithTimeout
C.connect (BoltCfg -> Bool
secure BoltCfg
bcfg) (BoltCfg -> String
host BoltCfg
bcfg) (Int -> PortNumber
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> PortNumber) -> Int -> PortNumber
forall a b. (a -> b) -> a -> b
$ BoltCfg -> Int
port BoltCfg
bcfg) (BoltCfg -> Int
socketTimeout BoltCfg
bcfg)
                       let pipe :: Pipe
pipe = ConnectionWithTimeout -> Word16 -> Word32 -> Pipe
Pipe ConnectionWithTimeout
conn (BoltCfg -> Word16
maxChunkSize BoltCfg
bcfg) (BoltCfg -> Word32
version BoltCfg
bcfg)
                       Pipe -> BoltCfg -> m ()
forall (m :: * -> *).
(MonadPipe m, HasCallStack) =>
Pipe -> BoltCfg -> m ()
handshake Pipe
pipe BoltCfg
bcfg
                       Pipe -> m Pipe
forall (f :: * -> *) a. Applicative f => a -> f a
pure Pipe
pipe

-- |Closes 'Pipe'
close :: MonadIO m => HasCallStack => Pipe -> m ()
close :: Pipe -> m ()
close Pipe
pipe = do Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Word32 -> Bool
isNewVersion (Word32 -> Bool) -> Word32 -> Bool
forall a b. (a -> b) -> a -> b
$ Pipe -> Word32
pipe_version Pipe
pipe) (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ (Pipe -> ExceptT BoltError m ()) -> Pipe -> m ()
forall (m :: * -> *) a b.
(MonadIO m, HasCallStack) =>
(a -> ExceptT BoltError m b) -> a -> m b
makeIO (Pipe -> Request -> ExceptT BoltError m ()
forall (m :: * -> *).
(MonadPipe m, HasCallStack) =>
Pipe -> Request -> m ()
`flush` Request
RequestGoodbye) Pipe
pipe
                ConnectionWithTimeout -> m ()
forall (m :: * -> *).
(MonadIO m, HasCallStack) =>
ConnectionWithTimeout -> m ()
C.close (ConnectionWithTimeout -> m ()) -> ConnectionWithTimeout -> m ()
forall a b. (a -> b) -> a -> b
$ Pipe -> ConnectionWithTimeout
connection Pipe
pipe

-- |Resets current sessions
reset :: MonadIO m => HasCallStack => Pipe -> m ()
reset :: Pipe -> m ()
reset = (Pipe -> ExceptT BoltError m ()) -> Pipe -> m ()
forall (m :: * -> *) a b.
(MonadIO m, HasCallStack) =>
(a -> ExceptT BoltError m b) -> a -> m b
makeIO Pipe -> ExceptT BoltError m ()
forall (m :: * -> *). MonadPipe m => Pipe -> m ()
reset'
  where
    reset' :: MonadPipe m => Pipe -> m ()
    reset' :: Pipe -> m ()
reset' Pipe
pipe = do Pipe -> Request -> m ()
forall (m :: * -> *).
(MonadPipe m, HasCallStack) =>
Pipe -> Request -> m ()
flush Pipe
pipe Request
RequestReset
                     Response
response <- Pipe -> m Response
forall (m :: * -> *).
(MonadPipe m, HasCallStack) =>
Pipe -> m Response
fetch Pipe
pipe
                     Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Response -> Bool
isFailure Response
response) (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$
                       BoltError -> m ()
forall e (m :: * -> *) a. MonadError e m => e -> m a
throwError BoltError
ResetFailed

-- Helper to make pipe operations in IO
makeIO :: MonadIO m => HasCallStack => (a -> ExceptT BoltError m b) -> a -> m b
makeIO :: (a -> ExceptT BoltError m b) -> a -> m b
makeIO a -> ExceptT BoltError m b
action a
arg = do Either BoltError b
actionIO <- ExceptT BoltError m b -> m (Either BoltError b)
forall e (m :: * -> *) a. ExceptT e m a -> m (Either e a)
runExceptT (a -> ExceptT BoltError m b
action a
arg)
                       case Either BoltError b
actionIO of
                         Right b
x -> b -> m b
forall (f :: * -> *) a. Applicative f => a -> f a
pure b
x
                         Left  BoltError
e -> IO b -> m b
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO b -> m b) -> IO b -> m b
forall a b. (a -> b) -> a -> b
$ BoltError -> IO b
forall e a. Exception e => e -> IO a
throwIO BoltError
e

-- = Internal interfaces

-- |Processes error via ackFailure or reset
processError :: MonadIO m => HasCallStack => Pipe -> m ()
processError :: Pipe -> m ()
processError pipe :: Pipe
pipe@Pipe{Word16
Word32
ConnectionWithTimeout
mcs :: Pipe -> Word16
pipe_version :: Word32
mcs :: Word16
connection :: ConnectionWithTimeout
connection :: Pipe -> ConnectionWithTimeout
pipe_version :: Pipe -> Word32
..} = if Word32 -> Bool
isNewVersion Word32
pipe_version
                               then Pipe -> m ()
forall (m :: * -> *). (MonadIO m, HasCallStack) => Pipe -> m ()
reset Pipe
pipe
                               else (Pipe -> ExceptT BoltError m ()) -> Pipe -> m ()
forall (m :: * -> *) a b.
(MonadIO m, HasCallStack) =>
(a -> ExceptT BoltError m b) -> a -> m b
makeIO Pipe -> ExceptT BoltError m ()
forall (m :: * -> *). (MonadPipe m, HasCallStack) => Pipe -> m ()
ackFailure Pipe
pipe

ackFailure :: MonadPipe m => HasCallStack => Pipe -> m ()
ackFailure :: Pipe -> m ()
ackFailure Pipe
pipe = Pipe -> Request -> m ()
forall (m :: * -> *).
(MonadPipe m, HasCallStack) =>
Pipe -> Request -> m ()
flush Pipe
pipe Request
RequestAckFailure m () -> m () -> m ()
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> m Response -> m ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (Pipe -> m Response
forall (m :: * -> *).
(MonadPipe m, HasCallStack) =>
Pipe -> m Response
fetch Pipe
pipe)

discardAll :: MonadPipe m => HasCallStack => Pipe -> m ()
discardAll :: Pipe -> m ()
discardAll Pipe
pipe = Pipe -> Request -> m ()
forall (m :: * -> *).
(MonadPipe m, HasCallStack) =>
Pipe -> Request -> m ()
flush Pipe
pipe Request
RequestDiscardAll m () -> m () -> m ()
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> m Response -> m ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (Pipe -> m Response
forall (m :: * -> *).
(MonadPipe m, HasCallStack) =>
Pipe -> m Response
fetch Pipe
pipe)

flush :: MonadPipe m => HasCallStack => Pipe -> Request -> m ()
flush :: Pipe -> Request -> m ()
flush Pipe
pipe Request
request = do [ByteString] -> (ByteString -> m ()) -> m ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [ByteString]
chunks ((ByteString -> m ()) -> m ()) -> (ByteString -> m ()) -> m ()
forall a b. (a -> b) -> a -> b
$ ConnectionWithTimeout -> [ByteString] -> m ()
forall (m :: * -> *).
(MonadIO m, HasCallStack) =>
ConnectionWithTimeout -> [ByteString] -> m ()
C.sendMany ConnectionWithTimeout
conn ([ByteString] -> m ())
-> (ByteString -> [ByteString]) -> ByteString -> m ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> [ByteString]
mkChunk
                        ConnectionWithTimeout -> ByteString -> m ()
forall (m :: * -> *).
(MonadIO m, HasCallStack) =>
ConnectionWithTimeout -> ByteString -> m ()
C.send ConnectionWithTimeout
conn ByteString
terminal
  where bs :: ByteString
bs        = Put -> ByteString
runPut (Put -> ByteString) -> Put -> ByteString
forall a b. (a -> b) -> a -> b
$ Structure -> Put
forall a. BoltValue a => a -> Put
pack (Structure -> Put) -> Structure -> Put
forall a b. (a -> b) -> a -> b
$ Request -> Structure
forall a. ToStructure a => a -> Structure
toStructure Request
request
        chunkSize :: Int64
chunkSize = Word16 -> ByteString -> Int64
chunkSizeFor (Pipe -> Word16
mcs Pipe
pipe) ByteString
bs
        chunks :: [ByteString]
chunks    = Int64 -> ByteString -> [ByteString]
split Int64
chunkSize ByteString
bs
        terminal :: ByteString
terminal  = Word16 -> ByteString
forall a. Binary a => a -> ByteString
encodeStrict (Word16
0 :: Word16)
        conn :: ConnectionWithTimeout
conn      = Pipe -> ConnectionWithTimeout
connection Pipe
pipe

        mkChunk :: ByteString -> [ByteString]
        mkChunk :: ByteString -> [ByteString]
mkChunk ByteString
chunk = let size :: Word16
size = Int -> Word16
forall a b. (Integral a, Num b) => a -> b
fromIntegral (ByteString -> Int
B.length ByteString
chunk) :: Word16
                        in  [Word16 -> ByteString
forall a. Binary a => a -> ByteString
encodeStrict Word16
size, ByteString
chunk]

fetch :: MonadPipe m => HasCallStack => Pipe -> m Response
fetch :: Pipe -> m Response
fetch Pipe
pipe = do ByteString
bs <- m ByteString
forall (m :: * -> *). MonadPipe m => m ByteString
chunks
                let response :: Either UnpackError Response
response = Get Structure -> ByteString -> Either UnpackError Structure
forall a. Get a -> ByteString -> Either UnpackError a
unpackAction Get Structure
forall a. BoltValue a => Get a
unpackT ByteString
bs Either UnpackError Structure
-> (Structure -> Either UnpackError Response)
-> Either UnpackError Response
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Structure -> Either UnpackError Response
forall a (m :: * -> *).
(FromStructure a, MonadError UnpackError m) =>
Structure -> m a
fromStructure
                case Either UnpackError Response
response of
                  Left UnpackError
ue -> BoltError -> m Response
forall e (m :: * -> *) a. MonadError e m => e -> m a
throwError (BoltError -> m Response) -> BoltError -> m Response
forall a b. (a -> b) -> a -> b
$ UnpackError -> BoltError
WrongMessageFormat UnpackError
ue
                  Right Response
a -> Response -> m Response
forall (f :: * -> *) a. Applicative f => a -> f a
pure Response
a

  where conn :: ConnectionWithTimeout
conn = Pipe -> ConnectionWithTimeout
connection Pipe
pipe

        chunks :: MonadPipe m => m BSL.ByteString
        chunks :: m ByteString
chunks = do Word16
size <- ByteString -> Word16
forall a. Binary a => ByteString -> a
decode (ByteString -> Word16) -> m ByteString -> m Word16
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ConnectionWithTimeout -> Word16 -> m ByteString
forall (m :: * -> *).
(MonadPipe m, HasCallStack) =>
ConnectionWithTimeout -> Word16 -> m ByteString
recvChunk ConnectionWithTimeout
conn Word16
2
                    ByteString
chunk <- ConnectionWithTimeout -> Word16 -> m ByteString
forall (m :: * -> *).
(MonadPipe m, HasCallStack) =>
ConnectionWithTimeout -> Word16 -> m ByteString
recvChunk ConnectionWithTimeout
conn Word16
size
                    if ByteString -> Bool
BSL.null ByteString
chunk
                      then ByteString -> m ByteString
forall (f :: * -> *) a. Applicative f => a -> f a
pure ByteString
BSL.empty
                      else do ByteString
rest <- m ByteString
forall (m :: * -> *). MonadPipe m => m ByteString
chunks
                              ByteString -> m ByteString
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ByteString -> m ByteString) -> ByteString -> m ByteString
forall a b. (a -> b) -> a -> b
$! ByteString
chunk ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> ByteString
rest

-- Helper functions

handshake :: MonadPipe m => HasCallStack => Pipe -> BoltCfg -> m ()
handshake :: Pipe -> BoltCfg -> m ()
handshake Pipe
pipe BoltCfg
bcfg = do let conn :: ConnectionWithTimeout
conn = Pipe -> ConnectionWithTimeout
connection Pipe
pipe
                         ConnectionWithTimeout -> ByteString -> m ()
forall (m :: * -> *).
(MonadIO m, HasCallStack) =>
ConnectionWithTimeout -> ByteString -> m ()
C.send ConnectionWithTimeout
conn (Word32 -> ByteString
forall a. Binary a => a -> ByteString
encodeStrict (Word32 -> ByteString) -> Word32 -> ByteString
forall a b. (a -> b) -> a -> b
$ BoltCfg -> Word32
magic BoltCfg
bcfg)
                         ConnectionWithTimeout -> ByteString -> m ()
forall (m :: * -> *).
(MonadIO m, HasCallStack) =>
ConnectionWithTimeout -> ByteString -> m ()
C.send ConnectionWithTimeout
conn (BoltCfg -> ByteString
boltVersionProposal BoltCfg
bcfg)
                         Word32
serverVersion <- ByteString -> Word32
forall a. Binary a => ByteString -> a
decode (ByteString -> Word32) -> m ByteString -> m Word32
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ConnectionWithTimeout -> Word16 -> m ByteString
forall (m :: * -> *).
(MonadPipe m, HasCallStack) =>
ConnectionWithTimeout -> Word16 -> m ByteString
recvChunk ConnectionWithTimeout
conn Word16
4
                         Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Word32
serverVersion Word32 -> Word32 -> Bool
forall a. Eq a => a -> a -> Bool
/= BoltCfg -> Word32
version BoltCfg
bcfg) (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$
                           BoltError -> m ()
forall e (m :: * -> *) a. MonadError e m => e -> m a
throwError BoltError
UnsupportedServerVersion
                         Pipe -> Request -> m ()
forall (m :: * -> *).
(MonadPipe m, HasCallStack) =>
Pipe -> Request -> m ()
flush Pipe
pipe (BoltCfg -> Request
createInit BoltCfg
bcfg)
                         Response
response <- Pipe -> m Response
forall (m :: * -> *).
(MonadPipe m, HasCallStack) =>
Pipe -> m Response
fetch Pipe
pipe
                         Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (Response -> Bool
isSuccess Response
response) (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$
                           BoltError -> m ()
forall e (m :: * -> *) a. MonadError e m => e -> m a
throwError BoltError
AuthentificationFailed

boltVersionProposal :: BoltCfg -> ByteString
boltVersionProposal :: BoltCfg -> ByteString
boltVersionProposal BoltCfg
bcfg = [ByteString] -> ByteString
B.concat ([ByteString] -> ByteString) -> [ByteString] -> ByteString
forall a b. (a -> b) -> a -> b
$ Word32 -> ByteString
forall a. Binary a => a -> ByteString
encodeStrict (Word32 -> ByteString) -> [Word32] -> [ByteString]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [BoltCfg -> Word32
version BoltCfg
bcfg, Word32
0, Word32
0, Word32
0]

recvChunk :: MonadPipe m => HasCallStack => ConnectionWithTimeout -> Word16 -> m BSL.ByteString
recvChunk :: ConnectionWithTimeout -> Word16 -> m ByteString
recvChunk ConnectionWithTimeout
conn Word16
size = Int -> m ByteString
forall (m :: * -> *). MonadPipe m => Int -> m ByteString
helper (Word16 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word16
size)
  where helper :: MonadPipe m => Int -> m BSL.ByteString
        helper :: Int -> m ByteString
helper Int
0  = ByteString -> m ByteString
forall (f :: * -> *) a. Applicative f => a -> f a
pure ByteString
BSL.empty
        helper Int
sz = do Maybe ByteString
mbChunk <- ConnectionWithTimeout -> Int -> m (Maybe ByteString)
forall (m :: * -> *).
(MonadIO m, HasCallStack) =>
ConnectionWithTimeout -> Int -> m (Maybe ByteString)
C.recv ConnectionWithTimeout
conn Int
sz
                       case Maybe ByteString
mbChunk of
                         Just ByteString
chunk -> ByteString -> ByteString -> ByteString
BSL.chunk ByteString
chunk (ByteString -> ByteString) -> m ByteString -> m ByteString
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> m ByteString
forall (m :: * -> *). MonadPipe m => Int -> m ByteString
helper (Int
sz Int -> Int -> Int
forall a. Num a => a -> a -> a
- ByteString -> Int
B.length ByteString
chunk)
                         Maybe ByteString
Nothing    -> BoltError -> m ByteString
forall e (m :: * -> *) a. MonadError e m => e -> m a
throwError BoltError
CannotReadChunk

chunkSizeFor :: Word16 -> BSL.ByteString -> Int64
chunkSizeFor :: Word16 -> ByteString -> Int64
chunkSizeFor Word16
maxSize ByteString
bs = Int64
1 Int64 -> Int64 -> Int64
forall a. Num a => a -> a -> a
+ Int64 -> Int64 -> Int64
forall a. Integral a => a -> a -> a
div Int64
len Int64
noc
  where len :: Int64
len = ByteString -> Int64
BSL.length ByteString
bs
        noc :: Int64
noc = Int64
1 Int64 -> Int64 -> Int64
forall a. Num a => a -> a -> a
+ Int64 -> Int64 -> Int64
forall a. Integral a => a -> a -> a
div Int64
len (Word16 -> Int64
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word16
maxSize)

split :: Int64 -> BSL.ByteString -> [ByteString]
split :: Int64 -> ByteString -> [ByteString]
split Int64
size ByteString
bs | ByteString -> Bool
BSL.null ByteString
bs = []
              | Bool
otherwise = let (ByteString
chunk, ByteString
rest) = Int64 -> ByteString -> (ByteString, ByteString)
BSL.splitAt Int64
size ByteString
bs
                            in ByteString -> ByteString
BSL.toStrict ByteString
chunk ByteString -> [ByteString] -> [ByteString]
forall a. a -> [a] -> [a]
: Int64 -> ByteString -> [ByteString]
split Int64
size ByteString
rest