{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE ScopedTypeVariables #-}

module Network.QUIC.Handshake where

import qualified Control.Exception as E
import Data.List (intersect)
import qualified Network.TLS as TLS
import Network.TLS.QUIC

import Network.QUIC.Config
import Network.QUIC.Connection
import Network.QUIC.Connector
import Network.QUIC.Crypto
import Network.QUIC.Imports
import Network.QUIC.Info
import Network.QUIC.Logger
import Network.QUIC.Parameters
import Network.QUIC.Qlog
import Network.QUIC.Recovery
import Network.QUIC.TLS
import Network.QUIC.Types

----------------------------------------------------------------

newtype HndState = HndState
    { hsRecvCnt :: Int -- number of 'recv' calls since last 'send'
    }

newHndStateRef :: IO (IORef HndState)
newHndStateRef = newIORef HndState{hsRecvCnt = 0}

sendCompleted :: IORef HndState -> IO ()
sendCompleted hsr = atomicModifyIORef'' hsr $ \hs -> hs{hsRecvCnt = 0}

recvCompleted :: IORef HndState -> IO Int
recvCompleted hsr = atomicModifyIORef' hsr $ \hs ->
    let cnt = hsRecvCnt hs in (hs{hsRecvCnt = cnt + 1}, cnt)

rxLevelChanged :: IORef HndState -> IO ()
rxLevelChanged = sendCompleted

----------------------------------------------------------------

sendCryptoData :: Connection -> Output -> IO ()
sendCryptoData = putOutput

recvCryptoData :: Connection -> IO Crypto
recvCryptoData = takeCrypto

recvTLS
    :: Connection
    -> IORef HndState
    -> CryptLevel
    -> IO (Either TLS.TLSError ByteString)
recvTLS conn hsr level =
    case level of
        CryptInitial -> go InitialLevel
        CryptMainSecret -> failure "QUIC does not receive data < TLS 1.3"
        CryptEarlySecret -> failure "QUIC does not send early data with TLS library"
        CryptHandshakeSecret -> go HandshakeLevel
        CryptApplicationSecret -> go RTT1Level
  where
    failure = return . Left . internalError

    go expected = do
        InpHandshake actual bs <- recvCryptoData conn
        if bs == ""
            then
                return $ Left TLS.Error_EOF
            else
                if actual /= expected
                    then
                        failure $
                            "encryption level mismatch: expected "
                                ++ show expected
                                ++ " but got "
                                ++ show actual
                    else do
                        when (isClient conn) $ do
                            n <- recvCompleted hsr
                            -- Sending ACKs for three times rule
                            when ((n `mod` 3) == 1) $
                                sendCryptoData conn $
                                    OutControl HandshakeLevel []
                        return $ Right bs

sendTLS :: Connection -> IORef HndState -> [(CryptLevel, ByteString)] -> IO ()
sendTLS conn hsr x = do
    mapM convertLevel x >>= sendCryptoData conn . OutHandshake
    sendCompleted hsr
  where
    convertLevel (CryptInitial, bs) = return (InitialLevel, bs)
    convertLevel (CryptMainSecret, _) = errorTLS "QUIC does not send data < TLS 1.3"
    convertLevel (CryptEarlySecret, _) = errorTLS "QUIC does not receive early data with TLS library"
    convertLevel (CryptHandshakeSecret, bs) = return (HandshakeLevel, bs)
    convertLevel (CryptApplicationSecret, bs) = return (RTT1Level, bs)

internalError :: String -> TLS.TLSError
internalError msg = TLS.Error_Protocol msg TLS.InternalError

----------------------------------------------------------------

handshakeClient :: ClientConfig -> Connection -> AuthCIDs -> IO (IO ())
handshakeClient conf conn myAuthCIDs = do
    qlogParamsSet conn (ccParameters conf, "local") -- fixme
    handshakeClient' conf conn myAuthCIDs <$> getVersion conn <*> newHndStateRef

handshakeClient'
    :: ClientConfig -> Connection -> AuthCIDs -> Version -> IORef HndState -> IO ()
handshakeClient' conf conn myAuthCIDs ver hsr = handshaker
  where
    handshaker =
        clientHandshaker qc conf ver myAuthCIDs setter use0RTT
            `E.catch` sendCCTLSError conn
    qc =
        QUICCallbacks
            { quicSend = sendTLS conn hsr
            , quicRecv = recvTLS conn hsr
            , quicInstallKeys = installKeysClient
            , quicNotifyExtensions = setPeerParams conn
            , quicDone = done
            }
    setter = setResumptionSession conn
    installKeysClient _ctx (InstallEarlyKeys Nothing) = return ()
    installKeysClient _ctx (InstallEarlyKeys (Just (EarlySecretInfo cphr cts))) = do
        setCipher conn RTT0Level cphr
        initializeCoder conn RTT0Level (cts, ServerTrafficSecret "")
        setConnection0RTTReady conn
    installKeysClient _ctx (InstallHandshakeKeys (HandshakeSecretInfo cphr tss)) = do
        setCipher conn HandshakeLevel cphr
        setCipher conn RTT1Level cphr
        initializeCoder conn HandshakeLevel tss
        setEncryptionLevel conn HandshakeLevel
        rxLevelChanged hsr
    installKeysClient ctx (InstallApplicationKeys appSecInf@(ApplicationSecretInfo tss)) = do
        storeNegotiated conn ctx appSecInf
        initializeCoder1RTT conn tss
        setEncryptionLevel conn RTT1Level
        rxLevelChanged hsr
        setConnection1RTTReady conn
        cidInfo <- getNewMyCID conn
        putOutput conn $ OutHandshake [] -- for h3spec testing
        sendFrames conn RTT1Level [NewConnectionID cidInfo 0]
    done _ctx = do
        -- Validating Chosen Version
        mPeerVerInfo <- versionInformation <$> getPeerParameters conn
        case mPeerVerInfo of
            Nothing -> return ()
            Just peerVerInfo -> do
                hdrVer <- getVersion conn
                when (hdrVer /= chosenVersion peerVerInfo) sendCCVNError
        info <- getConnectionInfo conn
        connDebugLog conn $ bhow info
    use0RTT = ccUse0RTT conf

----------------------------------------------------------------

handshakeServer
    :: ServerConfig -> Connection -> AuthCIDs -> StatelessResetToken -> IO (IO ())
handshakeServer conf conn myAuthCIDs srt =
    handshakeServer' conf conn
        <$> getVersion conn
        <*> newHndStateRef
        <*> newIORef params
  where
    params =
        (setCIDsToParameters myAuthCIDs $ scParameters conf)
            { statelessResetToken = Just srt
            }

handshakeServer'
    :: ServerConfig
    -> Connection
    -> Version
    -> IORef HndState
    -> IORef Parameters
    -> IO ()
handshakeServer' conf conn ver hsRef paramRef = handshaker
  where
    handshaker =
        serverHandshaker qc conf ver getParams
            `E.catch` sendCCTLSError conn
    qc =
        QUICCallbacks
            { quicSend = sendTLS conn hsRef
            , quicRecv = recvTLS conn hsRef
            , quicInstallKeys = installKeysServer
            , quicNotifyExtensions = setPeerParams conn
            , quicDone = done
            }
    installKeysServer _ctx (InstallEarlyKeys Nothing) = return ()
    installKeysServer _ctx (InstallEarlyKeys (Just (EarlySecretInfo cphr cts))) = do
        setCipher conn RTT0Level cphr
        initializeCoder conn RTT0Level (cts, ServerTrafficSecret "")
        setConnection0RTTReady conn
    installKeysServer _ctx (InstallHandshakeKeys (HandshakeSecretInfo cphr tss)) = do
        setCipher conn HandshakeLevel cphr
        setCipher conn RTT1Level cphr
        initializeCoder conn HandshakeLevel tss
        setEncryptionLevel conn HandshakeLevel
        rxLevelChanged hsRef
    installKeysServer ctx (InstallApplicationKeys appSecInf@(ApplicationSecretInfo tss)) = do
        storeNegotiated conn ctx appSecInf
        initializeCoder1RTT conn tss
    -- will switch to RTT1Level after client Finished
    -- is received and verified
    done ctx = do
        setEncryptionLevel conn RTT1Level
        TLS.getClientCertificateChain ctx >>= setCertificateChain conn
        fire conn (Microseconds 100000) $ do
            let ldcc = connLDCC conn
            discarded0 <- getAndSetPacketNumberSpaceDiscarded ldcc RTT0Level
            unless discarded0 $ dropSecrets conn RTT0Level
            discarded1 <- getAndSetPacketNumberSpaceDiscarded ldcc HandshakeLevel
            unless discarded1 $ do
                dropSecrets conn HandshakeLevel
                onPacketNumberSpaceDiscarded (connLDCC conn) HandshakeLevel
            clearCryptoStream conn HandshakeLevel
            clearCryptoStream conn RTT1Level
        setConnection1RTTReady conn
        setConnectionEstablished conn
        getConnectionInfo conn >>= onConnectionEstablished (connHooks conn)
        --        sendFrames conn RTT1Level [HandshakeDone]
        --
        info <- getConnectionInfo conn
        connDebugLog conn $ bhow info
    getParams = do
        params <- readIORef paramRef
        verInfo <- getVersionInfo conn
        return params{versionInformation = Just verInfo}

----------------------------------------------------------------

setPeerParams :: Connection -> TLS.Context -> [ExtensionRaw] -> IO ()
setPeerParams conn _ctx peerExts = do
    tpId <- extensionIDForTtransportParameter <$> getVersion conn
    case getTP tpId peerExts of
        Nothing ->
            sendCCTLSAlert conn TLS.MissingExtension "QUIC transport parameters are mssing"
        Just (ExtensionRaw _ bs) -> setPP bs
  where
    getTP n = find (\(ExtensionRaw extid _) -> extid == n)
    setPP bs = case decodeParameters bs of
        Nothing -> sendCCParamError
        Just params -> do
            checkAuthCIDs params
            checkInvalid params
            setParams params
            qlogParamsSet conn (params, "remote")
            when (isServer conn) $
                serverVersionNegotiation $
                    versionInformation params

    checkAuthCIDs params = do
        peerAuthCIDs <- getPeerAuthCIDs conn
        ensure (initialSourceConnectionId params) $ initSrcCID peerAuthCIDs
        when (isClient conn) $ do
            ensure (originalDestinationConnectionId params) $ origDstCID peerAuthCIDs
            ensure (retrySourceConnectionId params) $ retrySrcCID peerAuthCIDs
    ensure _ Nothing = return ()
    ensure v0 v1
        | v0 == v1 = return ()
        | otherwise = sendCCParamError
    checkInvalid params = do
        when (maxUdpPayloadSize params < 1200) sendCCParamError
        when (ackDelayExponent params > 20) sendCCParamError
        when (maxAckDelay params >= 2 ^ (14 :: Int)) sendCCParamError
        when (activeConnectionIdLimit params < 2) sendCCParamError
        when (isServer conn) $ do
            when (isJust $ originalDestinationConnectionId params) sendCCParamError
            when (isJust $ preferredAddress params) sendCCParamError
            when (isJust $ retrySourceConnectionId params) sendCCParamError
            when (isJust $ statelessResetToken params) sendCCParamError
        let vi = case versionInformation params of
                Nothing -> VersionInfo Version1 [Version1]
                Just vi0 -> vi0
        when (vi == brokenVersionInfo) sendCCParamError
        when (Negotiation `elem` otherVersions vi) sendCCParamError
        -- Always False for servers
        isICVN <- getIncompatibleVN conn
        when isICVN $ do
            -- Validating Other Version fields.
            verInfo <- getVersionInfo conn
            let myVer = chosenVersion verInfo
                myVers = filter (not . isGreasingVersion) $ otherVersions verInfo
                peerVers = otherVersions vi
            case myVers `intersect` peerVers of
                ver : _ | ver == myVer -> return ()
                _ -> sendCCVNError

    setParams params = do
        setPeerParameters conn params
        mapM_ (setPeerStatelessResetToken conn) $ statelessResetToken params
        setTxMaxData conn $ initialMaxData params
        -- My 'maxIdleTimeout' is already set 'minIdleTimeout'
        -- This selects the minimum of mine and peer's.
        setMinIdleTimeout conn $ milliToMicro $ maxIdleTimeout params
        setMaxAckDaley (connLDCC conn) $ milliToMicro $ maxAckDelay params
        setTxMaxStreams conn $ initialMaxStreamsBidi params
        setTxUniMaxStreams conn $ initialMaxStreamsUni params

    serverVersionNegotiation Nothing = return ()
    serverVersionNegotiation (Just peerVerInfo) = do
        myVerInfo <- getVersionInfo conn
        let clientVer = chosenVersion myVerInfo
            myVers = filter (not . isGreasingVersion) $ otherVersions myVerInfo
            peerVers = otherVersions peerVerInfo
        -- Server's preference should be preferred.
        case myVers `intersect` peerVers of
            vers@(serverVer : _)
                | clientVer /= serverVer -> do
                    setVersionInfo conn $ VersionInfo serverVer vers
                    dcid <- getClientDstCID conn
                    initializeCoder conn InitialLevel $ initialSecrets serverVer dcid
                    qlogDebug conn $ Debug "Version changed"
            _ -> return ()

storeNegotiated :: Connection -> TLS.Context -> ApplicationSecretInfo -> IO ()
storeNegotiated conn ctx appSecInf = do
    appPro <- TLS.getNegotiatedProtocol ctx
    minfo <- TLS.contextGetInformation ctx
    let mode = fromMaybe FullHandshake (minfo >>= TLS.infoTLS13HandshakeMode)
    setNegotiated conn mode appPro appSecInf

----------------------------------------------------------------

sendCCParamError :: IO ()
sendCCParamError = E.throwIO WrongTransportParameter

sendCCVNError :: IO ()
sendCCVNError = E.throwIO WrongVersionInformation

sendCCTLSError :: Connection -> TLS.TLSException -> IO ()
sendCCTLSError conn (TLS.HandshakeFailed (TLS.Error_Misc "WrongTransportParameter")) = closeConnection conn TransportParameterError "Transport parameter error"
sendCCTLSError conn (TLS.HandshakeFailed (TLS.Error_Misc "WrongVersionInformation")) = closeConnection conn VersionNegotiationError "Version negotiation error"
sendCCTLSError conn e = closeConnection conn err msg
  where
    tlserr = getErrorCause e
    err = cryptoError $ errorToAlertDescription tlserr
    msg = shortpack $ errorToAlertMessage tlserr

sendCCTLSAlert :: Connection -> TLS.AlertDescription -> ReasonPhrase -> IO ()
sendCCTLSAlert conn a msg = closeConnection conn (cryptoError a) msg

getErrorCause :: TLS.TLSException -> TLS.TLSError
getErrorCause (TLS.Terminated _ _ e) = e
getErrorCause (TLS.HandshakeFailed e) = e
getErrorCause (TLS.PostHandshake e) = e
getErrorCause (TLS.Uncontextualized e) = e
getErrorCause e =
    let msg = "unexpected TLS exception: " ++ show e
     in TLS.Error_Protocol msg TLS.InternalError