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.Trans
import Control.Monad.IO.Class (MonadIO, liftIO)
import Network.Socket
    ( addrAddress
    , getAddrInfo
    , defaultHints
    , addrFlags
    , AddrInfoFlag(AI_CANONNAME,AI_V4MAPPED,AI_NUMERICHOST)
    , SockAddr(..)
    )
import System.Endian (fromBE32)
import Data.List (nub)
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 (maybeToList,listToMaybe,mapMaybe)

import TraversableT
import UTmp (ProcessID,users)
import LocalPeerCred
import XMPPServer

type UserName = Text
type ResourceName = Text

unsplitJID :: (Maybe UserName,Text,Maybe ResourceName) -> Text
unsplitJID (n,h,r) = username <> h <> resource
 where
    username = maybe "" (<>"@") n
    resource = maybe "" ("/"<>) r

splitJID :: Text -> (Maybe UserName,Text,Maybe ResourceName)
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)

isPeerKey :: ConnectionKey -> Bool
isPeerKey k = case k of { PeerKey {} -> True ; _ -> False }

isClientKey :: ConnectionKey -> Bool
isClientKey k = case k of { ClientKey {} -> True ; _ -> False }

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
    , clientStatus :: TVar (Maybe Stanza)
    }


data LocalPresence = LocalPresence
    { networkClients :: Map ConnectionKey ClientState
    -- TODO: loginClients
    }

data RemotePresence = RemotePresence
    { resources :: Map Text ()
    -- , localSubscribers :: Map Text ()
    -- ^ subset of clientsByUser who should be
    -- notified about this presence.
    }



pcSingletonNetworkClient key client =
    LocalPresence
    { networkClients = Map.singleton key client
    } 

pcInsertNetworkClient key client pc =
    pc { networkClients = Map.insert key client (networkClients pc) }

pcRemoveNewtworkClient key pc = if pcIsEmpty pc' then Nothing
                                                 else Just pc'
 where
    pc' = pc { networkClients = Map.delete key (networkClients pc) }

pcIsEmpty pc = Map.null (networkClients pc)


data PresenceState = PresenceState
    { clients :: TVar (Map ConnectionKey ClientState)
    , clientsByUser :: TVar (Map Text LocalPresence)
    , remotesByPeer :: TVar (Map ConnectionKey
                                 (Map UserName
                                      RemotePresence))
    , associatedPeers :: TVar (Map SockAddr ())
    , server :: TMVar XMPPServer
    , keyToChan :: TVar (Map ConnectionKey Conn)
    }


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 $ strip_brackets addrtext)
                    (Just "5269")

strip_brackets s = 
    case Text.uncons s of
      Just ('[',t) -> Text.takeWhile (/=']') t
      _            -> s


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
    status <- atomically $ newTVar Nothing
    let client = ClientState { clientResource = maybe "fallback" id mtty
                             , clientUser = user
                             , clientPid = pid
                             , clientStatus = status }

    atomically $ do
        modifyTVar' (clients state) $ Map.insert k client
        modifyTVar' (clientsByUser state) $ flip Map.alter (clientUser client)
            $ \mb -> Just $ maybe (pcSingletonNetworkClient k client)
                                  (pcInsertNetworkClient k client)
                                  mb

    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 :: [Text] -> IO (Map SockAddr ())
resolveAllPeers hosts = fmap (toMapUnit . concat) $ Prelude.mapM (fmap (take 1) . resolvePeer) hosts

configText what u =
        fmap (map lazyByteStringToText)
            $ what (textToLazyByteString u)

rosterGetStuff
  :: (L.ByteString -> IO [L.ByteString])
     -> PresenceState -> ConnectionKey -> IO [Text]
