More refactoring of error handling and making better use of lenses

src/Pipes/IRC/Server/IrcMonad.hs

@@ -3,16 +3,18 @@
 module Pipes.IRC.Server.IrcMonad
 where
 
+import           Control.Error
 import           Control.Lens
 import           Control.Monad.RWS
 import           Data.ByteString.Char8    (ByteString)
 import qualified Data.ByteString.Char8    as BS
 import qualified Data.Map                 as M
-import           Data.Maybe               (catMaybes, fromJust, isJust)
+import           Data.Maybe               (fromJust)
 import qualified Data.Set                 as S
 
 import           Pipes.IRC.Message.Types
 import           Pipes.IRC.Server.Channel
+import           Pipes.IRC.Server.Server
 import           Pipes.IRC.Server.Types
 import           Pipes.IRC.Server.User
 
@@ -39,7 +41,11 @@ ppHostPreference hp = case hp of
 -- * Monadic utilities
 
 useChan :: ChanKey -> IrcMonad (Maybe IrcChannel)
-useChan cname = fmap (M.lookup cname) $ use (clientServer . ircChannels)
+useChan cname = use $ clientServer . ircChannels . at cname
+
+useUserChans :: NickKey -> IrcMonad (S.Set ChanKey)
+useUserChans nn =
+  use (clientServer . ircUsers . at nn . traverse . userChannels)
 
 useNick :: IrcMonad (Maybe NickKey)
 useNick = do
@@ -49,15 +55,6 @@ useNick = do
     RegUser (NickName nn _ _) -> Just nn
     _                         -> Nothing
 
-validateNick :: NickKey -> IrcMonad Bool
-validateNick nickname = do
-  nickSet <- use $ clientServer . ircNicks
-  if S.member nickname nickSet
-   then do
-      tellNumeric err_nicknameinuse [nickname, ":Nickname is already in use."]
-      return False
-   else return True
-
 channelTargets :: ChanKey -> IrcMonad [Int]
 channelTargets chname = do
   srv <- use clientServer
@@ -128,8 +125,78 @@ allChanEcho iMsg = do
       let chans = S.elems $ usr ^. userChannels
       chanEcho chans iMsg
 
+disconnectUser :: Int -> IrcMessage -> IrcMonad ()
+disconnectUser cid msg = do
+  allChanEcho msg
+  tell [Close cid]
+
+-- * Command validation utilities
+
+type ErrParam    = (IrcReply, [IrcParam])
+type IrcMonadErr = EitherT ErrParam IrcMonad
+
+tellErr :: Either (IrcReply, [IrcParam]) a -> IrcMonad ()
+tellErr (Left (r, ps)) = void $ tellNumeric r ps
+tellErr _ = return ()
+
+runValidation :: IrcMonadErr () -> IrcMonad ()
+runValidation = tellErr <=< runEitherT
+
+ensure :: Bool -> IrcReply -> [IrcParam] -> IrcMonadErr ()
+ensure p r ps = when p $ left (r, ps)
+
+ensureUse :: IrcMonad (Maybe a)  -> ErrParam -> IrcMonadErr a
+ensureUse u e = lift u >>= hoistEither . note e
+
+checkParamLength :: BS.ByteString -> [IrcParam] -> Int -> IrcMonadErr ()
+checkParamLength cmd ps n =
+  ensure (length ps >= n)
+    err_needmoreparams [cmd, ":Need more parameters"]
+
+checkSuppliedNickname :: [IrcParam] -> IrcMonadErr NickKey
+checkSuppliedNickname ps = do
+  ensure (not . null $ ps) err_nonicknamegiven [":must supply a nickname"]
+  return (head ps)
+
+checkRegistration :: IrcMonadErr NickKey
+checkRegistration =
+  ensureUse useNick (err_notregistered, [":You have not registered"])
+
+checkChannelPresence :: ChanKey -> IrcMonadErr IrcChannel
+checkChannelPresence ckey =
+  ensureUse (useChan ckey) (err_badchannelkey, [ckey, ":No such channel"])
+
+checkUserOnChan :: NickKey -> ChanKey -> IrcChannel -> IrcMonadErr ()
+checkUserOnChan nn c ch =
+  ensure (chanHasUser nn ch) err_notonchannel [c, ":Not on channel"]
+
+checkUserNotOnChan :: NickKey -> ChanKey -> IrcChannel -> IrcMonadErr ()
+checkUserNotOnChan nn c ch =
+  ensure (not . chanHasUser nn $ ch) err_notonchannel [c, ":Already on channel"]
+
+checkInvitation :: NickKey -> ChanKey -> IrcChannel -> IrcMonadErr ()
+checkInvitation nn c ch = ensure (ircInviteCheck nn ch)
+                            err_inviteonlychan [c, ":Cannot join channel (+i)"]
+
+checkPassKey :: Maybe PassKey -> ChanKey -> IrcChannel -> IrcMonadErr ()
+checkPassKey k c chan = ensure (ircPassCheck k chan)
+                          err_badchannelkey [c, ":Cannot join channel (+k)"]
+
+checkNickFree :: NickKey -> IrcMonadErr ()
+checkNickFree nickname = do
+  nickSet <- lift (use $ clientServer . ircNicks)
+  ensure (S.member nickname nickSet)
+         err_nicknameinuse [nickname, ":Nickname is already in use."]
+
 -- * Adding responses to the Writer portion of the monad
 
