ConnectionContext.hs

{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE DeriveDataTypeable #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE OverloadedStrings #-}
module Database.Redis.ConnectionContext (
    ConnectionContext(..)
  , ConnectTimeout(..)
  , ConnectionLostException(..)
  , ConnectAddr(..)
  , connect
  , disconnect
  , send
  , recv
  , errConnClosed
  , enableTLS
  , flush
  , ioErrorToConnLost
) where

import Control.Monad(when)
import qualified Data.ByteString as B
import qualified Data.ByteString.Char8 as Char8
import qualified Data.ByteString.Lazy as LB
import qualified Data.IORef as IOR
import Control.Concurrent.MVar(newMVar, readMVar, swapMVar)
import Control.Exception(bracketOnError, Exception, throwIO, try, finally, mask_)
import Data.Functor(void)
import qualified Network.Socket as NS
import qualified Network.TLS as TLS
import System.IO(Handle, hSetBinaryMode, hClose, IOMode(..), hFlush, hIsOpen)
import System.IO.Error(catchIOError)
import System.Timeout (timeout)

data ConnectionContext = NormalHandle Handle | TLSContext TLS.Context Handle

instance Show ConnectionContext where
    show (NormalHandle _) = "NormalHandle"
    show (TLSContext _ _) = "TLSContext"

data Connection = Connection
    { ctx :: ConnectionContext
    , lastRecvRef :: IOR.IORef (Maybe B.ByteString) }

instance Show Connection where
    show Connection{..} = "Connection{ ctx = " ++ show ctx ++ ", lastRecvRef = IORef}"

data ConnectPhase
  = PhaseUnknown
  | PhaseResolve
  | PhaseOpenSocket
  deriving (Show)

newtype ConnectTimeout = ConnectTimeout ConnectPhase
  deriving (Show)

instance Exception ConnectTimeout

data ConnectionLostException = ConnectionLost deriving Show
instance Exception ConnectionLostException

data ConnectAddr
  = ConnectAddrHostPort NS.HostName NS.PortNumber
  | ConnectAddrUnixSocket String
  deriving (Eq, Show)

connect :: ConnectAddr -> Maybe Int -> Maybe TLS.ClientParams -> IO ConnectionContext
connect connectAddr timeoutOpt mTlsParams =
  bracketOnError hConnect hClose $ \h -> do
    hSetBinaryMode h True
    case (mTlsParams, connectAddr) of
      (Just defaultTlsParams, ConnectAddrHostPort host port) -> do
        -- The defaultTlsParams are used to connect to the first
        -- host in the cluster, other hosts have different
        -- hostnames and so require a different server
        -- identification params
        let tlsParams = defaultTlsParams {
              TLS.clientServerIdentification =  (host, Char8.pack $ show port)
            }
        enableTLS tlsParams (NormalHandle h)
      _ -> return $ NormalHandle h
  where
        hConnect = do
          phaseMVar <- newMVar PhaseUnknown
          let doConnect = hConnect' phaseMVar
          case timeoutOpt of
            Nothing -> doConnect
            Just micros -> do
              result <- timeout micros doConnect
              case result of
                Just h -> return h
                Nothing -> do
                  phase <- readMVar phaseMVar
                  errConnectTimeout phase
        hConnect' mvar = bracketOnError createSock NS.close $ \sock -> do
          NS.setSocketOption sock NS.KeepAlive 1
          void $ swapMVar mvar PhaseResolve
          void $ swapMVar mvar PhaseOpenSocket
          NS.socketToHandle sock ReadWriteMode
          where
            createSock = case connectAddr of
              ConnectAddrHostPort hostName portNumber -> do
                addrInfo <- getHostAddrInfo hostName portNumber
                connectSocket addrInfo
              ConnectAddrUnixSocket addr -> bracketOnError
                (NS.socket NS.AF_UNIX NS.Stream NS.defaultProtocol)
                NS.close
                (\sock -> NS.connect sock (NS.SockAddrUnix addr) >> return sock)

getHostAddrInfo :: NS.HostName -> NS.PortNumber -> IO [NS.AddrInfo]
getHostAddrInfo hostname port = NS.getAddrInfo (Just hints) (Just hostname) (Just $ show port)
  where
    hints = NS.defaultHints
      { NS.addrSocketType = NS.Stream }

errConnectTimeout :: ConnectPhase -> IO a
errConnectTimeout phase = throwIO $ ConnectTimeout phase

connectSocket :: [NS.AddrInfo] -> IO NS.Socket
connectSocket [] = error "connectSocket: unexpected empty list"
connectSocket (addr:rest) = tryConnect >>= \case
  Right sock -> return sock
  Left err   -> if null rest
                then throwIO err
                else connectSocket rest
  where
    tryConnect :: IO (Either IOError NS.Socket)
    tryConnect = bracketOnError createSock NS.close $ \sock ->
      try (NS.connect sock $ NS.addrAddress addr) >>= \case
      Right () -> return (Right sock)
      Left err -> NS.close sock >> return (Left err)
      where
        createSock = NS.socket (NS.addrFamily addr)
                               (NS.addrSocketType addr)
                               (NS.addrProtocol addr)

send :: ConnectionContext -> B.ByteString -> IO ()
send (NormalHandle h) requestData =
    ioErrorToConnLost (B.hPut h requestData)
send (TLSContext ctx _) requestData =
    ioErrorToConnLost (TLS.sendData ctx (LB.fromStrict requestData))

recv :: ConnectionContext -> IO B.ByteString
recv (NormalHandle h) = ioErrorToConnLost $ B.hGetSome h 4096
recv (TLSContext ctx _) = TLS.recvData ctx


ioErrorToConnLost :: IO a -> IO a
ioErrorToConnLost a = a `catchIOError` const errConnClosed

errConnClosed :: IO a
errConnClosed = throwIO ConnectionLost


enableTLS :: TLS.ClientParams -> ConnectionContext -> IO ConnectionContext
enableTLS tlsParams (NormalHandle h) = do
  ctx <- TLS.contextNew h tlsParams
  TLS.handshake ctx
  return $! TLSContext ctx h
enableTLS _ c@(TLSContext _ _) = return c


disconnect :: ConnectionContext -> IO ()
disconnect (NormalHandle h) = mask_ $ do
  open <- hIsOpen h
  when open $ hClose h
disconnect (TLSContext ctx h) =
  TLS.bye ctx `finally` TLS.contextClose ctx `finally` (hIsOpen h >>= \open -> when open $ hClose h)

flush :: ConnectionContext -> IO ()
flush (NormalHandle h) = hFlush h
flush (TLSContext ctx _) = TLS.contextFlush ctx