{-# LANGUAGE BangPatterns #-}

module Main where

import Certificate
import Connection
import Control.Concurrent.Chan
import Data.Default (def)
import Data.IORef
import Data.X509
import Data.X509.Validation
import Gauge.Main
import Network.TLS
import Network.TLS.Extra.Cipher
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 SSL3 blockCipher) small "SSL3-256 bytes"
            , benchConnection (getParams TLS10 blockCipher) small "TLS10-256 bytes"
            , benchConnection (getParams TLS11 blockCipher) small "TLS11-256 bytes"
            , benchConnection (getParams TLS12 blockCipher) small "TLS12-256 bytes"
            ]
        , bgroup
            "resumption"
            [ benchResumption (getParams SSL3 blockCipher) small "SSL3-256 bytes"
            , benchResumption (getParams TLS10 blockCipher) small "TLS10-256 bytes"
            , benchResumption (getParams TLS11 blockCipher) small "TLS11-256 bytes"
            , 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_AES128GCM_SHA256
            , cipher_DHE_RSA_AES256GCM_SHA384
            , cipher_DHE_RSA_CHACHA20POLY1305_SHA256
            , cipher_DHE_RSA_AES128CCM_SHA256
            , cipher_DHE_RSA_AES128CCM8_SHA256
            , cipher_ECDHE_RSA_AES128GCM_SHA256
            , cipher_ECDHE_RSA_AES256GCM_SHA384
            , cipher_ECDHE_RSA_CHACHA20POLY1305_SHA256
            ]
        , benchCiphers
            "TLS13"
            TLS13
            large
            [ cipher_TLS13_AES128GCM_SHA256
            , cipher_TLS13_AES256GCM_SHA384
            , cipher_TLS13_CHACHA20POLY1305_SHA256
            , cipher_TLS13_AES128CCM_SHA256
            , cipher_TLS13_AES128CCM8_SHA256
            ]
        ]
  where
    small = B.replicate 256 0
    large = B.replicate 102400 0