+renderQuitMsg :: Maybe BS.ByteString -> IrcMonad BS.ByteString
+renderQuitMsg (Just msg) = return $ BS.append "Quit: " msg
+renderQuitMsg Nothing    = useNick >>= \nickname ->
+  return $ case nickname of
+    Just n -> BS.append "Quit: " n
+    Nothing   -> ""
+
 tellYOURHOST :: NickKey -> IrcMonad ()
 tellYOURHOST nickname = do
   srvname <- view ircHostName

src/Pipes/IRC/Server/MessageHandler.hs

@@ -5,20 +5,17 @@ module Pipes.IRC.Server.MessageHandler
   ( ircMessageHandler )
 where
 
+import           Control.Applicative       ((<$>))
 import           Control.Error
 import           Control.Lens
 import           Control.Monad.RWS
 import qualified Data.ByteString.Char8     as BS
-import           Data.Map                  as M
-import           Data.Maybe
 import           Data.Set                  as S
 
 import           Pipes.IRC.Message.Types
-import           Pipes.IRC.Server.Channel
 import           Pipes.IRC.Server.IrcMonad
 import           Pipes.IRC.Server.Server
 import           Pipes.IRC.Server.Types
-import           Pipes.IRC.Server.User
 import           Pipes.IRC.Server.Util
 
 ircMessageHandler :: IrcMessage -> IrcMonad ()
@@ -42,155 +39,6 @@ unregHandler msg@IrcMessage{..} =
     Left QUIT -> handleQUIT msg
     _         -> return ()
 
