SocketChannel.hs

{-
  Copyright (c) Meta Platforms, Inc. and affiliates.
  All rights reserved.

  This source code is licensed under the BSD-style license found in the
  LICENSE file in the root directory of this source tree.
-}

-- |
-- An implementation of a Thrift channel in Haskell using
-- 'Network.Socket'. Note that the wire format of messages is not
-- necessarily compatible with other transports, in particular you
-- can't mix and match SocketChannel clients/servers with
-- HeaderChannel clients/servers.

{-# LANGUAGE CPP #-}
module Thrift.Channel.SocketChannel
  ( -- * Socket channel types and API
    SocketChannel(..)
  , SocketConfig(..)
  , withSocketChannel
  , withSocketChannelIO

  , -- * Utilities used by test servers
    sendBS
  , teardownSock
  , threadWaitRecv
  , recvBlockBytes
  , localhost
  ) where

import Control.Concurrent
import Control.Exception
import Data.Proxy
import Network.Socket
import Network.Socket.ByteString
import System.Posix.Types (Fd(..))
import qualified Data.ByteString as BS
import qualified Data.Text as T

import Thrift.Channel
import Thrift.Monad
import Thrift.Protocol
import Thrift.Protocol.Id

newtype SocketChannel t = SocketChannel Socket

instance ClientChannel SocketChannel where
  sendRequest ch@(SocketChannel sock) req sendcb recvcb = do
    sendOnewayRequest ch req sendcb
    recvRes <- try (threadWaitRecv sock recvBlockBytes)
    case recvRes of
      Right s -> recvcb (Right $ Response s mempty)
      Left (e :: SomeException) ->
        recvcb (Left $ ChannelException $ T.pack (show e))

  sendOnewayRequest (SocketChannel sock) req sendcb = do
    res <- sendBS sock (reqMsg req)
    sendcb res

data SocketConfig a = SocketConfig
  { socketHost       :: HostName
  , socketPortNum    :: PortNumber
  , socketProtocolId :: ProtocolId
  }

deriving instance Show (SocketConfig a)

-- | Given a 'SocketConfig' that specifies where to find the
--   server we want to reach and what protocol it uses,
--   connect to the server and run the given client
--   computation with it.
withSocketChannel
  :: SocketConfig t
  -> (forall p . Protocol p => ThriftM p SocketChannel t a)
  -> IO a
withSocketChannel SocketConfig{..} f =
  withSocketChannelIO socketHost socketPortNum $ \ch -> do
    withProxy socketProtocolId (go ch f)

  where go :: Protocol p
           => SocketChannel t
           -> ThriftM p SocketChannel t a
           -> Proxy p
           -> IO a
        go ch f _proxy = runThrift f ch

-- | A lower level variant of 'withSocketChannel' that connects
--   to the server at the given host and port and passes the
--   resulting 'SocketChannel' to the third argument.
withSocketChannelIO
  :: HostName
  -> PortNumber
  -> (SocketChannel t -> IO a)
  -> IO a
withSocketChannelIO host port f = withSocketClient host port $ \sock _ ->
  f (SocketChannel sock)

sendBS :: Socket -> BS.ByteString -> IO (Maybe ChannelException)
sendBS sock bs = try (sendAll sock bs) >>= \res -> case res of
  Left  e -> return (Just e)
  Right _ -> return Nothing

-- | Block (using 'threadWaitRead') until some data is available,
--   and then call 'recv' to grab it.
threadWaitRecv
  :: Socket
  -> Int -- ^ max. number of bytes we want to receive
  -> IO BS.ByteString
threadWaitRecv sock numBytes = withFdSocket sock $ \fd ->
  threadWaitRead (Fd fd) >> recv sock numBytes

withSocketClient
  :: HostName
  -> PortNumber
  -> (Socket -> AddrInfo -> IO a)
  -> IO a
withSocketClient host port f = withSocketsDo $ do
  addr <- resolveClient host port
  bracket (setupClientSock addr) teardownSock $ \sock ->
    f sock addr

setupClientSock :: AddrInfo -> IO Socket
setupClientSock addr = do
  s <- openSocket addr
  connect s (addrAddress addr)
  return s

teardownSock :: Socket -> IO ()
teardownSock sock = shutdown sock ShutdownBoth `finally` close sock

resolveClient :: HostName -> PortNumber -> IO AddrInfo
resolveClient host port = do
  results <- getAddrInfo (Just hints) (Just host) (Just $ show port)
  case results of
    [] -> error $ "SocketChannel.resolveClient: " <>
      "getAddrInfo returned an empty list"
    (a:_) -> return a

  where hints = defaultHints { addrSocketType = Stream }

localhost :: HostName
localhost =
#ifdef IPV4
  "127.0.0.1"
#else
  "::1"
#endif

recvBlockBytes :: Int
recvBlockBytes = 4096

#if !MIN_VERSION_network(3,1,2)
openSocket :: AddrInfo -> IO Socket
openSocket addr =
  socket (addrFamily addr) (addrSocketType addr) (addrProtocol addr)
#endif