State.hs

{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE OverloadedStrings #-}

module Network.TLS.Handshake.State (
    HandshakeState (..),
    TransHashState (..),
    HandshakeMode13 (..),
    RTT0Status (..),
    CertReqCBdata,
    HandshakeM,
    newEmptyHandshake,
    runHandshake,

    -- * key accessors
    setPublicKey,
    setPublicPrivateKeys,
    getLocalPublicPrivateKeys,
    getRemotePublicKey,
    setServerDHParams,
    getServerDHParams,
    setServerECDHParams,
    getServerECDHParams,
    setDHPrivate,
    getDHPrivate,
    setGroupPrivate,
    getGroupPrivate,

    -- * cert accessors
    setClientRandom,
    getClientRandom,
    setOuterClientRandom,
    getOuterClientRandom,
    setClientHello,
    getClientHello,
    setECHEE,
    getECHEE,
    setECHAccepted,
    getECHAccepted,
    setClientCertSent,
    getClientCertSent,
    setCertReqSent,
    getCertReqSent,
    setClientCertChain,
    getClientCertChain,
    setCertReqToken,
    getCertReqToken,
    setCertReqCBdata,
    getCertReqCBdata,
    setCertReqSigAlgsCert,
    getCertReqSigAlgsCert,

    -- * digest accessors
    addHandshakeMessage,
    getHandshakeMessages,

    -- * main secret
    setMainSecret,
    setMainSecretFromPre,

    -- * misc accessor
    getPendingCipher,
    setExtendedMainSecret,
    getExtendedMainSecret,
    setSupportedGroup,
    getSupportedGroup,
    setTLS13HandshakeMode,
    getTLS13HandshakeMode,
    setTLS13RTT0Status,
    getTLS13RTT0Status,
    setTLS13EarlySecret,
    getTLS13EarlySecret,
    setTLS13ResumptionSecret,
    getTLS13ResumptionSecret,
    setTLS13CertComp,
    getTLS13CertComp,
    setCCS13Sent,
    getCCS13Sent,
    setCCS13Recv,
    getCCS13Recv,
) where

import Control.Monad.State.Strict
import Data.ByteArray (ByteArrayAccess)
import Data.X509 (CertificateChain)

import Network.TLS.Cipher
import Network.TLS.Compression
import Network.TLS.Crypto
import Network.TLS.Imports
import Network.TLS.Packet
import Network.TLS.Record.State
import Network.TLS.Struct
import Network.TLS.Types
import Network.TLS.Util

data HandshakeKeyState = HandshakeKeyState
    { hksRemotePublicKey :: Maybe PubKey
    , hksLocalPublicPrivateKeys :: Maybe (PubKey, PrivKey)
    }
    deriving (Show)

data TransHashState
    = -- | Initial state
      TransHashState0
    | -- | A raw CH is stored since hash algo is not chosen yet.
      TransHashState1 [ByteString]
    | -- | Hashed
      TransHashState2 HashCtx

{- FOURMOLU_DISABLE -}
instance Show TransHashState where
    show  TransHashState0       = "State0 "
    show (TransHashState1 _)    = "State1 CH"
    show (TransHashState2 hctx) = showBytesHex $ hashFinal hctx
{- FOURMOLU_ENABLE -}

data HandshakeState = HandshakeState
    { hstClientVersion :: Version
    , hstClientRandom :: ClientRandom
    -- ^ For ECH, inner client random.
    , hstServerRandom :: Maybe ServerRandom
    , hstMainSecret :: Maybe ByteString
    , hstKeyState :: HandshakeKeyState
    , hstServerDHParams :: Maybe ServerDHParams
    , hstDHPrivate :: Maybe DHPrivate
    , hstServerECDHParams :: Maybe ServerECDHParams
    , hstGroupPrivate :: Maybe GroupPrivate
    , hstTransHashState :: TransHashState
    , hstTransHashStateI :: TransHashState -- Inner CH for client ECH
    , hstHandshakeMessages :: [ByteString]
    -- ^ To create certificate verify for TLS 1.2.
    --   This should be removed when TLS 1.2 is dropped.
    , hstCertReqToken :: Maybe ByteString
    -- ^ Set to Just-value when a TLS13 certificate request is received
    , hstCertReqCBdata :: Maybe CertReqCBdata
    -- ^ Set to Just-value when a certificate request is received
    , hstCertReqSigAlgsCert :: Maybe [HashAndSignatureAlgorithm]
    -- ^ In TLS 1.3, these are separate from the certificate
    -- issuer signature algorithm hints in the callback data.
    -- In TLS 1.2 the same list is overloaded for both purposes.
    -- Not present in TLS 1.1 and earlier
    , hstClientCertSent :: Bool
    -- ^ Set to true when a client certificate chain was sent
    , hstCertReqSent :: Bool
    -- ^ Set to true when a certificate request was sent.  This applies
    -- only to requests sent during handshake (not post-handshake).
    , hstClientCertChain :: Maybe CertificateChain
    , hstPendingTxState :: Maybe RecordState
    , hstPendingRxState :: Maybe RecordState
    , hstPendingCipher :: Maybe Cipher
    , hstPendingCompression :: Compression
    , hstExtendedMainSecret :: Bool
    , hstSupportedGroup :: Maybe Group
    , hstTLS13HandshakeMode :: HandshakeMode13
    , hstTLS13RTT0Status :: RTT0Status
    , hstTLS13EarlySecret :: Maybe (BaseSecret EarlySecret) -- xxx
    , hstTLS13ResumptionSecret :: Maybe (BaseSecret ResumptionSecret)
    , hstTLS13CertComp :: Bool
    , hstCCS13Sent :: Bool
    , hstCCS13Recv :: Bool
    , hstTLS13OuterClientRandom :: Maybe ClientRandom
    -- ^ Used for key logging in the case of ECH.
    , hstTLS13ClientHello :: Maybe (ClientHello, [ByteString])
    -- ^ Inner client hello in the case of ECH.
    , hstTLS13ECHAccepted :: Bool
    , hstTLS13ECHEE :: Bool
    }
    deriving (Show)

-- | When we receive a CertificateRequest from a server, a just-in-time
--    callback is issued to the application to obtain a suitable certificate.
--    Somewhat unfortunately, the callback parameters don't abstract away the
--    details of the TLS 1.2 Certificate Request message, which combines the
--    legacy @certificate_types@ and new @supported_signature_algorithms@
--    parameters is a rather subtle way.
--
--    TLS 1.2 also (again unfortunately, in the opinion of the author of this
--    comment) overloads the signature algorithms parameter to constrain not only
--    the algorithms used in TLS, but also the algorithms used by issuing CAs in
--    the X.509 chain.  Best practice is to NOT treat such that restriction as a
--    MUST, but rather take it as merely a preference, when a choice exists.  If
--    the best chain available does not match the provided signature algorithm
--    list, go ahead and use it anyway, it will probably work, and the server may
--    not even care about the issuer CAs at all, it may be doing DANE or have
--    explicit mappings for the client's public key, ...
--
--    The TLS 1.3 @CertificateRequest@ message, drops @certificate_types@ and no
--    longer overloads @supported_signature_algorithms@ to cover X.509.  It also
--    includes a new opaque context token that the client must echo back, which
--    makes certain client authentication replay attacks more difficult.  We will
--    store that context separately, it does not need to be presented in the user
--    callback.  The certificate signature algorithms preferred by the peer are
--    now in the separate @signature_algorithms_cert@ extension, but we cannot
--    report these to the application callback without an API change.  The good
--    news is that filtering the X.509 signature types is generally unnecessary,
--    unwise and difficult.  So we just ignore this extension.
--
--    As a result, the information we provide to the callback is no longer a
--    verbatim copy of the certificate request payload.  In the case of TLS 1.3
--    The 'CertificateType' list is synthetically generated from the server's
--    @signature_algorithms@ extension, and the @signature_algorithms_certs@
--    extension is ignored.
--
--    Since the original TLS 1.2 'CertificateType' has no provision for the newer
--    certificate types that have appeared in TLS 1.3 we're adding some synthetic
--    values that have no equivalent values in the TLS 1.2 'CertificateType' as
--    defined in the IANA
--    <https://www.iana.org/assignments/tls-parameters/tls-parameters.xhtml#tls-parameters-2
--    TLS ClientCertificateType Identifiers> registry.  These values are inferred
--    from the TLS 1.3 @signature_algorithms@ extension, and will allow clients to
--    present Ed25519 and Ed448 certificates when these become supported.
type CertReqCBdata =
    ( [CertificateType]
    , Maybe [HashAndSignatureAlgorithm]
    , [DistinguishedName]
    )

newtype HandshakeM a = HandshakeM {runHandshakeM :: State HandshakeState a}
    deriving (Functor, Applicative, Monad)

instance MonadState HandshakeState HandshakeM where
    put x = HandshakeM (put x)
    get = HandshakeM get
    state f = HandshakeM (state f)

-- create a new empty handshake state
newEmptyHandshake :: Version -> ClientRandom -> HandshakeState
newEmptyHandshake ver crand =
    HandshakeState
        { hstClientVersion = ver
        , hstClientRandom = crand
        , hstServerRandom = Nothing
        , hstMainSecret = Nothing
        , hstKeyState = HandshakeKeyState Nothing Nothing
        , hstServerDHParams = Nothing
        , hstDHPrivate = Nothing
        , hstServerECDHParams = Nothing
        , hstGroupPrivate = Nothing
        , hstTransHashState = TransHashState0
        , hstTransHashStateI = TransHashState0
        , hstHandshakeMessages = []
        , hstCertReqToken = Nothing
        , hstCertReqCBdata = Nothing
        , hstCertReqSigAlgsCert = Nothing
        , hstClientCertSent = False
        , hstCertReqSent = False
        , hstClientCertChain = Nothing
        , hstPendingTxState = Nothing
        , hstPendingRxState = Nothing
        , hstPendingCipher = Nothing
        , hstPendingCompression = nullCompression
        , hstExtendedMainSecret = False
        , hstSupportedGroup = Nothing
        , hstTLS13HandshakeMode = FullHandshake
        , hstTLS13RTT0Status = RTT0None
        , hstTLS13EarlySecret = Nothing
        , hstTLS13ResumptionSecret = Nothing
        , hstTLS13CertComp = False
        , hstCCS13Sent = False
        , hstCCS13Recv = False
        , hstTLS13OuterClientRandom = Nothing
        , hstTLS13ClientHello = Nothing
        , hstTLS13ECHAccepted = False
        , hstTLS13ECHEE = False
        }

runHandshake :: HandshakeState -> HandshakeM a -> (a, HandshakeState)
runHandshake hst f = runState (runHandshakeM f) hst

setPublicKey :: PubKey -> HandshakeM ()
setPublicKey pk = modify' (\hst -> hst{hstKeyState = setPK (hstKeyState hst)})
  where
    setPK hks = hks{hksRemotePublicKey = Just pk}

setPublicPrivateKeys :: (PubKey, PrivKey) -> HandshakeM ()
setPublicPrivateKeys keys = modify' (\hst -> hst{hstKeyState = setKeys (hstKeyState hst)})
  where
    setKeys hks = hks{hksLocalPublicPrivateKeys = Just keys}

getRemotePublicKey :: HandshakeM PubKey
getRemotePublicKey = fromJust <$> gets (hksRemotePublicKey . hstKeyState)

getLocalPublicPrivateKeys :: HandshakeM (PubKey, PrivKey)
getLocalPublicPrivateKeys =
    fromJust <$> gets (hksLocalPublicPrivateKeys . hstKeyState)

setServerDHParams :: ServerDHParams -> HandshakeM ()
setServerDHParams shp = modify' (\hst -> hst{hstServerDHParams = Just shp})

getServerDHParams :: HandshakeM ServerDHParams
getServerDHParams = fromJust <$> gets hstServerDHParams

setServerECDHParams :: ServerECDHParams -> HandshakeM ()
setServerECDHParams shp = modify' (\hst -> hst{hstServerECDHParams = Just shp})

getServerECDHParams :: HandshakeM ServerECDHParams
getServerECDHParams = fromJust <$> gets hstServerECDHParams

setDHPrivate :: DHPrivate -> HandshakeM ()
setDHPrivate shp = modify' (\hst -> hst{hstDHPrivate = Just shp})

getDHPrivate :: HandshakeM DHPrivate
getDHPrivate = fromJust <$> gets hstDHPrivate

getGroupPrivate :: HandshakeM GroupPrivate
getGroupPrivate = fromJust <$> gets hstGroupPrivate

setGroupPrivate :: GroupPrivate -> HandshakeM ()
setGroupPrivate shp = modify' (\hst -> hst{hstGroupPrivate = Just shp})

setExtendedMainSecret :: Bool -> HandshakeM ()
setExtendedMainSecret b = modify' (\hst -> hst{hstExtendedMainSecret = b})

getExtendedMainSecret :: HandshakeM Bool
getExtendedMainSecret = gets hstExtendedMainSecret

setSupportedGroup :: Group -> HandshakeM ()
setSupportedGroup g = modify' (\hst -> hst{hstSupportedGroup = Just g})

getSupportedGroup :: HandshakeM (Maybe Group)
getSupportedGroup = gets hstSupportedGroup

-- | Type to show which handshake mode is used in TLS 1.3.
data HandshakeMode13
    = -- | Full handshake is used.
      FullHandshake
    | -- | Full handshake is used with hello retry request.
      HelloRetryRequest
    | -- | Server authentication is skipped.
      PreSharedKey
    | -- | Server authentication is skipped and early data is sent.
      RTT0
    deriving (Show, Eq)

setTLS13HandshakeMode :: HandshakeMode13 -> HandshakeM ()
setTLS13HandshakeMode s = modify' (\hst -> hst{hstTLS13HandshakeMode = s})

getTLS13HandshakeMode :: HandshakeM HandshakeMode13
getTLS13HandshakeMode = gets hstTLS13HandshakeMode

data RTT0Status
    = RTT0None
    | RTT0Sent
    | RTT0Accepted
    | RTT0Rejected
    deriving (Show, Eq)

setTLS13RTT0Status :: RTT0Status -> HandshakeM ()
setTLS13RTT0Status s = modify' (\hst -> hst{hstTLS13RTT0Status = s})

getTLS13RTT0Status :: HandshakeM RTT0Status
getTLS13RTT0Status = gets hstTLS13RTT0Status

setTLS13EarlySecret :: BaseSecret EarlySecret -> HandshakeM ()
setTLS13EarlySecret secret = modify' (\hst -> hst{hstTLS13EarlySecret = Just secret})

getTLS13EarlySecret :: HandshakeM (Maybe (BaseSecret EarlySecret))
getTLS13EarlySecret = gets hstTLS13EarlySecret

setTLS13ResumptionSecret :: BaseSecret ResumptionSecret -> HandshakeM ()
setTLS13ResumptionSecret secret = modify' (\hst -> hst{hstTLS13ResumptionSecret = Just secret})

getTLS13ResumptionSecret :: HandshakeM (Maybe (BaseSecret ResumptionSecret))
getTLS13ResumptionSecret = gets hstTLS13ResumptionSecret

setTLS13CertComp :: Bool -> HandshakeM ()
setTLS13CertComp comp = modify' (\hst -> hst{hstTLS13CertComp = comp})

getTLS13CertComp :: HandshakeM Bool
getTLS13CertComp = gets hstTLS13CertComp

setCCS13Sent :: Bool -> HandshakeM ()
setCCS13Sent sent = modify' (\hst -> hst{hstCCS13Sent = sent})

getCCS13Sent :: HandshakeM Bool
getCCS13Sent = gets hstCCS13Sent

setCCS13Recv :: Bool -> HandshakeM ()
setCCS13Recv sent = modify' (\hst -> hst{hstCCS13Recv = sent})

getCCS13Recv :: HandshakeM Bool
getCCS13Recv = gets hstCCS13Recv

setCertReqSent :: Bool -> HandshakeM ()
setCertReqSent b = modify' (\hst -> hst{hstCertReqSent = b})

getCertReqSent :: HandshakeM Bool
getCertReqSent = gets hstCertReqSent

setClientCertSent :: Bool -> HandshakeM ()
setClientCertSent b = modify' (\hst -> hst{hstClientCertSent = b})

getClientRandom :: HandshakeM ClientRandom
getClientRandom = gets hstClientRandom

setClientRandom :: ClientRandom -> HandshakeM ()
setClientRandom cr = modify' $ \hst -> hst{hstClientRandom = cr}

getOuterClientRandom :: HandshakeM (Maybe ClientRandom)
getOuterClientRandom = gets hstTLS13OuterClientRandom

setOuterClientRandom :: Maybe ClientRandom -> HandshakeM ()
setOuterClientRandom mcr = modify' (\hst -> hst{hstTLS13OuterClientRandom = mcr})

getClientHello :: HandshakeM (Maybe (ClientHello, [ByteString]))
getClientHello = gets hstTLS13ClientHello

setClientHello :: ClientHello -> [ByteString] -> HandshakeM ()
setClientHello ch b = modify' $ \hst -> hst{hstTLS13ClientHello = Just (ch, b)}

getECHAccepted :: HandshakeM Bool
getECHAccepted = gets hstTLS13ECHAccepted

setECHAccepted :: Bool -> HandshakeM ()
setECHAccepted b = modify' $ \hst -> hst{hstTLS13ECHAccepted = b}

getECHEE :: HandshakeM Bool
getECHEE = gets hstTLS13ECHEE

setECHEE :: Bool -> HandshakeM ()
setECHEE b = modify' $ \hst -> hst{hstTLS13ECHEE = b}

getClientCertSent :: HandshakeM Bool
getClientCertSent = gets hstClientCertSent

setClientCertChain :: CertificateChain -> HandshakeM ()
setClientCertChain b = modify' (\hst -> hst{hstClientCertChain = Just b})

getClientCertChain :: HandshakeM (Maybe CertificateChain)
getClientCertChain = gets hstClientCertChain

--
setCertReqToken :: Maybe ByteString -> HandshakeM ()
setCertReqToken token = modify' $ \hst -> hst{hstCertReqToken = token}

getCertReqToken :: HandshakeM (Maybe ByteString)
getCertReqToken = gets hstCertReqToken

--
setCertReqCBdata :: Maybe CertReqCBdata -> HandshakeM ()
setCertReqCBdata d = modify' (\hst -> hst{hstCertReqCBdata = d})

getCertReqCBdata :: HandshakeM (Maybe CertReqCBdata)
getCertReqCBdata = gets hstCertReqCBdata

-- Dead code, until we find some use for the extension
setCertReqSigAlgsCert :: Maybe [HashAndSignatureAlgorithm] -> HandshakeM ()
setCertReqSigAlgsCert as = modify' $ \hst -> hst{hstCertReqSigAlgsCert = as}

getCertReqSigAlgsCert :: HandshakeM (Maybe [HashAndSignatureAlgorithm])
getCertReqSigAlgsCert = gets hstCertReqSigAlgsCert

--
getPendingCipher :: HandshakeM Cipher
getPendingCipher = fromJust <$> gets hstPendingCipher

addHandshakeMessage :: ByteString -> HandshakeM ()
addHandshakeMessage content = modify' $ \hs -> hs{hstHandshakeMessages = content : hstHandshakeMessages hs}

getHandshakeMessages :: HandshakeM [ByteString]
getHandshakeMessages = gets (reverse . hstHandshakeMessages)

-- | Generate the main secret from the pre-main secret.
setMainSecretFromPre
    :: ByteArrayAccess preMain
    => Version
    -- ^ chosen transmission version
    -> Role
    -- ^ the role (Client or Server) of the generating side
    -> preMain
    -- ^ the pre-main secret
    -> HandshakeM ByteString
setMainSecretFromPre ver role preMainSecret = do
    ems <- getExtendedMainSecret
    secret <- if ems then get >>= genExtendedSecret else genSecret <$> get
    setMainSecret ver role secret
    return secret
  where
    genSecret hst =
        generateMainSecret
            ver
            (fromJust $ hstPendingCipher hst)
            preMainSecret
            (hstClientRandom hst)
            (fromJust $ hstServerRandom hst)
    genExtendedSecret hst =
        generateExtendedMainSecret
            ver
            (fromJust $ hstPendingCipher hst)
            preMainSecret
            <$> getSessionHash

getSessionHash :: HandshakeM ByteString
getSessionHash = gets $ \hst ->
    case hstTransHashState hst of
        TransHashState2 hashCtx -> hashFinal hashCtx
        _ -> error "un-initialized session hash"

-- | Set main secret and as a side effect generate the key block
-- with all the right parameters, and setup the pending tx/rx state.
setMainSecret :: Version -> Role -> ByteString -> HandshakeM ()
setMainSecret ver role mainSecret = modify' $ \hst ->
    let (pendingTx, pendingRx) = computeKeyBlock hst mainSecret ver role
     in hst
            { hstMainSecret = Just mainSecret
            , hstPendingTxState = Just pendingTx
            , hstPendingRxState = Just pendingRx
            }

computeKeyBlock
    :: HandshakeState -> ByteString -> Version -> Role -> (RecordState, RecordState)
computeKeyBlock hst mainSecret ver cc = (pendingTx, pendingRx)
  where
    cipher = fromJust $ hstPendingCipher hst
    keyblockSize = cipherKeyBlockSize cipher

    bulk = cipherBulk cipher
    digestSize =
        if hasMAC (bulkF bulk)
            then hashDigestSize (cipherHash cipher)
            else 0
    keySize = bulkKeySize bulk
    ivSize = bulkIVSize bulk
    kb =
        generateKeyBlock
            ver
            cipher
            (hstClientRandom hst)
            (fromJust $ hstServerRandom hst)
            mainSecret
            keyblockSize

    (cMACSecret, sMACSecret, cWriteKey, sWriteKey, cWriteIV, sWriteIV) =
        fromJust $
            partition6 kb (digestSize, digestSize, keySize, keySize, ivSize, ivSize)

    cstClient =
        CryptState
            { cstKey = bulkInit bulk (BulkEncrypt `orOnServer` BulkDecrypt) cWriteKey
            , cstIV = cWriteIV
            , cstMacSecret = cMACSecret
            }
    cstServer =
        CryptState
            { cstKey = bulkInit bulk (BulkDecrypt `orOnServer` BulkEncrypt) sWriteKey
            , cstIV = sWriteIV
            , cstMacSecret = sMACSecret
            }
    msClient = MacState{msSequence = 0}
    msServer = MacState{msSequence = 0}

    pendingTx =
        RecordState
            { stCryptState = if cc == ClientRole then cstClient else cstServer
            , stMacState = if cc == ClientRole then msClient else msServer
            , stCryptLevel = CryptMainSecret
            , stCipher = Just cipher
            , stCompression = hstPendingCompression hst
            }
    pendingRx =
        RecordState
            { stCryptState = if cc == ClientRole then cstServer else cstClient
            , stMacState = if cc == ClientRole then msServer else msClient
            , stCryptLevel = CryptMainSecret
            , stCipher = Just cipher
            , stCompression = hstPendingCompression hst
            }

    orOnServer f g = if cc == ClientRole then f else g