path: root/Presence/main.hs

{-# LANGUAGE CPP               #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE TypeFamilies      #-}
module Main where

import Debug.Trace
import System.Directory
import Control.Monad
import System.Posix.Signals
import System.Posix.Types
import System.Posix.Process
import Data.Maybe
import Data.Char
import ConfigFiles

import System.INotify
#ifndef NOUTMP
import UTmp 
-- Breaks profiling build with error:
--     Dynamic linking required, but this is a non-standard build (eg. prof).
--         You need to build the program twice: once the normal way, and then
--             in the desired way using -osuf to set the object file suffix.
--
-- TODO: Figure out wtf ghc is trying to tell me.
--       In the mean time, use -DNOTMP to build for profiling.
-- 
#endif
import FGConsole
import XMPPServer
import Data.HList
import Network.Socket (sClose)
import Control.Exception
import LocalPeerCred
import System.Posix.User
import qualified Data.Set as Set
import           Data.Set as Set (Set,(\\))
import qualified Data.Map as Map
import           Data.Map as Map (Map)

import Control.Concurrent.STM
import Control.Concurrent
import Control.Monad.Trans.Maybe
import Control.Monad.IO.Class

import ByteStringOperators
import qualified Data.ByteString.Lazy.Char8 as L
import           Data.ByteString.Lazy.Char8 as L (ByteString,putStrLn)
import qualified Prelude
import           Prelude hiding (putStrLn)
import System.Environment
import qualified Text.Show.ByteString as L


data UnixSession = UnixSession {
    localhost :: ByteString,
    unix_uid :: (IORef (Maybe UserID)),
    unix_resource :: (IORef (Maybe L.ByteString)),
    presence_state :: PresenceState
}

instance XMPPSession UnixSession where
    data XMPPClass UnixSession = UnixSessions PresenceState
    newSession (UnixSessions state) sock handle = do
        muid <- getLocalPeerCred sock
        L.putStrLn $ "SESSION: open " <++> bshow muid
        uid_ref <- newIORef muid
        res_ref <- newIORef Nothing
        return $ UnixSession (hostname state) uid_ref res_ref state
    setResource  s  resource    = do
        writeIORef (unix_resource s) (Just resource)
        L.putStrLn $ "SESSION: resource " <++> resource
    getJID       s             = do
        let host = localhost s
        muid <- readIORef (unix_uid s)
        user <- maybe (return "nobody") 
                   (\uid -> 
                        handle (\(SomeException _) -> 
                                return . L.append "uid." . L.pack . show $ uid)
                          $ do
                             user <- fmap userName $ getUserEntryForID uid
                             return (L.pack user)
                        )
                   muid
        rsc <- readIORef (unix_resource s)
        -- let jid = user <++> "@" <++> host <++?> "/" <++$> rsc
        L.putStrLn $ "SESSION: jid " <++> L.show (JID (Just user) host rsc)
        return (JID (Just user) host rsc)
    closeSession _             = L.putStrLn "SESSION: close"
    subscribe session Nothing = do
        let tmvar = localSubscriber (presence_state session)
        atomically $ subscribeToChan tmvar
    subscribe session (Just jid) = do
        let tvar = subscriberMap (presence_state session)
        atomically $ subscribeToMap tvar jid
    announcePresence session (Presence jid status) = do
        subs <- readTVarIO $ subscriberMap (presence_state session)
        update_presence Nothing (fmap snd subs) (Set.singleton jid) (const status)


subscribeToChan tmvar =
            (do (cnt,chan) <- takeTMVar tmvar
                putTMVar tmvar (cnt+1,chan)
                chan' <- dupTChan chan
                return chan' )
          `orElse`
            (do chan <- newTChan
                putTMVar tmvar (1,chan)
                return chan )

subscribeToMap tvar jid = do
            subs <- readTVar tvar
            let mbchan       = Map.lookup jid subs
            (chan,subs') <- 
                do case mbchan of
                    Nothing -> do
                      newchan <- newTChan
                      return (newchan, Map.insert jid (1,newchan) subs)
                    Just (cnt,chan) -> do
                      chan' <- dupTChan chan
                      return (chan', Map.insert jid (cnt+1,chan) subs)
            writeTVar tvar subs'
            return chan

matchResource tty jid = maybe Away (avail . (==tty)) $ resource jid
 where
    avail True  = Available
    avail False = Away

sendPresence chan jid status = 
    (liftIO  . atomically . writeTChan chan . Presence jid $ status) :: MaybeT IO ()

lookupT jid subscribers = MaybeT . return $ Map.lookup jid subscribers 

update_presence locals_greedy subscribers state getStatus =
    forM_ (Set.toList state) $ \jid -> do
      let status = getStatus jid 
      runMaybeT $ do
        chan <- lookupT jid subscribers
        sendPresence chan jid status
      runMaybeT $ do
        chan <- MaybeT . return $ locals_greedy
        sendPresence chan jid status
      putStrLn $ bshow jid <++> " " <++> bshow status

type RefCount = Int

data PresenceState = PresenceState 
    { hostname :: ByteString
    , currentTTY    :: TVar ByteString
    , activeUsers   :: TVar (Set JID)
    , subscriberMap :: TVar (Map JID (RefCount,TChan Presence))
    , localSubscriber :: TMVar (RefCount,TChan Presence)
    }

newPresenceState hostname = atomically $ do
    tty  <- newTVar ""
    us   <- newTVar (Set.empty)
    subs <- newTVar (Map.empty)
    locals_greedy <- newEmptyTMVar
    return $ PresenceState hostname tty us subs locals_greedy

track_login host state e = do
#ifndef NOUTMP
    us <- UTmp.users
#else
    let us = []
#endif
    let toJabberId host (user,tty,_) = 
            if L.take 3 tty == "tty"
              then Just (jid user host tty)
              else Nothing
        new_users = Set.fromList $ mapMaybe (toJabberId host) us
    (tty,known_users,subs,locals_greedy) <- atomically $ do
            tty <- readTVar                $ currentTTY state
            st  <- flip swapTVar new_users $ activeUsers state
            xs  <- readTVar                $ subscriberMap state
            locals_greedy <- tryReadTMVar $ localSubscriber state
            return (tty,st,fmap snd xs,fmap snd locals_greedy)
    let arrivals   = new_users \\ known_users
        departures = known_users \\ new_users
    update_presence locals_greedy subs departures $ const Offline
    update_presence locals_greedy subs arrivals   $ matchResource tty

on_chvt state vtnum = do
    let tty = L.snoc "tty" $ intToDigit (fromIntegral vtnum)
    L.putStrLn $ "VT switch: " <++> tty
    (users,subs,locals_greedy) <- atomically $ do
        us   <- readTVar $ activeUsers state
        subs <- readTVar $ subscriberMap state
        writeTVar (currentTTY state) tty
        locals_greedy <- tryReadTMVar $ localSubscriber state
        return (us,fmap snd subs,fmap snd locals_greedy)
    update_presence locals_greedy subs users $ matchResource tty

data UnixConfig = UnixConfig

instance XMPPConfig UnixConfig where
    getBuddies     _ user = ConfigFiles.getBuddies user
    getSubscribers _ user = ConfigFiles.getSubscribers user

start :: ByteString -> IO ()
start host = do
    tracked <- newPresenceState host
    let dologin e = track_login host tracked e
        dologin :: t -> IO ()

    chan <- atomically $ subscribeToChan (localSubscriber tracked)
    remotes <- forkIO $ seekRemotePeers (/=host) UnixConfig chan

    installHandler sigUSR1 (Catch (dologin (userError "signaled"))) Nothing
    -- installHandler sigTERM (CatchOnce (dologin (userError "term signaled"))) Nothing
    mtty <- monitorTTY (on_chvt tracked)
    inotify <- initINotify
#ifndef NOUTMP
    wd <- addWatch
            inotify
            [CloseWrite] -- [Open,Close,Access,Modify,Move]
            utmp_file
            dologin
#endif
    sockLocals <- listenForXmppClients (UnixSessions tracked) 5222 HNil
    sockRemotes <- listenForRemotePeers (UnixSessions tracked) 5269 HNil

    threadDelay 1000 -- wait a moment to obtain current tty
    dologin ()
    putStrLn "\nHit enter to terminate...\n"
    getLine
    killThread remotes
    sClose sockLocals
    sClose sockRemotes
    -- threadDelay 1000
    putStrLn "closed listener."
    unmonitorTTY mtty
    putStrLn "unhooked tty monitor."
#ifndef NOUTMP
    removeWatch wd
#endif
    putStrLn "Normal termination."

sendUSR1 pid = do
    signalProcess sigUSR1 pid

getStartupAction []     = throw (userError "pid file?") >> return (Right "")
getStartupAction (p:ps) = do
    handle onEr $
        ( do
            pid <- fmap CPid (readFile p >>= readIO)
            -- signal pid
            return (Left pid) )
 where
    onEr (SomeException _) = do
        pid <- getProcessID
        putStrLn $ "starting pid = " <++> bshow pid
        handle (\(SomeException _) -> getStartupAction ps)
            (do 
                writeFile p (show pid)
                putStrLn $ "writing " <++> bshow p
                -- start daemon
                return (Right p) )

runOnce ps run notify = getStartupAction ps >>= doit
 where
    doit (Left pid      ) = notify pid
    doit (Right pidfile ) = do 
        run
        removeFile pidfile

getOptions [] = Map.empty

getOptions (('-':opt_name):xs) =
    if xs/=[] && xs!!0!!0/='-'
      then Map.insert (L.pack opt_name) (L.pack (xs!!0)) (getOptions (tail xs))
      else Map.insert (L.pack opt_name) "" (getOptions xs)

getOptions (x0:xs) = getOptions xs


main = do
    opts <- fmap getOptions getArgs
    let hostname = maybe "localhost" id (Map.lookup "n" opts)
    L.putStrLn $ "hostname = " <++> hostname
    runOnce ["/var/run/presence.pid","/tmp/presence.pid"] (start hostname) sendUSR1