{-# LANGUAGE BangPatterns #-}

module Main where

import Certificate
import Control.Concurrent.Chan
import Data.Default (def)
import Data.IORef
import Data.X509
import Data.X509.Validation
import Test.Tasty.Bench
import Network.TLS
import Network.TLS.Extra.Cipher
import Session
import Run
import PubKey

import qualified Data.ByteString as B
import qualified Data.ByteString.Lazy as L

blockCipher :: Cipher
blockCipher =
    Cipher
        { cipherID = 0xff12
        , cipherName = "rsa-id-const"
        , cipherBulk =
            Bulk
                { bulkName = "id"
                , bulkKeySize = 16
                , bulkIVSize = 16
                , bulkExplicitIV = 0
                , bulkAuthTagLen = 0
                , bulkBlockSize = 16
                , bulkF = BulkBlockF $ \_ _ _ m -> (m, B.empty)
                }
        , cipherHash = MD5
        , cipherPRFHash = Nothing
        , cipherKeyExchange = CipherKeyExchange_RSA
        , cipherMinVer = Nothing
        }

getParams :: Version -> Cipher -> (ClientParams, ServerParams)
getParams connectVer cipher = (cParams, sParams)
  where
    sParams =
        def
            { serverSupported = supported
            , serverShared =
                def
                    { sharedCredentials =
                        Credentials
                            [(CertificateChain [simpleX509 $ PubKeyRSA pubKey], PrivKeyRSA privKey)]
                    }
            }
    cParams =
        (defaultParamsClient "" B.empty)
            { clientSupported = supported
            , clientShared =
                def
                    { sharedValidationCache =
                        ValidationCache
                            { cacheAdd = \_ _ _ -> return ()
                            , cacheQuery = \_ _ _ -> return ValidationCachePass
                            }
                    }
            }
    supported =
        def
            { supportedCiphers = [cipher]
            , supportedVersions = [connectVer]
            , supportedGroups = [X25519, FFDHE2048]
            }
    (pubKey, privKey) = getGlobalRSAPair

runTLSPipe
    :: (ClientParams, ServerParams)
    -> (Context -> Chan b -> IO ())
    -> (Chan a -> Context -> IO ())
    -> a
    -> IO b
runTLSPipe params tlsServer tlsClient d = do
    withDataPipe params tlsServer tlsClient $ \(writeStart, readResult) -> do
        writeStart d
        readResult

runTLSPipeSimple
    :: (ClientParams, ServerParams) -> B.ByteString -> IO B.ByteString
runTLSPipeSimple params = runTLSPipe params tlsServer tlsClient
  where
    tlsServer ctx queue = do
        handshake ctx
        d <- recvData ctx
        writeChan queue d
        bye ctx
    tlsClient queue ctx = do
        handshake ctx
        d <- readChan queue
        sendData ctx (L.fromChunks [d])
        byeBye ctx

benchConnection
    :: (ClientParams, ServerParams) -> B.ByteString -> String -> Benchmark
benchConnection params !d name = bench name . nfIO $ runTLSPipeSimple params d

benchResumption
    :: (ClientParams, ServerParams) -> B.ByteString -> String -> Benchmark
benchResumption params !d name = env initializeSession runResumption
  where
    initializeSession = do
        sessionRefs <- twoSessionRefs
        let sessionManagers = twoSessionManagers sessionRefs
            params1 = setPairParamsSessionManagers sessionManagers params
        _ <- runTLSPipeSimple params1 d

        Just sessionParams <- readClientSessionRef sessionRefs
        let params2 = setPairParamsSessionResuming sessionParams params1
        newIORef params2

    runResumption paramsRef = bench name . nfIO $ do
        params2 <- readIORef paramsRef
        runTLSPipeSimple params2 d

benchResumption13
    :: (ClientParams, ServerParams) -> B.ByteString -> String -> Benchmark
benchResumption13 params !d name = env initializeSession runResumption
  where
    initializeSession = do
        sessionRefs <- twoSessionRefs
        let sessionManagers = twoSessionManagers sessionRefs
            params1 = setPairParamsSessionManagers sessionManagers params
        _ <- runTLSPipeSimple params1 d
        newIORef (params1, sessionRefs)

    -- with TLS13 the sessionId is constantly changing so we must update
    -- our parameters at each iteration unfortunately
    runResumption paramsRef = bench name . nfIO $ do
        (params1, sessionRefs) <- readIORef paramsRef
        Just sessionParams <- readClientSessionRef sessionRefs
        let params2 = setPairParamsSessionResuming sessionParams params1
        runTLSPipeSimple params2 d

benchCiphers :: String -> Version -> B.ByteString -> [Cipher] -> Benchmark
benchCiphers name connectVer d = bgroup name . map doBench
  where
    doBench cipher =
        benchResumption13 (getParams connectVer cipher) d (cipherName cipher)

main :: IO ()
main =
    defaultMain
        [ bgroup
            "connection"
            -- not sure the number actually make sense for anything. improve ..
            [ benchConnection (getParams TLS12 blockCipher) small "TLS12-256 bytes"
            ]
        , bgroup
            "resumption"
            [ benchResumption (getParams TLS12 blockCipher) small "TLS12-256 bytes"
            ]
        , -- Here we try to measure TLS12 and TLS13 performance with AEAD ciphers.
          -- Resumption and a larger message can be a demonstration of the symmetric
          -- crypto but for TLS13 this does not work so well because of dhe_psk.
          benchCiphers
            "TLS12"
            TLS12
            large
            [ cipher_DHE_RSA_WITH_AES_128_GCM_SHA256
            , cipher_DHE_RSA_WITH_AES_256_GCM_SHA384
            , cipher_DHE_RSA_WITH_CHACHA20_POLY1305_SHA256
            , cipher_ECDHE_RSA_WITH_AES_128_GCM_SHA256
            , cipher_ECDHE_RSA_WITH_AES_256_GCM_SHA384
            , cipher_ECDHE_RSA_WITH_CHACHA20_POLY1305_SHA256
            ]
        , benchCiphers
            "TLS13"
            TLS13
            large
            [ cipher13_AES_128_GCM_SHA256
            , cipher13_AES_256_GCM_SHA384
            , cipher13_CHACHA20_POLY1305_SHA256
            , cipher13_AES_128_CCM_SHA256
            , cipher13_AES_128_CCM_8_SHA256
            ]
        ]
  where
    small = B.replicate 256 0
    large = B.replicate 102400 0