-unregPASS :: IrcMessage -> IrcMonad ()
-unregPASS IrcMessage{..} =
-  if length params < 1
-  then tellNumeric err_needmoreparams ["PASS", ":Need more parameters"]
-  else do
-    clientReg . rcvdPass .= (Just $ head params)
-    return ()
-
-unregNICK :: IrcMessage -> IrcMonad ()
-unregNICK IrcMessage{..} =
-  if length params /= 1
-  then tellNumeric err_nonicknamegiven [":must supply a nickname"]
-  else let nickname = head params in
-    validateNick nickname >>= flip when
-      (clientReg . rcvdNick .= Just nickname >>
-       clientServer . ircNicks %= S.insert nickname >>
-       tryRegistration)
-
-
-unregUSER :: IrcMessage -> IrcMonad ()
-unregUSER IrcMessage{..} =
-  if length params < 4
-  then tellNumeric err_needmoreparams ["USER", ":need more parameters"]
-  else do
-    clientReg . rcvdName .= (Just $ head params)
-    tryRegistration
-
-renderQuitMsg :: Maybe BS.ByteString -> IrcMonad BS.ByteString
-renderQuitMsg (Just msg) = return $ BS.append "Quit: " msg
-renderQuitMsg Nothing    = useNick >>= \nickname ->
-  return $ case nickname of
-    Just nick -> BS.append "Quit: " nick
-    Nothing   -> ""
-
-handlePING :: IrcMessage -> IrcMonad ()
-handlePING _ = do
-  srvname <- view ircHostName
-  tellCommand PONG [":" <> srvname]
-
-handlePONG :: IrcMessage -> IrcMonad ()
-handlePONG _ = do
-  cid <- use clientConn
-  tell [Pong cid]
-
--- JOIN, MODE, KICK, PART, QUIT and of course PRIVMSG/NOTICE need to be
--- echoed to channels that the user belongs to
-
-doQuit :: IrcMessage -> Maybe BS.ByteString -> IrcMonad ()
-doQuit msg quitParamIn = do
-  connId <- use clientConn
-  quitMsg <- renderQuitMsg quitParamIn
-  hostname <- use clientHost
-  let hoststr = fromMaybe "unknown hostname" hostname
-  let quitParam = BS.concat [ ":Closing Link: ", hoststr
-                            , " (", quitMsg, ")"]
-  tellCommand ERROR [quitParam]
-  allChanEcho msg{params = [quitParam]}
-  tell [Close connId]
-  return ()
-
-handleQUIT :: IrcMessage -> IrcMonad ()
-handleQUIT msg@IrcMessage{..} =
-  doQuit msg $ case params of
-    []  -> Nothing
-    p:_ -> Just p
-
-handleJOIN :: IrcMessage -> IrcMonad ()
-handleJOIN msg@IrcMessage{..} =
-  case params of
-    []      -> tellNumeric err_needmoreparams ["JOIN", ":Not enough parameters"]
-    ["0"]   -> do
-      Just nn <- useNick
-      cs <- use $ clientServer . ircUsers . at nn
-      let Just chans = fmap (^. userChannels) cs
-      doPart msg{command=Left PART} (S.elems chans) Nothing
-    cs:[]   -> doJoin msg $ zipParams (parseParamList cs) []
-    cs:ks:_ -> doJoin msg $ zipParams (parseParamList cs) (parseParamList ks)
-
-handlePART :: IrcMessage -> IrcMonad ()
-handlePART msg@IrcMessage{..} =
-  case params of
-    []      -> tellNumeric err_needmoreparams ["PART", ":Not enough parameters"]
-    cs:[]   -> doPart msg (parseParamList cs) Nothing
-    cs:pm:_ -> doPart msg (parseParamList cs) (Just pm)
-
-tellErr :: IrcReply -> [IrcParam] -> IrcMonadErr ()
-tellErr r ps = lift (tellNumeric r ps) >> left []
-
-checkRegistration :: IrcMonadErr BS.ByteString
-checkRegistration = do
-  mNick <- lift useNick
-  when (isNothing mNick) $
-    tellErr err_notregistered [":You have not registered"]
-  right $ fromJust mNick
-
-doJoin :: IrcMessage -> [(BS.ByteString, Maybe BS.ByteString)] -> IrcMonad ()
-doJoin msg chans = forM_ chans $ \(c, k) -> runEitherT $ do
-  nick <- checkRegistration
-  mChan <- lift $ useChan c
-  when (isJust mChan) $ do
-    let chan = fromJust mChan
-
-    when (chanHasUser nick chan) $ left "already on channel"
-
-    when (ircInviteCheck nick chan) $
-      tellErr err_inviteonlychan [c, ":Cannot join channel (+i)"]
-
-    when (ircPassCheck k chan) $
-      tellErr err_badchannelkey [c, ":Cannot join channel (+k)"]
-
-  lift $ do
-    clientServer %= ircJoin nick c
-    tellTOPIC c
-    tellNAMES [c]
-    chanEcho [c] msg
-
-doPart :: IrcMessage -> [BS.ByteString] -> Maybe BS.ByteString -> IrcMonad ()
-doPart msg chans pmsg = forM_ chans $ \c -> runEitherT $ do
-  nn <- checkRegistration
-  mChan <- lift . use $ clientServer . ircChannels . at c
-  when (isNothing mChan) $ tellErr err_nosuchchannel [c, ":No such channel"]
-  let ch = fromJust mChan
-  unless (chanHasUser nn ch) $ tellErr err_notonchannel [c, ":Not on channel"]
-
-  lift $ do
-    let plist = case pmsg of
-          Just bs -> [c, BS.append ":" bs]
-          Nothing -> [c]
-    let newMsg = msg{params = plist}
-    chanEcho [c] newMsg
-    findReceivers [nn] >>= fwdMsgNoReplace msg
-    clientServer %= ircPart nn c
-
-tryRegistration :: IrcMonad ()
-tryRegistration = do
-  regState <- use clientReg
-  hostname <- use clientHost
-  case regState of
-    Unreg _ (Just nickname) (Just name) -> do
-      let nn = NickName nickname (Just name) hostname
-      clientReg .= RegUser nn
-      usr <- mkUser
-      clientServer . ircUsers %= M.insert nickname usr
-
-      tellWELCOME  nickname
-      tellYOURHOST nickname
-      tellMOTD     nickname
-    _                               -> return ()
-
 regHandler :: IrcMessage -> IrcMonad ()
 regHandler msg@IrcMessage{..} = do
   pMsg <- addUserPrefix msg
