-- | This module implements the lossless Tox session protocol.
{-# LANGUAGE GADTs         #-}
{-# LANGUAGE TupleSections #-}
module Network.Tox.Session
    ( SessionParams(..)
    , SessionKey
    , Session(..)
    , sTheirUserKey
    , sClose
    , handshakeH
    ) where

import Control.Concurrent.STM
import Control.Monad
import Control.Exception
import Data.Dependent.Sum
import Data.Functor.Identity
import Data.Word
import Network.Socket         (SockAddr)

import Crypto.Tox
import Data.PacketBuffer            (PacketInboundEvent (..))
import qualified Data.Tox.DHT.Multi as Multi
import Data.Tox.Msg
import DebugTag
import DPut
import Network.Lossless
import Network.QueryResponse
import Network.SessionTransports
import Network.Tox.Crypto.Transport
import Network.Tox.DHT.Transport    (Cookie (..), CookieData (..), key2id,
                                     longTermKey)
import Network.Tox.Handshake
import Network.Tox.TCP              (ViaRelay (..))

-- | Alias for 'SecretKey' to document that it is used as the temporary Tox
-- session key corresponding to the 'PublicKey' we sent in the handshake.
type SessionKey = SecretKey

-- | These inputs to 'handshakeH' indicate how to respond to handshakes, how to
-- assign packets to sessions, and what to do with established sessions after
-- they are made lossless by queuing packets and appending sequence numbers.
data SessionParams = SessionParams
    { -- | The database of secret keys necessary to encrypt handshake packets.
      spCrypto           :: TransportCrypto
      -- | This is used to create sessions and dispatch packets to them.
    , spSessions         :: Sessions (CryptoPacket Encrypted)
      -- | This method returns the session information corresponding to the
      -- cookie pair for the remote address.  If no handshake was sent, this
      -- should send one immediately.  It should return 'Nothing' if anything
      -- goes wrong.
    , spGetSentHandshake :: SecretKey -> Multi.SessionAddress
                                      -> Cookie Identity
                                      -> Cookie Encrypted
                                      -> IO (Maybe (Multi.SessionAddress, (SessionKey, HandshakeData)))
      -- | This method is invoked on each new session and is responsible for
      -- launching any threads necessary to keep the session alive.
    , spOnNewSession     :: Session -> IO ()
    }

-- | After a session is established, this information is given to the
-- 'spOnNewSession' callback.
data Session = Session
    { -- | This is the secret user (toxid) key that corresponds to the
      -- local-end of this session.
      sOurKey            :: SecretKey
      -- | The remote address for this session. (Not unique, see 'sSessionID').
    , sTheirAddr         :: Multi.SessionAddress
      -- | The information we sent in the handshake for this session.
    , sSentHandshake     :: HandshakeData
      -- | The information we received in a handshake for this session.
    , sReceivedHandshake :: Handshake Identity
      -- | This method can be used to trigger packets to be re-sent given a
      -- list of their sequence numbers.  It should be used when the remote end
      -- indicates they lost packets.
    , sResendPackets     :: [Word32] -> IO ()
      -- | This list of sequence numbers should be periodically polled and if
      -- it is not empty, we should request they re-send these packets.  For
      -- convenience, a lower bound for the numbers in the list is also
      -- returned.  Suggested polling interval: a few seconds.
    , sMissingInbound    :: IO ([Word32],Word32)
      -- | A lossless transport for sending and receiving packets in this
      -- session.  It is up to the caller to spawn the await-loop to handle
      -- inbound packets.
    , sTransport         :: Transport String () CryptoMessage
      -- | A unique small integer that identifies this session for as long as
      -- it is established.
    , sSessionID         :: Int
    }

-- | Helper to obtain the remote ToxID key from the locally-issued cookie
-- associated with the session.
sTheirUserKey :: Session -> PublicKey
sTheirUserKey s = longTermKey $ runIdentity cookie
 where
    Cookie _ cookie = handshakeCookie (sReceivedHandshake s)

-- | Helper to close the 'Transport' associated with a session.
sClose :: Session -> IO ()
sClose s = do
    sendMessage (sTransport s) () (Pkt KillPacket ==> ())
    closeTransport (sTransport s)


-- | Call this whenever a new handshake arrives so that a session is
-- negotiated.  It always returns Nothing which makes it convenient to use with
-- 'Network.QueryResponse.addHandler'.
handshakeH :: SessionParams
              -> Arrival err Multi.SessionAddress (Handshake Encrypted)
              -> STM (Arrival err Multi.SessionAddress (Handshake Encrypted), IO ())
handshakeH sp (Arrival saddr handshake) = return $ (,) Discarded $ do
    decryptHandshake (spCrypto sp) handshake
        >>= either (\err -> return ())
                   (uncurry $ plainHandshakeH sp saddr)
handshakeH _ m = return (m, return ())


plainHandshakeH :: SessionParams
                   -> Multi.SessionAddress
                   -> SecretKey
                   -> Handshake Identity
                   -> IO ()
plainHandshakeH sp saddr0 skey handshake = do
    let hd = runIdentity $ handshakeData handshake
        prelude = show saddr0 ++ " --> "
    dput XNetCrypto $ unlines $ map (prelude ++)
        [ "handshake: auth=" ++ show (handshakeCookie handshake)
        , "         : issuing=" ++ show (otherCookie hd)
        , "         : baseNonce=" ++ show (baseNonce hd)
        ]
    sent <- spGetSentHandshake sp skey saddr0 (handshakeCookie handshake) (otherCookie hd)
    dput XNetCrypto $ " <-- (cached) handshake baseNonce " ++ show (fmap (baseNonce . snd . snd) sent)
    sent' <- case sent of
        Just (Multi.SessionTCP :=> Identity (ViaRelay Nothing _ _),_) -> do
            dput XNetCrypto $ "Rejecting OOB netcrypto session because it is incompatible with toxcore."
            return Nothing
        _ -> return sent
    forM_ sent' $ \(saddr, (hd_skey,hd_sent)) -> do
    let Cookie _ (Identity CookieData{ longTermKey = them }) = handshakeCookie handshake
    sk <- SessionKeys (spCrypto sp)
                      hd_skey
                      (sessionKey hd)
                      <$> atomically (newTVar $ baseNonce hd)
                      <*> atomically (newTVar $ baseNonce hd_sent)
    let addr_lbl = Multi.showSessionAddr saddr
    m <- newSession (spSessions sp) (\() p -> return p) (\_ -> decryptPacket sk addr_lbl) saddr
    dput XNetCrypto $ prelude ++ "plainHandshakeH: session " ++ maybe "Nothing" (const "Just") m
    forM_ m $ \(sid, t) -> do
        (t2,resend,getMissing)
            <- lossless (take 8 (showKey256 them) ++ "." ++ Multi.showSessionAddr saddr)
                        (\cp a -> return $ fmap (,a) $ checkLossless $ runIdentity $ pktData cp)
                        (\seqno p@(Pkt m :=> _) _ -> do
                            y <- encryptPacket sk addr_lbl $ bookKeeping seqno p
                            return OutgoingInfo
                                { oIsLossy         = lossyness m == Lossy
                                , oEncoded         = y
                                , oHandleException = Just $ \e -> do
                                        dput XUnexpected $ unlines
                                            [ "<-- " ++ show e
                                            , "<--   while sending " ++ show (seqno,p) ]
                                        throwIO e
                                })
                        ()
                        t
        let _ = t :: TransportA String () (CryptoPacket Identity) (CryptoPacket Encrypted)
            _ = t2 :: Transport String () CryptoMessage
        sendMessage t2 () $ (Pkt ONLINE ==> ())
        spOnNewSession sp Session
            { sOurKey            = skey
            , sTheirAddr         = saddr
            , sSentHandshake     = hd_sent
            , sReceivedHandshake = handshake
            , sResendPackets     = resend
            , sMissingInbound    = getMissing
            , sTransport         = t2
            , sSessionID         = sid
            }
    return ()


-- | The per-session nonce and key state maintained by 'decryptPacket' and
-- 'encryptPacket'.
data SessionKeys = SessionKeys
    { skCrypto        :: TransportCrypto -- ^ Cache of shared-secrets.
    , skMe            :: SessionKey      -- ^ My session key
    , skThem          :: PublicKey       -- ^ Their session key
    , skNonceIncoming :: TVar Nonce24    -- ^ +21845 when a threshold is reached.
    , skNonceOutgoing :: TVar Nonce24    -- ^ +1 on every packet
    }

-- | Decrypt an inbound session packet and update the nonce for the next one.
decryptPacket :: SessionKeys -> String -> CryptoPacket Encrypted -> IO (Maybe (CryptoPacket Identity, ()))
decryptPacket sk lbl (CryptoPacket n16 ciphered) = do
    (n24,δ) <- atomically $ do
        n <- readTVar (skNonceIncoming sk)
        let δ = n16 - nonce24ToWord16 n
        return ( n `addtoNonce24` fromIntegral δ, δ )
    secret <- lookupSharedSecret (skCrypto sk) (skMe sk) (skThem sk) n24
    case decodePlain =<< decrypt secret ciphered of
        Left e  -> return Nothing
        Right x -> do
            when ( δ  > 43690 )
                $ atomically $ writeTVar (skNonceIncoming sk) (n24 `addtoNonce24` 21845)

            do let them = key2id $ skThem sk
                   CryptoData ack seqno _ = x
                   cm = decodeRawCryptoMsg x
               dput XNetCrypto $ unwords [lbl,"-->",show (msgID cm),show (n24,ack,seqno)]

            return $ Just ( CryptoPacket n16 (pure x), () )

-- | Encrypt an outbound session packet and update the nonce for the next one.
encryptPacket :: SessionKeys -> String -> CryptoData -> IO (CryptoPacket Encrypted)
encryptPacket sk lbl plain = do
    n24 <- atomically $ do
        n24 <- readTVar (skNonceOutgoing sk)
        modifyTVar' (skNonceOutgoing sk) incrementNonce24
        return n24
    secret <- lookupSharedSecret (skCrypto sk) (skMe sk) (skThem sk) n24
    let ciphered = encrypt secret $ encodePlain $ plain

    do let them = key2id $ skThem sk
           CryptoData ack seqno cm = plain
       dput XNetCrypto $ unwords [lbl,"<--",show (msgID cm),show (n24,ack,seqno)]

    return $ CryptoPacket (nonce24ToWord16 n24) ciphered


-- | Add sequence information to an outbound packet.
--
-- From spec.md:
--
--   Data in the encrypted packets:
--
--    [our recvbuffers buffer_start, (highest packet number handled + 1), (big endian)]
--    [uint32_t packet number if lossless, sendbuffer buffer_end if lossy, (big endian)]
--    [data]
bookKeeping :: SequenceInfo -> CryptoMessage -> CryptoData
bookKeeping (SequenceInfo seqno ack) m = CryptoData
    { bufferStart = ack   :: Word32
    , bufferEnd   = seqno :: Word32
    , bufferData  = m
    }

-- | Classify an inbound packet as lossy or lossless based on its id byte.
checkLossless :: CryptoData -> PacketInboundEvent CryptoMessage
checkLossless cd@CryptoData{ bufferStart = ack
                           , bufferEnd   = no
                           , bufferData  = x   } = tag no x' ack
 where
    x' = decodeRawCryptoMsg cd
    tag = case someLossyness (msgID x') of Lossy -> PacketReceivedLossy
                                           _     -> PacketReceived