pipes-irc-server/src/Pipes/IRC/Server.hs

{-# LANGUAGE OverloadedStrings #-}

module Pipes.IRC.Server
  ( startIrcServer
  )
where

import           Control.Concurrent              (threadDelay)
import           Control.Concurrent.Async
import           Control.Concurrent.STM
import           Control.Lens                    as L
import           Control.Monad
import           Control.Monad.RWS
import           Data.ByteString.Char8           as BS
import           Data.Map                        as M
import           Data.Maybe
import           Data.Set                        as S
import           Data.Time.Clock
import           Network.Socket                  as NS
import           Pipes
import           Pipes.Attoparsec
import           Pipes.Concurrent                as PC
import           Pipes.IRC.Message
import           Pipes.IRC.Server.EventHandler
import           Pipes.IRC.Server.IrcMonad
import           Pipes.IRC.Server.Log
import           Pipes.IRC.Server.MessageHandler
import           Pipes.IRC.Server.Server
import           Pipes.IRC.Server.Types
import           Pipes.Network.TCP               as PN

version :: ByteString
version = "0.1a"

parseMessage :: Producer ByteString IO ()
             -> Producer (Either ByteString IrcMessage) IO ()
parseMessage prod = do
  void $ for (parsed parseMsgOrLine prod) $ \res ->
    case res of
      (Left _)    -> yield $ Left "ERROR Bad Parse"
      (Right val) -> yield $ Right val
  return ()

renderMessage :: Pipe IrcMessage ByteString IO ()
renderMessage = forever $ do
  msg <- await
  let output = renderIrcMessage msg
  yield output

filterMsgs :: Pipe (Either ByteString IrcMessage) IrcMessage IO ()
filterMsgs = forever $ do
  cmd <- await
  case cmd of
    Left bs -> liftIO $ logLine $ BS.concat ["BAD COMMAND: ", bs]
    Right c -> yield c

addIrcConnection :: ServerState -> IrcConnection -> IO Int
addIrcConnection srv client = do
  let clients = srv ^. ircConnections
      ids = srv ^. ircConnIds
  cid <- atomically $ do
    lastId <- readTVar ids
    let newId = lastId + 1
    writeTVar ids newId
    return newId
  atomically $ modifyTVar' clients $ M.insert cid client
  return cid

delIrcConnection :: ServerState -> Int -> IO ()
delIrcConnection srv cid = atomically $ do
  cs <- readTVar (srv ^. ircConnections)
  case M.lookup cid cs of
    Just conn -> do
      let nn = case conn ^. reg of
            Unreg{ _rcvdNick = Just n }           -> n
            RegUser{ _regdNick = NickName n _ _ } -> n
            _                                     -> ""
      modifyTVar' (srv ^. ircState) $ ircDelUser nn
    _ -> return ()
  modifyTVar' (srv ^. ircConnections) $ M.delete cid

ircMonadTransaction :: ServerState -> Int -> RegState -> IrcMonad ()
                    -> IO (RegState, [IrcEvent])
ircMonadTransaction srv cid userReg action = do
  cmap <- readTVarIO $ srv ^. ircConnections
  let hostname = view (at cid . traverse . hname) cmap
  curTime <- getCurrentTime
  atomically $ do
    sState <- readTVar $ srv ^. ircState
    let sConf = srv ^. ircConfig
    let cState = ClientState { _clientReg = userReg
                             , _clientServer = sState
                             , _clientHost = hostname
                             , _clientConn = cid }
    let (_, newState, events) = runRWS action sConf cState
    writeTVar (srv ^. ircState) (newState ^. clientServer)
    modifyTVar' (srv ^. ircConnections) $
      M.adjust ((reg .~ newState ^. clientReg).(lastCom .~ curTime)) cid
    return (newState ^. clientReg, events)

runIrcMonad :: ServerState -> Int -> IrcMonad () -> IO ()
runIrcMonad srv cid action = do
  cs <- readTVarIO (srv ^. ircConnections)
  let cReg = maybe (Unreg Nothing Nothing Nothing) (view reg) (cs ^. at cid)
  (_, events) <- ircMonadTransaction srv cid cReg action
  forM_ events $ ircEventHandler srv

cmdHandler :: ServerState -> Int -> Consumer IrcMessage IO ()
cmdHandler srv cid =
  let cReg = Unreg Nothing Nothing Nothing
  in do
    conns <- liftIO $ readTVarIO $ srv ^. ircConnections
    case M.lookup cid conns of
      Just c  -> handle (c ^. hname) cReg
      Nothing -> return ()
  where
    handle h userReg = do
      -- wait for the next command
      nextMsg <- await

      liftIO $ logMsg nextMsg (fromMaybe "unknown" h) userReg

      -- run the handler in a transaction
      (newReg, events) <-
        liftIO $ ircMonadTransaction srv cid userReg (ircMessageHandler nextMsg)

      -- handle resulting events
      aliveL <- liftIO $ forM events $ ircEventHandler srv

      -- loop for the next command
      when (and aliveL) $ handle h newReg

idlePinger :: ServerState -> Int -> IO ()
idlePinger srv cid =
  let
    pingMsg  = IrcMessage Nothing (Left PING)
                 [":" `append` (srv ^. ircConfig . ircHostName)]
    oneMinute = 60 * 1000000 -- microseconds
    getLastCom = do conns <- readTVarIO (srv ^. ircConnections)
                    return $ conns ! cid ^. lastCom
    resetPong = atomically $ modifyTVar' (srv ^. ircConnections) $
                               M.adjust (gotPong .~ False) cid
    checkPong = do conns <- readTVarIO (srv ^. ircConnections)
                   return $ conns ! cid ^. gotPong
    timeoutLoop b = when b $ do
      threadDelay oneMinute
      curTime <- getCurrentTime
      time <- getLastCom
      let diffTime = toRational . diffUTCTime curTime $ time
      if diffTime > 60
        then do
          resetPong

          atomically $ do
            conns <- readTVar (srv ^. ircConnections)
            void $ PC.send (conns ! cid ^. out) pingMsg
            return ()

          threadDelay oneMinute

          checkPong >>= timeoutLoop
        else timeoutLoop True
  in do
    -- timeoutLoop will call itself repeatedly until a timeout occurs
    timeoutLoop True
    -- Then we need to send a timeout message
    runIrcMonad srv cid $ doQuit (Just "Ping timeout")

listenHandler :: ServerState -> (Socket, SockAddr) -> IO ()
listenHandler srv (lsock, _) =
  forever $ acceptFork lsock $ \(csock, caddr) -> do
    let sockWriter = toSocket csock
        sockReader = fromSocket csock 4096

    (hName, _) <- getNameInfo [] True False caddr

    (writeEnd, readEnd) <- spawn unbounded
    curTime <- getCurrentTime

    logLine $ BS.pack $
      "Accepted connection from " ++ fromMaybe "unknown" hName

    let client = IrcConnection
                  { _sock  = csock
                  , _addr  = caddr
                  , _hname = fmap BS.pack hName
                  , _out   = writeEnd
                  , _reg   = Unreg Nothing Nothing Nothing
                  , _lastCom = curTime
                  , _gotPong = False
                  }

    cid <- addIrcConnection srv client

    let handler = cmdHandler srv cid

    r <- async $ runEffect $
           parseMessage sockReader >-> filterMsgs >-> handler
    link r

    idle <- async $ idlePinger srv cid
    link idle

    w <- async $ runEffect $
           fromInput readEnd >-> renderMessage >-> sockWriter
    link w

    void $ waitAnyCancel [r, w, idle]

    logLine $ BS.pack $
      "Connection from " ++ fromMaybe "unknown" hName ++ " terminated"

    delIrcConnection srv cid

mkIrcServer :: IrcConfig -> IO ServerState
mkIrcServer config = do
  let nks = S.empty
      urs = M.empty
      chs = M.empty
      srv = IrcServer nks urs chs version
  tvState <- newTVarIO srv
  tvCns <- newTVarIO M.empty
  tvRef <- newTVarIO 0
  return $ ServerState tvState config tvCns tvRef

startIrcServer :: IrcConfig -> IO (Async ())
startIrcServer config = do
  srv <- mkIrcServer config
  let sHost = srv ^. ircConfig . ircHost
      sPort = srv ^. ircConfig . ircPort
  logLine $ BS.pack $
    mconcat ["Starting server on ", show sHost, " ", show sPort]
  async $ PN.listen sHost sPort (listenHandler srv)