@@ -205,6 +53,112 @@ regHandler msg@IrcMessage{..} = do
     Left QUIT    -> handleQUIT pMsg
     _            -> return ()
 
+-- * Handlers for messages coming from unregistered users
+
+unregPASS :: IrcMessage -> IrcMonad ()
+unregPASS IrcMessage{..} = runValidation $ do
+  checkParamLength "PASS" params 1
+  lift $ clientReg . rcvdPass ?= head params
+
+unregNICK :: IrcMessage -> IrcMonad ()
+unregNICK IrcMessage{..} = runValidation $ do
+  nickname <- checkSuppliedNickname params
+  checkNickFree nickname
+  lift $ do
+    clientReg . rcvdNick ?= nickname
+    clientServer . ircNicks . contains nickname .= True
+    tryRegistration
+
+unregUSER :: IrcMessage -> IrcMonad ()
+unregUSER IrcMessage{..} = runValidation $ do
+  checkParamLength "USER" params 4
+  lift $ do
+    clientReg . rcvdName ?= head params
+    tryRegistration
+
+tryRegistration :: IrcMonad ()
+tryRegistration = do
+  regState <- use clientReg
+  hostname <- use clientHost
+  case regState of
+    Unreg _ (Just nickname) (Just name) -> do
+      usr <- mkUser
+      clientReg .= (RegUser $ NickName nickname (Just name) hostname)
+      clientServer . ircUsers . at nickname ?= usr
+
+      tellWELCOME  nickname
+      tellYOURHOST nickname
+      tellMOTD     nickname
+    _                                   -> return ()
+
+-- * Handlers for messages coming from registered users
+
+handlePING :: IrcMessage -> IrcMonad ()
+handlePING _ = do
+  srvname <- view ircHostName
+  tellCommand PONG [":" <> srvname]
+
+handlePONG :: IrcMessage -> IrcMonad ()
+handlePONG _ = do
+  cid <- use clientConn
+  tell [Pong cid]
+
+-- JOIN, MODE, KICK, PART, QUIT and of course PRIVMSG/NOTICE need to be
+-- echoed to channels that the user belongs to
+
+handleQUIT :: IrcMessage -> IrcMonad ()
+handleQUIT msg@IrcMessage{..} = do
+  connId   <- use clientConn
+  hostname <- fromMaybe "unknown hostname" <$> use clientHost
+  quitMsg  <- renderQuitMsg (headMay params)
+  let quitParam = mconcat [":Closing Link: ", hostname, " (", quitMsg, ")"]
+  tellCommand ERROR [quitParam]
+  disconnectUser connId msg{params = [quitParam]}
+
+handleJOIN :: IrcMessage -> IrcMonad ()
+handleJOIN msg@IrcMessage{..} = runValidation $ do
+  checkParamLength "JOIN" params 1
+  nn <- checkRegistration
+  lift $ case params of
+    ["0"]   -> do cs <- useUserChans nn
+                  doPart msg{command=Left PART} (S.elems cs) Nothing
+    cs:[]   -> doJoin msg $ zipParams (parseParamList cs) []
+    cs:ks:_ -> doJoin msg $ zipParams (parseParamList cs) (parseParamList ks)
+
+handlePART :: IrcMessage -> IrcMonad ()
+handlePART msg@IrcMessage{..} = runValidation $ do
+  checkParamLength "PART" params 1
+  lift $ case params of
+    cs:[]   -> doPart msg (parseParamList cs) Nothing
+    cs:pm:_ -> doPart msg (parseParamList cs) (Just pm)
+
+doJoin :: IrcMessage -> [(ChanKey, Maybe PassKey)] -> IrcMonad ()
+doJoin msg chans = forM_ chans $ \(c, k) -> runValidation $ do
+  nick <- checkRegistration
+  chan <- checkChannelPresence c
+  checkUserNotOnChan nick c chan
+  checkInvitation nick c chan
+  checkPassKey k c chan
+
+  lift $ do
+    clientServer %= ircJoin nick c
+    tellTOPIC c
+    tellNAMES [c]
+    chanEcho [c] msg
+
+doPart :: IrcMessage -> [BS.ByteString] -> Maybe BS.ByteString -> IrcMonad ()
+doPart msg chans pmsg = forM_ chans $ \c -> runValidation $ do
+  nn <- checkRegistration
+  ch <- checkChannelPresence c
+  checkUserOnChan nn c ch
+  let newMsg = msg{params = c : maybe [] ((:[]) . BS.append ":") pmsg}
+
+  lift $ do
+    chanEcho [c] newMsg
+    rs <- findReceivers [nn]
+    fwdMsgNoReplace msg rs
+    clientServer %= ircPart nn c
+
 handlePRIVMSG :: IrcMessage -> IrcMonad ()
 handlePRIVMSG msg@IrcMessage{..} = do
   case params of

src/Pipes/IRC/Server/Server.hs

@@ -117,7 +117,7 @@ ircInviteCheck n chan =
 -- | Determine whether the given channel will disallow joining due to
 -- a missing or incorrect password. A 'True' value indicates that
 -- conditions for joining are met.
-ircPassCheck :: Maybe ByteString -- ^ password supplied by user
+ircPassCheck :: Maybe PassKey    -- ^ password supplied by user
              -> IrcChannel       -- ^ channel to check
              -> Bool             -- ^ may the user join?
 ircPassCheck k chan =

src/Pipes/IRC/Server/Types.hs

@@ -22,6 +22,7 @@ import           Pipes.Network.TCP       (HostPreference (..), ServiceName,
 
 type NickKey = ByteString
 type ChanKey = ByteString
+type PassKey = ByteString
 
 type IrcEvents = [IrcEvent]
 
@@ -112,5 +113,4 @@ data ClientState =
               } deriving (Show)
 makeLenses ''ClientState
 
-type IrcMonad a = RWS IrcConfig IrcEvents ClientState a
+type IrcMonad = RWS IrcConfig IrcEvents ClientState
-type IrcMonadErr a = EitherT String (RWS IrcConfig IrcEvents ClientState) a