rosterGetStuff what state k = forClient state k (return [])
    $ \client -> do
    jids <- configText what (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
    -- Grok peers to associate with from the roster:
    forM_ (Map.keys addrs) $ \addr -> do
        putStrLn $ "new addr: "++show addr
        addPeer sv addr
    -- Update local set of associated peers
    atomically $ do
        writeTVar (associatedPeers state) (addrs `Map.union` peers)
        putTMVar (server state) sv
    return jids

rosterGetBuddies :: PresenceState -> ConnectionKey -> IO [Text]
rosterGetBuddies state k = do
    buds <- rosterGetStuff ConfigFiles.getBuddies state k
    return buds

rosterGetSolicited = rosterGetStuff ConfigFiles.getSolicited
rosterGetOthers = rosterGetStuff ConfigFiles.getOthers
rosterGetSubscribers = rosterGetStuff ConfigFiles.getSubscribers

data Conn = Conn { connChan :: TChan Stanza
                 , auxAddr :: SockAddr }

textAdapter what u = fmap (map lazyByteStringToText)
                        $ what (textToLazyByteString u)

getBuddies' :: Text -> IO [Text]
getBuddies' = textAdapter ConfigFiles.getBuddies
getSolicited' :: Text -> IO [Text]
getSolicited' = textAdapter ConfigFiles.getSolicited

sendProbesAndSolicitations state k laddr chan = do
    -- get all buddies & solicited matching k for all users
    xs <- runTraversableT $ do
        cbu <- lift $ atomically $ readTVar $ clientsByUser state
        user <- liftT $ Map.keys cbu
        (isbud,getter) <- liftT [(True ,getBuddies'  )
                                ,(False,getSolicited')]
        bud  <- liftMT $ getter user
        let (u,h,r) = splitJID bud
        addr <- liftMT $ nub `fmap` resolvePeer h
        liftT $ guard (PeerKey addr == k)
        -- Note: Earlier I was tempted to do all the IO
        -- within the TraversableT monad.  That apparently
        -- is a bad idea.  Perhaps due to laziness and an
        -- unforced list?  Instead, we will return a list
        -- of (Bool,Text) for processing outside.
        return (isbud,u)
    -- XXX: The following O(n²) nub may be a little
    --      too onerous.
    forM_ (nub xs) $ \(isbud,u) -> do
        let make = if isbud then presenceProbe 
                            else presenceSolicitation
            toh = peerKeyToText k
            jid = unsplitJID (u,toh,Nothing)
            me = addrToText laddr
        stanza <- make me jid
        -- send probes for buddies, solicitations for solicited.
        putStrLn $ "probing "++show k++" for: " ++ show (isbud,jid)
        atomically $ writeTChan chan stanza
    -- reverse xs `seq` return ()

newConn state k addr outchan = do
    atomically $ modifyTVar' (keyToChan state)
        $ Map.insert k Conn { connChan = outchan
                            , auxAddr  = addr }
    when (isPeerKey k)
        $ sendProbesAndSolicitations state k addr outchan

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

rewriteJIDForClient1:: Text -> IO (Maybe ((Maybe Text,Text,Maybe Text),SockAddr))
rewriteJIDForClient1 jid = do
    let (n,h,r) = splitJID jid
    maddr <- fmap listToMaybe $ resolvePeer h
    flip (maybe $ return Nothing) maddr $ \addr -> do
        h' <- peerKeyToResolvedName (PeerKey addr)
        return $ Just ((n,h',r), addr)

parseAddress :: Text -> IO (Maybe SockAddr)
parseAddress addr_str = do
    info <- getAddrInfo (Just $ defaultHints { addrFlags = [ AI_NUMERICHOST ] }) 
                        (Just . Text.unpack $ addr_str) 
                        (Just "0")
    return . listToMaybe $ map addrAddress info

todo = error "Unimplemented"


-- | for example:  2001-db8-85a3-8d3-1319-8a2e-370-7348.ipv6-literal.net
ip6literal :: Text -> Text
ip6literal addr = Text.map dash addr <> ".ipv6-literal.net"
 where
    dash ':' = '-'
    dash x   = x

withPort (SockAddrInet _ a)      port = SockAddrInet (toEnum port) a
withPort (SockAddrInet6 _ a b c) port = SockAddrInet6 (toEnum port) a b c

-- | The given address is taken to be the local address for the socket this JID
-- came in on.  The returned JID parts are suitable for unsplitJID to create a
-- valid JID for communicating to a client.  The returned Bool is True when the
-- host part refers to this local host (i.e. it equals the given SockAddr).
rewriteJIDForClient :: SockAddr -> Text -> IO (Bool,(Maybe Text,Text,Maybe Text))
rewriteJIDForClient laddr jid = do
    let (n,h,r) = splitJID jid
    maddr <- parseAddress (strip_brackets h)
    flip (maybe $ return (False,(n,ip6literal h,r))) maddr $ \addr -> do
    let mine =  laddr `withPort` 0  == addr `withPort` 0
    h' <- if mine then textHostName
                  else peerKeyToResolvedName (PeerKey addr)
    return (mine,(n,h',r))

addrTextToKey h = do
    maddr <- parseAddress (strip_brackets h)
    return (fmap PeerKey maddr)

guardPortStrippedAddress h laddr = do
    maddr <- fmap (fmap (`withPort` 0)) $ parseAddress (strip_brackets h)
    let laddr' = laddr `withPort` 0
    return $ maddr >>= guard . (==laddr')


-- | Accepts a textual representation of a domainname
-- JID suitable for client connections, and returns the
-- coresponding ipv6 address JID suitable for peers paired
-- with a SockAddr with the address part of that JID in
-- binary form.  If no suitable address could be resolved
-- for the given name, Nothing is returned.
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 = 
    case stanzaOrigin msg of
      NetworkOrigin senderk@(ClientKey {}) _ -> do
        -- Case 1.  Client -> Peer
        mto <- do
            flip (maybe $ return Nothing) (stanzaTo msg) $ \to -> do
            rewriteJIDForPeer to
        flip (maybe fail {- reverse lookup failure -})
             mto
            $ \(to',addr) -> do
        let k = PeerKey addr
        chans <- atomically $ readTVar (keyToChan state)
        flip (maybe fail) (Map.lookup k chans) $ \(Conn { connChan=chan
                                                        , auxAddr=laddr }) -> do
        (n,r) <- forClient state senderk (return (Nothing,Nothing))
                        $ \c -> return (Just (clientUser c), Just (clientResource c))
        -- original 'from' address is discarded.
        let from' = unsplitJID (n,addrToText laddr,r)
        -- dup <- atomically $ cloneStanza (msg { stanzaTo=Just to', stanzaFrom=Just from' })
        let dup = (msg { stanzaTo=Just to', stanzaFrom=Just from' })
        sendModifiedStanzaToPeer dup chan
      NetworkOrigin senderk@(PeerKey {}) _ -> do
        key_to_chan <- atomically $ readTVar (keyToChan state)
        flip (maybe fail) (Map.lookup senderk key_to_chan)
            $ \(Conn { connChan=sender_chan
                     , auxAddr=laddr }) -> do
        flip (maybe fail) (stanzaTo msg) $ \to -> do
        (mine,(n,h,r)) <- rewriteJIDForClient laddr to
        if not mine then fail else do
        let to' = unsplitJID (n,h,r)
        from' <- do
            flip (maybe $ return Nothing) (stanzaFrom msg) $ \from -> do
            (_,trip) <- rewriteJIDForClient laddr from
            return . Just $ unsplitJID trip
        cmap <- atomically . readTVar $ clientsByUser state
        flip (maybe fail) n $ \n -> do
        flip (maybe fail) (Map.lookup n cmap) $ \presence_container -> do
        let ks = Map.keys (networkClients presence_container)
            chans = mapMaybe (flip Map.lookup key_to_chan) ks
        if null chans then fail else do
        forM_ chans $ \Conn { connChan=chan} -> do
        putStrLn $ "sending "++show (stanzaId msg)++" to clients "++show ks
        -- TODO: Cloning isn't really neccessary unless there are multiple
        -- destinations and we should probably transition to minimal cloning,
        -- or else we should distinguish between announcable stanzas and
        -- consumable stanzas and announcables use write-only broadcast
        -- channels that must be cloned in order to be consumed.
        -- For now, we are doing redundant cloning.
        dup <- cloneStanza (msg { stanzaTo=Just to'
                                , stanzaFrom=from' })
        sendModifiedStanzaToClient dup
                                  chan

informClientPresence state k stanza = do
    dup <- cloneStanza stanza
    atomically $ do
        mb <- fmap (Map.lookup k) $ readTVar (clients state)
        flip (maybe $ return ()) mb $ \cstate -> do
        writeTVar (clientStatus cstate) $ Just dup
    forClient state k (return ()) $ \client -> do
    jids <- configText ConfigFiles.getSubscribers (clientUser client)
    let hosts = map ((\(_,h,_)->h) . splitJID) jids
    addrs <- fmap Map.keys $ resolveAllPeers hosts 
    ktc <- atomically $ readTVar (keyToChan state)
    let connected = mapMaybe (flip Map.lookup ktc . PeerKey) addrs
    forM_ connected $ \con -> do
    let from' = unsplitJID ( Just $ clientUser client
                           , addrToText $ auxAddr con
                           , Just $ clientResource client)
    mto <- runTraversableT $ do
            to <- liftT $ stanzaTo stanza
            (to',_) <- liftMT $ rewriteJIDForPeer to
            return to'
    dup <- cloneStanza stanza
    sendModifiedStanzaToPeer dup { stanzaFrom = Just from'
                                 , stanzaTo = mto }
                             (connChan con)

informPeerPresence state k stanza = do
    -- Presence must indicate full JID with resource...
    putStrLn $ "xmppInformPeerPresence checking from address..."
    flip (maybe $ return ()) (stanzaFrom stanza) $ \from -> do
    let (muser,h,mresource) = splitJID from
    flip (maybe $ return ()) mresource       $ \resource -> do
    flip (maybe $ return ()) muser           $ \user -> do
    
    clients <- atomically $ do
        -- TODO: Store the stanza
        -- For now, all clients:
        -- (TODO: interested/authorized clients only.)
        ktc <- readTVar (keyToChan state)
        runTraversableT $ do
        (ck,client) <- liftMT $ fmap Map.toList $ readTVar (clients state)
        con <- liftMaybe $ Map.lookup ck ktc 
        return (ck,con,client)
    putStrLn $ "xmppInformPeerPresence (length clients="++show (length clients)++")"
    forM_ clients $ \(ck,con,client) -> do
        from' <- do
            let ClientKey laddr = ck
            (_,trip) <- rewriteJIDForClient laddr from
            return (unsplitJID trip)
        putStrLn $ "sending to client: " ++ show (stanzaType stanza)
        dup <- cloneStanza stanza
        sendModifiedStanzaToClient (dup { stanzaFrom=Just from' })
                                   (connChan con)

answerProbe state k stanza chan = do
    putStrLn $ "answerProbe! " ++ show (stanzaType stanza)
    ktc <- atomically $ readTVar (keyToChan state)
    replies <- runTraversableT $ do
        to <- liftMaybe $ stanzaTo stanza
        conn <- liftMaybe $ Map.lookup k ktc
        let (mu,h,_) = splitJID to -- TODO: currently resource-id is ignored on presence
                                   --  probes.  Is this correct? Check the spec.
        liftIOMaybe $ guardPortStrippedAddress h (auxAddr conn)
        u <- liftMaybe mu
        cbu <- lift . atomically $ readTVar (clientsByUser state)
        lpres <- liftMaybe $ Map.lookup u cbu
        clientState <- liftT $ Map.elems (networkClients lpres)
        stanza <- liftIOMaybe $ atomically (readTVar (clientStatus clientState))
        stanza <- lift $ cloneStanza stanza
        let jid = unsplitJID (Just $ clientUser clientState
                             , ch
                             ,Just $ clientResource clientState)
            ch = addrToText (auxAddr conn)
        return stanza { stanzaFrom = Just jid }
    forM_ replies $ \reply -> do
        sendModifiedStanzaToPeer reply chan
    -- TODO: if null replies, send offline message

main = runResourceT $ do
    state <- liftIO . atomically $ do
        clients <- newTVar Map.empty
        clientsByUser <- newTVar Map.empty
        remotesByPeer <- newTVar Map.empty
        associatedPeers <- newTVar Map.empty
        xmpp <- newEmptyTMVar
        keyToChan <- newTVar Map.empty
        return PresenceState
            { clients = clients
            , clientsByUser = clientsByUser
            , remotesByPeer = remotesByPeer
            , associatedPeers = associatedPeers
            , keyToChan = keyToChan
            , 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 ()
        , xmppDeliverMessage = deliverMessage state
        , xmppInformClientPresence = informClientPresence state
        , xmppInformPeerPresence = informPeerPresence state
        , xmppAnswerProbe = answerProbe state
        }
    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 ()