path: root/xmppServer.hs

{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE TupleSections #-}
import System.Posix.Signals
import Control.Concurrent.STM
import Control.Concurrent.STM.TMVar
import Control.Monad.Trans.Resource (runResourceT)
import Control.Monad.IO.Class (MonadIO, liftIO)
import Network.Socket
    ( addrAddress
    , getAddrInfo
    , defaultHints
    , addrFlags
    , AddrInfoFlag(AI_CANONNAME,AI_V4MAPPED)
    , SockAddr(..)
    )
import System.Endian (fromBE32)
import Data.Monoid ( (<>) )
import qualified Data.Text as Text
import qualified Data.Text.IO as Text
import qualified Data.Text.Encoding as Text
import Control.Monad
import qualified Network.BSD as BSD
import qualified Data.Text as Text
import Data.Text (Text)
import qualified Data.Map as Map
import Data.Map (Map)
import Control.Exception ({-evaluate,-}handle,SomeException(..),bracketOnError,ErrorCall(..))
import System.Posix.User (getUserEntryForID,userName)
import qualified Data.ByteString.Lazy.Char8 as L
import qualified ConfigFiles
import Data.Maybe (listToMaybe)

import UTmp (ProcessID,users)
import LocalPeerCred
import XMPPServer
-- import Server

unsplitJID (n,h,r) = jid
 where
    jid0 = maybe h (\n->n<>"@"<>h) n
    jid = maybe jid0 (\r->jid0<>"/"<>r) r

splitJID :: Text -> (Maybe Text,Text,Maybe Text)
splitJID bjid = 
    let xs = splitAll '@' bjid
        ys = splitAll '/' (last xs)
        splitAll c bjid = take 1 xs0 ++ map (Text.drop 1) (drop 1 xs0)
            where xs0 = Text.groupBy (\x y-> y/=c) bjid
        server = head ys
        name
            = case xs of
                (n:s:_) -> Just n
                (s:_)   -> Nothing
        rsrc = case ys of
                (s:_:_) -> Just $ last ys
                _       -> Nothing
    in (name,server,rsrc)


textHostName = fmap Text.pack BSD.getHostName

localJID user resource = do
    hostname <- textHostName
    return $ user <> "@" <> hostname <> "/" <> resource

data ClientState = ClientState
    { clientResource :: Text
    , clientUser :: Text
    , clientPid :: Maybe ProcessID
    }

data PresenceState = PresenceState
    { clients :: TVar (Map ConnectionKey ClientState)
    , associatedPeers :: TVar (Map SockAddr ())
    , server :: TMVar XMPPServer
    , writeTos :: TVar (Map ConnectionKey (TChan Stanza))
    }


make6mapped4 addr@(SockAddrInet6 {}) = addr
make6mapped4 addr@(SockAddrInet port a) = SockAddrInet6 port 0 (0,0,0xFFFF,fromBE32 a) 0

resolvePeer :: Text -> IO [SockAddr]
resolvePeer addrtext = do
        fmap (map $ make6mapped4 . addrAddress) $ 
            getAddrInfo (Just $ defaultHints { addrFlags = [ AI_CANONNAME, AI_V4MAPPED ]})
                    (Just $ Text.unpack addrtext)
                    (Just "5269")


getConsolePids :: PresenceState -> IO [(Text,ProcessID)]
getConsolePids state = do
    us <- UTmp.users
    return $ map (\(_,tty,pid)->(lazyByteStringToText tty,pid)) us

lazyByteStringToText = (foldr1 (<>) . map Text.decodeUtf8 . L.toChunks)
textToLazyByteString s = L.fromChunks [Text.encodeUtf8 s]

identifyTTY' ttypids uid inode = ttypid
 where ttypids' = map (\(tty,pid)->(L.fromChunks [Text.encodeUtf8 tty], pid)) ttypids
       ttypid = fmap textify $ identifyTTY ttypids' uid inode
       textify (tty,pid) = (fmap lazyByteStringToText tty, pid)

chooseResourceName state k addr desired = do
    muid <- getLocalPeerCred' addr
    (mtty,pid) <- getTTYandPID muid
    user <- getJabberUserForId muid
    let client = ClientState { clientResource = maybe "fallback" id mtty
                             , clientUser = user
                             , clientPid = pid }

    atomically $
        modifyTVar' (clients state) $ Map.insert k client
    localJID (clientUser client) (clientResource client)
    
  where
    getTTYandPID muid = do
        -- us <- fmap (map (second fst) . Map.toList) . readTVarIO $ activeUsers state
        ttypids <- getConsolePids state
        -- let tailOf3 ((_,a),b) = (a,b)
        (t,pid) <- case muid of
                      Just (uid,inode) -> identifyTTY' ttypids uid inode
                      Nothing          -> return (Nothing,Nothing)
        let rsc = t `mplus` fmap ( ("pid."<>) . Text.pack . show ) pid
        return (rsc,pid)

    getJabberUserForId muid =
             maybe (return "nobody") 
                   (\(uid,_) -> 
                        handle (\(SomeException _) -> 
                                return . (<> "uid.") . Text.pack . show $ uid)
                          $ do
                             user <- fmap userName $ getUserEntryForID uid
                             return (Text.pack user)
                        )
                   muid

forClient state k fallback f = do
    mclient <- atomically $ do
        cs <- readTVar (clients state)
        return $ Map.lookup k cs
    maybe fallback f mclient

tellClientHisName state k = forClient state k fallback go
 where
    fallback  = localJID "nobody" "fallback"
    go client = localJID (clientUser client) (clientResource client)

toMapUnit xs = Map.fromList $ map (,()) xs

resolveAllPeers hosts = fmap (toMapUnit . concat) $ mapM (fmap (take 1) . resolvePeer) hosts

rosterGetStuff what state k = forClient state k (return [])
    $ \client -> do
    jids <- 
        fmap (map lazyByteStringToText)
            $ what (textToLazyByteString $ clientUser client)
    let hosts = map ((\(_,h,_)->h) . splitJID) jids
    addrs <- resolveAllPeers hosts 
    peers <- atomically $ readTVar (associatedPeers state)
    addrs <- return $ addrs `Map.difference` peers
    sv <- atomically $ takeTMVar $ server state
    forM_ (Map.keys addrs) $ \addr -> do
        putStrLn $ "new addr: "++show addr
        addPeer sv addr
    atomically $ do
        writeTVar (associatedPeers state) (addrs `Map.union` peers)
        putTMVar (server state) sv
    return jids

rosterGetBuddies :: PresenceState -> ConnectionKey -> IO [Text]
rosterGetBuddies = rosterGetStuff ConfigFiles.getBuddies
rosterGetSolicited = rosterGetStuff ConfigFiles.getSolicited
rosterGetOthers = rosterGetStuff ConfigFiles.getOthers
rosterGetSubscribers = rosterGetStuff ConfigFiles.getSubscribers

newConn state k outchan =
    atomically $ modifyTVar' (writeTos state) $ Map.insert k outchan

eofConn state k = atomically $ modifyTVar' (writeTos state) $ Map.delete k

rewriteJIDForPeer :: Text -> IO (Maybe (Text,SockAddr))
rewriteJIDForPeer jid = do
    let (n,h,r) = splitJID jid
    maddr <- fmap listToMaybe $ resolvePeer h
    return $ flip fmap maddr $ \addr ->
        let h' = addrToText addr
            to' = unsplitJID (n,h',r)
        in (to',addr)

deliverMessage state fail msg = do
    mto <- do
        flip (maybe $ return Nothing) (stanzaTo msg) $ \to -> do
        rewriteJIDForPeer to
    flip (maybe fail) mto $ \(to',addr) -> do
    from' <- do
        flip (maybe $ return Nothing) (stanzaFrom msg) $ \from -> do
        m <- rewriteJIDForPeer from
        return $ fmap fst m
    let k = PeerKey addr
    chans <- atomically $ readTVar (writeTos state)
    flip (maybe fail) (Map.lookup k chans) $ \chan -> do
    sendModifiedStanzaToPeer (msg { stanzaTo=Just to', stanzaFrom=from' }) chan

main = runResourceT $ do
    state <- liftIO . atomically $ do
        clients <- newTVar Map.empty
        associatedPeers <- newTVar Map.empty
        xmpp <- newEmptyTMVar
        writeTos <- newTVar Map.empty
        return PresenceState
            { clients = clients
            , associatedPeers = associatedPeers
            , writeTos = writeTos
            , server = xmpp
            }
    sv <- xmppServer
        XMPPServerParameters
        { xmppChooseResourceName = chooseResourceName state
        , xmppTellClientHisName = tellClientHisName state
        , xmppTellMyNameToClient = textHostName
        , xmppTellMyNameToPeer = \addr -> return $ addrToText addr
        , xmppTellPeerHisName = return . peerKeyToText
        , xmppTellClientNameOfPeer = peerKeyToResolvedName
        , xmppNewConnection = newConn state
        , xmppEOF = eofConn state
        , xmppRosterBuddies = rosterGetBuddies state
        , xmppRosterSubscribers = rosterGetSubscribers state
        , xmppRosterSolicited = rosterGetSolicited state
        , xmppRosterOthers = rosterGetOthers state
        , xmppSubscribeToRoster = \k -> return ()
        -- , xmppLookupClientJID = \k -> return $ "nobody@" <> hostname <> "/tty666"
        , {- xmppDeliverMessage = \fail msg -> do
            let msgs = msgLangMap (stanzaType msg)
                body = fmap (maybe "" id . msgBody . snd) $ take 1 msgs
            when (not $ null body) $ do
            Text.putStrLn $ "MESSAGE " <> head body
            return ()
          -}
          xmppDeliverMessage = deliverMessage state
        , xmppInformClientPresence = \k stanza -> return ()
        }
    liftIO $ do
    atomically $ putTMVar (server state) sv

    quitVar <- newEmptyTMVarIO
    installHandler sigTERM (CatchOnce (atomically $ putTMVar quitVar True)) Nothing
    installHandler sigINT (CatchOnce (atomically $ putTMVar quitVar True)) Nothing
    quitMessage <- atomically $ takeTMVar quitVar

    putStrLn "goodbye."
    return ()