Connection.hs

-- | Internal module. Not part of the public API.
module Database.Bolty.Connection.Connection
  ( connect
  , close
  , send
  , receiveBytestring
  , receiveBinary
  ) where

import           Control.Monad.Trans    (MonadIO(..))
import           Data.Default           (Default(..))
import           Data.Persist           (HasEndianness, runGet, getBE)
import           Data.Word              (Word16)
import           GHC.Stack              (HasCallStack, withFrozenCallStack)
import           Network.Connection     (Connection, ConnectionParams(..), connectFromSocket, connectionClose,
                                        connectionGetExact, connectionPut,
                                        initConnectionContext)
import qualified Data.ByteString        as BS
import qualified Data.Text              as T
import qualified Network.Socket         as NS
import qualified System.Timeout         as ST (timeout)
import           TextShow               (TextShow, showt)


timeoutThrow :: (HasCallStack, TextShow b) => Int -> b -> IO a -> IO a
timeoutThrow milliseconds associated_data action = withFrozenCallStack $ do
  res <- ST.timeout (milliseconds * 1_000) action
  case res of
    Just a  -> return a
    Nothing -> fail $ "Timeout: " ++ show (showt associated_data)


-- | Open a raw TCP connection to the given host and port with optional TLS.
connect :: (MonadIO m, HasCallStack) => Bool -> T.Text -> Word16 -> Int -> m (Connection, Int)
connect use_tls host port timeout = liftIO $ do
  ctx <- initConnectionContext
  let params = ConnectionParams
        { connectionHostname  = T.unpack host
        , connectionPort      = fromIntegral port
        , connectionUseSecure = if use_tls then Just def else Nothing
        , connectionUseSocks  = Nothing
        }
  let hints = NS.defaultHints { NS.addrSocketType = NS.Stream }
  addrs <- NS.getAddrInfo (Just hints) (Just $ T.unpack host) (Just $ show port)
  case addrs of
    [] -> fail $ "Cannot resolve " ++ T.unpack host
    (addr:_) -> do
      sock <- NS.openSocket addr
      NS.setSocketOption sock NS.NoDelay 1
      timeoutThrow timeout ("connecting to " <> host <> ":" <> showt port) $ do
        NS.connect sock (NS.addrAddress addr)
        conn <- connectFromSocket ctx sock params
        pure (conn, timeout)


-- | Close a raw TCP connection.
close :: (MonadIO m, HasCallStack) => Connection -> Int -> m ()
close conn timeout = liftIO $ timeoutThrow timeout ("closing connection" :: T.Text) $ connectionClose conn


-- | Receive exactly @size@ bytes from the connection with a timeout.
receiveBytestring :: (MonadIO m, HasCallStack) => Connection -> Int -> Int -> m BS.ByteString
receiveBytestring conn timeout size =
  liftIO $ timeoutThrow timeout ("receiving " <> showt size <> " bytes") $ connectionGetExact conn size


-- | Receive and decode a big-endian binary value from the connection.
receiveBinary :: forall a m. (HasEndianness a, MonadIO m, HasCallStack) => Connection -> Int -> Int -> m a
receiveBinary conn timeout size =
  fmap decodeBE $ liftIO $ timeoutThrow timeout ("receiving data" :: T.Text) $ connectionGetExact conn size
  where
    decodeBE bs = case runGet (getBE @a) bs of
      Left e  -> error $ "receiveBinary: " ++ e
      Right a -> a


-- | Send raw bytes over the connection with a timeout.
send :: (MonadIO m, HasCallStack) => Connection -> Int -> BS.ByteString -> m ()
send conn timeout bytes = do
  let timeout_message = "sending " <> showt (BS.length bytes) <> " bytes"
  liftIO $ timeoutThrow timeout timeout_message $ connectionPut conn bytes