--------------------------------------------------------------------------------
-- Logstash client for Haskell                                                --
--------------------------------------------------------------------------------
-- This source code is licensed under the MIT license found in the LICENSE    --
-- file in the root directory of this source tree.                            --
--------------------------------------------------------------------------------

-- | This module implements a logstash client for the tcp input plugin:
-- https://www.elastic.co/guide/en/logstash/7.10/plugins-inputs-tcp.html
module Logstash.TCP (
    LogstashTcpConfig(..),
    logstashTcp,
    logstashTcpPool,
    logstashTls,
    logstashTlsPool
) where 

--------------------------------------------------------------------------------

import Control.Monad.IO.Class

import Data.Acquire
import qualified Data.ByteString as BS
import qualified Data.ByteString.Lazy as BSL
import Data.Default.Class
import Data.Time
import Data.Pool

import Network.Socket
import Network.Socket.ByteString (sendAll)
import Network.TLS

import Logstash.Connection

--------------------------------------------------------------------------------

-- | Represents configurations for Logstash TCP inputs.
data LogstashTcpConfig = LogstashTcpConfig {
    -- | The hostname of the server to connect to.
    LogstashTcpConfig -> String
logstashTcpHost :: String,
    -- | The port of the server to connect to.
    LogstashTcpConfig -> Int
logstashTcpPort :: Int
} deriving (LogstashTcpConfig -> LogstashTcpConfig -> Bool
(LogstashTcpConfig -> LogstashTcpConfig -> Bool)
-> (LogstashTcpConfig -> LogstashTcpConfig -> Bool)
-> Eq LogstashTcpConfig
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: LogstashTcpConfig -> LogstashTcpConfig -> Bool
$c/= :: LogstashTcpConfig -> LogstashTcpConfig -> Bool
== :: LogstashTcpConfig -> LogstashTcpConfig -> Bool
$c== :: LogstashTcpConfig -> LogstashTcpConfig -> Bool
Eq, Int -> LogstashTcpConfig -> ShowS
[LogstashTcpConfig] -> ShowS
LogstashTcpConfig -> String
(Int -> LogstashTcpConfig -> ShowS)
-> (LogstashTcpConfig -> String)
-> ([LogstashTcpConfig] -> ShowS)
-> Show LogstashTcpConfig
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [LogstashTcpConfig] -> ShowS
$cshowList :: [LogstashTcpConfig] -> ShowS
show :: LogstashTcpConfig -> String
$cshow :: LogstashTcpConfig -> String
showsPrec :: Int -> LogstashTcpConfig -> ShowS
$cshowsPrec :: Int -> LogstashTcpConfig -> ShowS
Show)

instance Default LogstashTcpConfig where 
    def :: LogstashTcpConfig
def = LogstashTcpConfig :: String -> Int -> LogstashTcpConfig
LogstashTcpConfig{
        logstashTcpHost :: String
logstashTcpHost = String
"127.0.0.1",
        logstashTcpPort :: Int
logstashTcpPort = Int
5000
    }

-- | `connectTCP` @config@ establishes a TCP socket connection to the server
-- configured by @config@.
connectTcpSocket 
    :: MonadIO m 
    => LogstashTcpConfig 
    -> m Socket
connectTcpSocket :: LogstashTcpConfig -> m Socket
connectTcpSocket LogstashTcpConfig{Int
String
logstashTcpPort :: Int
logstashTcpHost :: String
logstashTcpPort :: LogstashTcpConfig -> Int
logstashTcpHost :: LogstashTcpConfig -> String
..} = IO Socket -> m Socket
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO Socket -> m Socket) -> IO Socket -> m Socket
forall a b. (a -> b) -> a -> b
$ IO Socket -> IO Socket
forall a. IO a -> IO a
withSocketsDo (IO Socket -> IO Socket) -> IO Socket -> IO Socket
forall a b. (a -> b) -> a -> b
$ do
    -- initialise a TCP socket for the given host and port
    let hints :: AddrInfo
hints = AddrInfo
defaultHints{ addrSocketType :: SocketType
addrSocketType = SocketType
Stream }
    AddrInfo
