A demo of a simple multi-threaded IRC bot in Haskell using stm and async.

stm-ircbot-demo.hs

{-# LANGUAGE NondecreasingIndentation, OverloadedStrings #-}

-- This demo program is placed into the Public Domain.

module Main where

-- imports are grouped by the package they come from

-- base
import Control.Monad (mapM_, forever)
import Data.Monoid ((<>))
import Data.List (isPrefixOf)
import Control.Concurrent (threadDelay)
-- network-simple
import Network.Simple.TCP (withSocketsDo, connect, send, recv)
-- bytestring
import Data.ByteString (ByteString)
import qualified Data.ByteString as B
import qualified Data.ByteString.Char8 as C
-- stm
import Control.Concurrent.STM
import Control.Concurrent.STM.TQueue
-- async
import Control.Concurrent.Async

type IRCMessage = String

-- | For sake of simplicity, we will not use Text, which supports
-- utf8 and all that. Instead we'll assume that all bytes are 8-bit
-- characters and turn the bytes into a String that way. Don't do
-- this in production code, that would be unwise.
parseIRCMessage :: ByteString -> (ByteString,[IRCMessage])
parseIRCMessage bytes = let
    -- try to break off just one message frame
    (h,t) = B.breakSubstring "\r\n" bytes
    -- recursively check for any more frames
    (spare, rest) = parseIRCMessage (B.drop 2 t)
    -- If 't' is null then we don't have even a single complete frame, so we
    -- just hand back all the bytes as spare.
    in if B.null t
        then (bytes, [])
        else (spare, C.unpack h:rest)

main :: IO ()
main = withSocketsDo $ do
    inQ <- newTQueueIO
    outQ <- newTQueueIO
    
    let incoming socket spare = do
            maybeBytes <- recv socket 512
            case maybeBytes of
                Nothing -> putStrLn "The Socket Closed!"
                Just bytes -> do
                    let (newSpare, messages) = parseIRCMessage $ spare <> bytes
                    atomically $ mapM_ (writeTQueue inQ) messages
                    incoming socket newSpare
    
        outgoing socket = forever $ do
            ircMessage <- atomically $ readTQueue outQ
            putStrLn $ "O " <> ircMessage
            send socket (C.pack ircMessage)
            send socket "\r\n"
    
        worker = forever $ do
            message <- atomically $ readTQueue inQ
            putStrLn $ "I " <> message
            let pingPrefix = "PING :"
                pongPrefix = "PONG :"
                prefixLength = length pingPrefix
            -- you could do more complex work here of course, but for now
            -- we'll just respond to any pings and that's it.
            if pingPrefix `isPrefixOf` message
                then atomically $ writeTQueue outQ $ pongPrefix <> drop prefixLength message
                else return ()
    
    connect "irc.freenode.net" "6667" $ \(socket, remoteAddr) -> do
        withAsync (incoming socket "") $ \inAsync -> do
        withAsync (outgoing socket) $ \outAsync -> do
        withAsync worker $ \workAsync -> do
        -- send the openers
        atomically $ do
            writeTQueue outQ "NICK farmbotirc"
            writeTQueue outQ "USER farmbotirc 8 * :farmbotirc the bot"
            writeTQueue outQ "JOIN #lokathor"
        -- wait 20 seconds and then say something in the channel unprompted.
        threadDelay (1000000*20)
        atomically $ do
            writeTQueue outQ "PRIVMSG #lokathor :Hello from Haskell."
        -- this just blocks in the main thread while the program runs
        wait inAsync