path: root/src/Network/Tox/Crypto/Handlers.hs

{-# LANGUAGE NamedFieldPuns #-}
{-# LANGUAGE TupleSections #-}
module Network.Tox.Crypto.Handlers where

import Network.Tox.Crypto.Transport
import Network.Tox.DHT.Transport (Cookie(..),CookieData(..))
import Crypto.Tox
import Control.Concurrent.STM
import Network.Address
import qualified Data.Map.Strict as Map
import Crypto.Hash
import Control.Applicative
import Control.Monad
import Data.Time.Clock.POSIX
import qualified Data.ByteString  as B
import Control.Lens
import Data.Function
import Data.Serialize             as S
import Data.Word
import GHC.Conc (unsafeIOToSTM)
import qualified Data.Set         as Set

-- util, todo: move to another module
maybeToEither :: Maybe b -> Either String b
maybeToEither (Just x) = Right x
maybeToEither Nothing  = Left "maybeToEither"

data NetCryptoSessionStatus = Unaccepted | Accepted | Confirmed
    deriving (Eq,Ord,Show,Enum)


type IOHook addr x = addr -> x -> IO (Maybe (x -> x))
type NetCryptoHook = IOHook NetCryptoSession CryptoData


data NetCryptoSession = NCrypto { ncState         :: TVar NetCryptoSessionStatus
                                , ncTheirBaseNonce :: TVar Nonce24 -- base nonce + packet number
                                , ncMyPacketNonce :: TVar Nonce24 -- base nonce + packet number
                                , ncHandShake     :: TVar (Maybe (Handshake Encrypted))
                                , ncCookie        :: TVar (Maybe Cookie)
                                , ncTheirDHTKey   :: PublicKey
                                , ncTheirSessionPublic :: Maybe PublicKey
                                , ncSessionSecret :: SecretKey
                                , ncSockAddr      :: SockAddr
                                , ncHooks         :: TVar (Map.Map MessageType [NetCryptoHook])
                                , ncUnrecognizedHook :: TVar (MessageType -> NetCryptoHook)
                                , ncAllSessions :: NetCryptoSessions -- ^ may be needed if one net-crypto session
                                                                     --   needs to possibly start another, as is
                                                                     --   the case in group chats
                                , ncGroups :: TVar (Map.Map GroupChatId (Set.Set SockAddr))
                                }

data NetCryptoSessions = NCSessions { netCryptoSessions :: TVar (Map.Map SockAddr NetCryptoSession)
                                    , transportCrypto :: TransportCrypto
                                    , defaultHooks :: Map.Map MessageType [NetCryptoHook]
                                    , defaultUnrecognizedHook :: MessageType -> NetCryptoHook
                                    }

newSessionsState :: TransportCrypto -> (MessageType -> NetCryptoHook) -> Map.Map MessageType [NetCryptoHook] -> IO NetCryptoSessions
newSessionsState crypto unrechook hooks = do
    x <- atomically $ newTVar Map.empty
    return NCSessions { netCryptoSessions = x
                      , transportCrypto = crypto
                      , defaultHooks = hooks
                      , defaultUnrecognizedHook = unrechook
                      }

data HandshakeParams
        = HParam
            { hpTheirBaseNonce :: Maybe Nonce24 -- ignore and generate your own
            , hpOtherCookie :: Maybe Cookie
            , hpTheirSessionKeyPublic :: PublicKey
            , hpMySecretKey :: SecretKey
            , hpCookieRemotePubkey :: PublicKey
            , hpCookieRemoteDhtkey :: PublicKey
            }
newHandShakeData :: TransportCrypto -> HandshakeParams -> HandshakeData
newHandShakeData = error "todo"

-- | called when we recieve a crypto handshake with valid cookie
freshCryptoSession :: NetCryptoSessions -> SockAddr -> HandshakeParams -> IO ()
freshCryptoSession  sessions
                    addr
                    hp@(HParam
                              { hpTheirBaseNonce = Just theirBaseNonce
                              , hpOtherCookie    = Just otherCookie
                              , hpTheirSessionKeyPublic = theirSessionKey
                              , hpMySecretKey        = key
                              , hpCookieRemotePubkey = remotePublicKey
                              , hpCookieRemoteDhtkey = remoteDhtPublicKey
                              }) = do
    let crypto = transportCrypto sessions
        allsessions = netCryptoSessions sessions
    ncState0 <- atomically $ newTVar Accepted
    ncTheirBaseNonce0 <- atomically $ newTVar theirBaseNonce
    n24 <- atomically $ transportNewNonce crypto
    state <- lookupSharedSecret crypto key remoteDhtPublicKey n24
    let myhandshakeData = newHandShakeData crypto hp
        plain = encodePlain myhandshakeData
        encrypted = encrypt state plain
        myhandshake = Handshake { handshakeCookie = otherCookie
                                , handshakeNonce = n24
                                , handshakeData = encrypted
                                }
    ncMyPacketNonce0 <- atomically $ newTVar (baseNonce myhandshakeData)
    ncHandShake0 <- atomically $ newTVar (Just myhandshake)
    cookie0 <- atomically $ newTVar (Just otherCookie)
    newsession <- generateSecretKey
    ncHooks0 <- atomically $ newTVar (defaultHooks sessions)
    ncUnrecognizedHook0 <- atomically $ newTVar (defaultUnrecognizedHook sessions)
    ncGroups0 <- atomically $ newTVar (Map.empty)
    let netCryptoSession =
            NCrypto { ncState         = ncState0
                    , ncTheirBaseNonce= ncTheirBaseNonce0
                    , ncMyPacketNonce   = ncMyPacketNonce0
                    , ncHandShake     = ncHandShake0
                    , ncCookie        = cookie0
                    , ncTheirDHTKey   = remoteDhtPublicKey
                    , ncTheirSessionPublic = Just theirSessionKey
                    , ncSessionSecret = newsession
                    , ncSockAddr      = addr
                    , ncHooks         = ncHooks0
                    , ncUnrecognizedHook = ncUnrecognizedHook0
                    , ncAllSessions = sessions
                    , ncGroups = ncGroups0
                    }
    atomically $ modifyTVar allsessions (Map.insert addr netCryptoSession)

-- | Called when we get a handshake, but there's already a session entry.
updateCryptoSession :: NetCryptoSessions -> SockAddr -> HandshakeParams -> NetCryptoSession -> IO ()
updateCryptoSession sessions addr hp session = do
    ncState0 <- atomically $ readTVar (ncState session)
    ncTheirBaseNonce0 <- atomically $ readTVar (ncTheirBaseNonce session)
    if (ncState0 >= Accepted)
        -- If the nonce in the handshake and the dht key are both the same as
        -- the ones we have saved, assume we already handled this and this is a
        -- duplicate handshake packet, otherwise disregard everything, and
        -- refresh all state.
        --
     then when (  Just ncTheirBaseNonce0 /= hpTheirBaseNonce hp
               || ncTheirDHTKey session /= hpCookieRemoteDhtkey hp
               ) $ freshCryptoSession sessions addr hp
     else if ( Just ncTheirBaseNonce0 /= hpTheirBaseNonce hp)
            then freshCryptoSession sessions addr hp -- basenonce mismatch, trigger refresh
            else atomically $ writeTVar (ncState session) Accepted


cryptoNetHandler :: NetCryptoSessions -> SockAddr -> NetCrypto -> IO (Maybe (NetCrypto -> NetCrypto))
cryptoNetHandler sessions addr (NetHandshake (Handshake (Cookie n24 ecookie) nonce24 encrypted)) = do
    -- Handle Handshake Message
    let crypto = transportCrypto sessions
        allsessions = netCryptoSessions sessions
        anyRight [] f     = return $ Left "missing key"
        anyRight (x:xs) f = f x >>= either (const $ anyRight xs f) (return . Right)
    seckeys <- map fst <$> atomically (readTVar (userKeys crypto))
    symkey <- atomically $ transportSymmetric crypto
    now <- getPOSIXTime
    lr <- fmap join . sequence $ do -- Either Monad
            (CookieData cookieTime remotePubkey remoteDhtkey) <- (decodePlain =<< decryptSymmetric symkey n24 ecookie)
            Right $ do -- IO Monad
            decrypted <- anyRight seckeys $ \key -> do
                            secret <- lookupSharedSecret crypto key remotePubkey nonce24
                            return $ (key,) <$> (decodePlain =<< decrypt secret encrypted)
            return $ do -- Either Monad
            (key,HandshakeData { baseNonce, sessionKey, cookieHash, otherCookie }) <- decrypted
            -- check cookie time < 15 seconds ago
            guard (now - fromIntegral cookieTime < 15)
            -- cookie hash is valid? sha512 of ecookie
            let hinit = hashInit
                hctx = hashUpdate hinit n24
                hctx' = hashUpdate hctx ecookie
                digest = hashFinalize hctx'
            guard (cookieHash == digest)
            -- known friend?
            -- todo
            return
             HParam
                { hpTheirBaseNonce = Just baseNonce
                , hpOtherCookie    = Just otherCookie
                , hpTheirSessionKeyPublic =  sessionKey
                , hpMySecretKey        = key
                , hpCookieRemotePubkey = remotePubkey
                , hpCookieRemoteDhtkey = remoteDhtkey
                }
    case lr of
        Left _ -> return ()
        Right hp@(HParam
                      { hpTheirBaseNonce = Just theirBaseNonce
                      , hpOtherCookie    = Just otherCookie
                      , hpTheirSessionKeyPublic = theirSessionKey
                      , hpMySecretKey        = key
                      , hpCookieRemotePubkey = remotePublicKey
                      , hpCookieRemoteDhtkey = remoteDhtPublicKey
                      }) -> do
            sessionsmap <- atomically $ readTVar allsessions
            -- Do a lookup, so we can handle the update case differently
            case Map.lookup addr sessionsmap of
                Nothing -> freshCryptoSession sessions addr hp -- create new session
                Just session -> updateCryptoSession sessions addr hp session -- update existing session
    return Nothing


cryptoNetHandler sessions addr (NetCrypto (CryptoPacket nonce16 encrypted)) = do
    let crypto = transportCrypto sessions
        allsessions = netCryptoSessions sessions
    sessionsmap <- atomically $ readTVar allsessions
    -- Handle Encrypted Message
    case Map.lookup addr sessionsmap of
        Nothing -> return Nothing -- drop packet, we have no session
        Just session@(NCrypto {ncState, ncHooks,ncSessionSecret,ncTheirSessionPublic,ncTheirBaseNonce}) -> do
            theirBaseNonce <- atomically $ readTVar ncTheirBaseNonce
            -- Try to decrypt message
            let diff :: Word16
                diff = nonce16 - fromIntegral (last2Bytes theirBaseNonce) -- truncating to Word16
            tempNonce <- addtoNonce24 theirBaseNonce (fromIntegral diff) -- expanding to Word
            lr <- fmap join $ sequence $ do -- Either Monad --
                        pubkey <- maybeToEither ncTheirSessionPublic
                        Right $ do -- IO Monad
                        secret <- lookupSharedSecret crypto ncSessionSecret pubkey tempNonce
                        return $ decodePlain =<< decrypt secret encrypted
            case lr of
                Left _ -> return Nothing -- decryption failed, ignore packet
                Right cd@(CryptoData {bufferStart, bufferEnd, bufferData=cm}) -> do -- decryption succeeded,
                    -- TODO: Why do I need bufferStart & bufferEnd?
                    --
                    --   buffer_start = highest packet number handled + 1
                    --                , recvbuffers buffer_start
                    --
                    --   bufferEnd = sendbuffer buffer_end if lossy, otherwise packet number
                    -- update ncTheirBaseNonce if necessary
                    when (diff > 2 * dATA_NUM_THRESHOLD)$
                        atomically $ do
                            y <- readTVar ncTheirBaseNonce
                            -- all because Storable forces IO...
                            x <- unsafeIOToSTM $ addtoNonce24 y (fromIntegral dATA_NUM_THRESHOLD)
                            writeTVar ncTheirBaseNonce y
                    -- then set session confirmed,
                    atomically $ writeTVar ncState Confirmed
                    hookmap <- atomically $ readTVar ncHooks
                    -- run hook
                    flip fix (cd,hookmap) $ \lookupAgain (cd,hookmap) -> do
                        let msgTyp = cd ^. messageType
                        case Map.lookup msgTyp hookmap of
                                Nothing -> do -- no recognizing hook, run ncUnrecognizedHook0, loopAgain on result
                                             unrecognize <- atomically $ readTVar (ncUnrecognizedHook session)
                                             mbConsume <- unrecognize msgTyp session cd
                                             case mbConsume of
                                                Just f -> do
                                                           -- ncUnrecognizedHook0 may have updated the hookmap
                                                           hookmap' <- atomically $ readTVar ncHooks
                                                           lookupAgain (f cd,hookmap')
                                                Nothing -> return Nothing
                                Just hooks -> flip fix (hooks,cd,msgTyp) $ \loop (hooks,cd,typ) -> do
                                                let _ = cd :: CryptoData
                                                case (hooks,cd) of
                                                    ([],_) ->  return Nothing
                                                    (hook:more,cd) -> do
                                                        r <- hook session cd :: IO (Maybe (CryptoData -> CryptoData))
                                                        case r of
                                                            Just f -> let newcd = f cd
                                                                          newtyp = newcd ^. messageType
                                                                          in if newtyp == typ then loop (more,newcd,newtyp)
                                                                                              else lookupAgain (newcd,hookmap)
                                                            Nothing -> return Nothing -- message consumed
    where
        last2Bytes :: Nonce24 -> Word
        last2Bytes (Nonce24 bs) = case S.decode (B.drop 22 bs) of
                                    Right n -> n
                                    _ -> error "unreachable-last2Bytes"
        dATA_NUM_THRESHOLD = 21845 -- = 65535 / 3

-- | handles nothing
defaultCryptoDataHooks :: Map.Map MessageType [NetCryptoHook]
defaultCryptoDataHooks = Map.empty

-- | discards all unrecognized packets
defaultUnRecHook :: MessageType -> NetCryptoHook
defaultUnRecHook _ _ _ = return Nothing

-- | use to add a single hook to a specific session.
addCryptoDataHook1 :: Map.Map MessageType [NetCryptoHook] -> MessageType -> NetCryptoHook -> Map.Map MessageType [NetCryptoHook]
addCryptoDataHook1 mp typ hook = case Map.lookup typ mp of
                                    Nothing -> Map.insert typ [hook] mp
                                    Just hooks -> Map.insert typ (hook:hooks) mp