addr <- [AddrInfo] -> AddrInfo
forall a. [a] -> a
head ([AddrInfo] -> AddrInfo) -> IO [AddrInfo] -> IO AddrInfo
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe AddrInfo -> Maybe String -> Maybe String -> IO [AddrInfo]
getAddrInfo 
        (AddrInfo -> Maybe AddrInfo
forall a. a -> Maybe a
Just AddrInfo
hints) (String -> Maybe String
forall a. a -> Maybe a
Just String
logstashTcpHost) (String -> Maybe String
forall a. a -> Maybe a
Just (String -> Maybe String) -> String -> Maybe String
forall a b. (a -> b) -> a -> b
$ Int -> String
forall a. Show a => a -> String
show Int
logstashTcpPort)
    Socket
sock <- Family -> SocketType -> ProtocolNumber -> IO Socket
socket (AddrInfo -> Family
addrFamily AddrInfo
addr) (AddrInfo -> SocketType
addrSocketType AddrInfo
addr) (AddrInfo -> ProtocolNumber
addrProtocol AddrInfo
addr)
    Socket -> SockAddr -> IO ()
connect Socket
sock (SockAddr -> IO ()) -> SockAddr -> IO ()
forall a b. (a -> b) -> a -> b
$ AddrInfo -> SockAddr
addrAddress AddrInfo
addr

    -- return the socket
    Socket -> IO Socket
forall (f :: * -> *) a. Applicative f => a -> f a
pure Socket
sock

-- | `createTcpConnection` @config@ establishes a `LogstashConnection` via
-- TCP to the server configured by @config@.
createTcpConnection 
    :: MonadIO m 
    => LogstashTcpConfig 
    -> m LogstashConnection
createTcpConnection :: LogstashTcpConfig -> m LogstashConnection
createTcpConnection LogstashTcpConfig
cfg = do 
    -- establish a TCP connection
    Socket
sock <- LogstashTcpConfig -> m Socket
forall (m :: * -> *). MonadIO m => LogstashTcpConfig -> m Socket
connectTcpSocket LogstashTcpConfig
cfg

    -- return the Logstash connection
    LogstashConnection -> m LogstashConnection
forall (f :: * -> *) a. Applicative f => a -> f a
pure (LogstashConnection -> m LogstashConnection)
-> LogstashConnection -> m LogstashConnection
forall a b. (a -> b) -> a -> b
$ LogstashConnection :: (ByteString -> IO ()) -> IO () -> LogstashConnection
LogstashConnection{
        writeData :: ByteString -> IO ()
writeData = Socket -> ByteString -> IO ()
sendAll Socket
sock (ByteString -> IO ())
-> (ByteString -> ByteString) -> ByteString -> IO ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> ByteString
BSL.toStrict,
        closeConnection :: IO ()
closeConnection = Socket -> IO ()
close Socket
sock
    }

-- | `logstashTcp` @config@ produces an `Acquire` for establishing
-- `LogstashConnection` values for the given TCP @config@.
logstashTcp 
    :: LogstashTcpConfig 
    -> Acquire LogstashConnection
logstashTcp :: LogstashTcpConfig -> Acquire LogstashConnection
logstashTcp LogstashTcpConfig
cfg = IO LogstashConnection
-> (LogstashConnection -> IO ()) -> Acquire LogstashConnection
forall a. IO a -> (a -> IO ()) -> Acquire a
mkAcquire (LogstashTcpConfig -> IO LogstashConnection
forall (m :: * -> *).
MonadIO m =>
LogstashTcpConfig -> m LogstashConnection
createTcpConnection LogstashTcpConfig
cfg) LogstashConnection -> IO ()
closeConnection 

-- | `logstashTcpPool` @config stripes ttl resources@ produces a `Pool`
-- of `LogstashConnection` values for the given TCP @config@. The other
-- parameters are passed on directly to `createPool`.
logstashTcpPool 
    :: LogstashTcpConfig
    -> Int 
    -> NominalDiffTime
    -> Int 
    -> IO LogstashPool
logstashTcpPool :: LogstashTcpConfig
-> Int -> NominalDiffTime -> Int -> IO LogstashPool
logstashTcpPool LogstashTcpConfig
cfg = IO LogstashConnection
-> (LogstashConnection -> IO ())
-> Int
-> NominalDiffTime
-> Int
-> IO LogstashPool
forall a.
IO a
-> (a -> IO ()) -> Int -> NominalDiffTime -> Int -> IO (Pool a)
createPool (LogstashTcpConfig -> IO LogstashConnection
forall (m :: * -> *).
MonadIO m =>
LogstashTcpConfig -> m LogstashConnection
createTcpConnection LogstashTcpConfig
cfg) LogstashConnection -> IO ()
closeConnection

--------------------------------------------------------------------------------

-- | `createTlsConnection` @config params@ establishes a `LogstashConnection` via
-- TLS to the server configured by @config@ and the TLS configuration given by
-- @params@.
createTlsConnection 
    :: MonadIO m 
    => LogstashTcpConfig 
    -> ClientParams
    -> m LogstashConnection
createTlsConnection :: LogstashTcpConfig -> ClientParams -> m LogstashConnection
createTlsConnection LogstashTcpConfig
cfg ClientParams
params = do 
    -- establish a TCP connection
    Socket
sock <- LogstashTcpConfig -> m Socket
forall (m :: * -> *). MonadIO m => LogstashTcpConfig -> m Socket
connectTcpSocket LogstashTcpConfig
cfg

    -- establish a TLS connection on top
    Context
ctx <- Socket -> ClientParams -> m Context
forall (m :: * -> *) backend params.
(MonadIO m, HasBackend backend, TLSParams params) =>
backend -> params -> m Context
contextNew Socket
sock ClientParams
params
    Context -> m ()
forall (m :: * -> *). MonadIO m => Context -> m ()
handshake Context
ctx

    -- return the Logstash connection
    LogstashConnection -> m LogstashConnection
forall (f :: * -> *) a. Applicative f => a -> f a
pure (LogstashConnection -> m LogstashConnection)
-> LogstashConnection -> m LogstashConnection
forall a b. (a -> b) -> a -> b
$ LogstashConnection :: (ByteString -> IO ()) -> IO () -> LogstashConnection
LogstashConnection{
        writeData :: ByteString -> IO ()
writeData = Context -> ByteString -> IO ()
forall (m :: * -> *). MonadIO m => Context -> ByteString -> m ()
sendData Context
ctx,
        closeConnection :: IO ()
closeConnection = do 
            Context -> IO ()
forall (m :: * -> *). MonadIO m => Context -> m ()
bye Context
ctx 
            Socket -> IO ()
close Socket
sock
    }

-- | `logstashTls` @config params@ produces an `Acquire` for establishing
-- `LogstashConnection` values for the given TCP @config@ and TLS @params@.
logstashTls 
    :: LogstashTcpConfig 
    -> ClientParams 
    -> Acquire LogstashConnection 
logstashTls :: LogstashTcpConfig -> ClientParams -> Acquire LogstashConnection
logstashTls LogstashTcpConfig
cfg ClientParams
params = 
    IO LogstashConnection
-> (LogstashConnection -> IO ()) -> Acquire LogstashConnection
forall a. IO a -> (a -> IO ()) -> Acquire a
mkAcquire (LogstashTcpConfig -> ClientParams -> IO LogstashConnection
forall (m :: * -> *).
MonadIO m =>
LogstashTcpConfig -> ClientParams -> m LogstashConnection
createTlsConnection LogstashTcpConfig
cfg ClientParams
params) LogstashConnection -> IO ()
closeConnection 

-- | `logstashTlsPool` @config params stripes ttl resources@ produces a 
-- `Pool` of `LogstashConnection` values for the given TCP @config@ and TLS
-- @params@. The other parameters are passed on directly to `createPool`.
logstashTlsPool
    :: LogstashTcpConfig 
    -> ClientParams 
    -> Int 
    -> NominalDiffTime
    -> Int
    -> IO LogstashPool
logstashTlsPool :: LogstashTcpConfig
-> ClientParams -> Int -> NominalDiffTime -> Int -> IO LogstashPool
logstashTlsPool LogstashTcpConfig
cfg ClientParams
params = 
    IO LogstashConnection
-> (LogstashConnection -> IO ())
-> Int
-> NominalDiffTime
-> Int
-> IO LogstashPool
forall a.
IO a
-> (a -> IO ()) -> Int -> NominalDiffTime -> Int -> IO (Pool a)
createPool (LogstashTcpConfig -> ClientParams -> IO LogstashConnection
forall (m :: * -> *).
MonadIO m =>
LogstashTcpConfig -> ClientParams -> m LogstashConnection
createTlsConnection LogstashTcpConfig
cfg ClientParams
params) LogstashConnection -> IO ()
closeConnection 

--------------------------------------------------------------------------------