hs-tls/Network/TLS/State.hs

{-# LANGUAGE GeneralizedNewtypeDeriving, FlexibleContexts, MultiParamTypeClasses, ExistentialQuantification, RankNTypes, CPP #-}
-- |
-- Module      : Network.TLS.State
-- License     : BSD-style
-- Maintainer  : Vincent Hanquez <vincent@snarc.org>
-- Stability   : experimental
-- Portability : unknown
--
-- the State module contains calls related to state initialization/manipulation
-- which is use by the Receiving module and the Sending module.
--
module Network.TLS.State
        ( TLSState(..)
        , TLSSt
        , runTLSState
        , TLSHandshakeState(..)
        , TLSCryptState(..)
        , TLSMacState(..)
        , newTLSState
        , genTLSRandom
        , withTLSRNG
        , withCompression
        , assert -- FIXME move somewhere else (Internal.hs ?)
        , updateVerifiedData
        , finishHandshakeTypeMaterial
        , finishHandshakeMaterial
        , makeDigest
        , setMasterSecret
        , setMasterSecretFromPre
        , setPublicKey
        , setPrivateKey
        , setClientPublicKey
        , setClientPrivateKey
        , setClientCertSent
        , getClientCertSent
        , setClientCertChain
        , getClientCertChain
        , setClientCertRequest
        , getClientCertRequest
        , setKeyBlock
        , setVersion
        , setCipher
        , setServerRandom
        , setSecureRenegotiation
        , getSecureRenegotiation
        , setExtensionNPN
        , getExtensionNPN
        , setNegotiatedProtocol
        , getNegotiatedProtocol
        , setServerNextProtocolSuggest
        , getServerNextProtocolSuggest
        , getClientCertificateChain
        , setClientCertificateChain
        , getVerifiedData
        , setSession
        , getSession
        , getSessionData
        , isSessionResuming
        , switchTxEncryption
        , switchRxEncryption
        , getCipherKeyExchangeType
        , isClientContext
        , startHandshakeClient
        , updateHandshakeDigest
        , getHandshakeDigest
        , getCertVerifyDigest
        , endHandshake
        ) where

import Data.Word
import Data.Maybe (isNothing)
import Network.TLS.Util
import Network.TLS.Struct
import Network.TLS.Wire
import Network.TLS.Packet
import Network.TLS.Crypto
import Network.TLS.Cipher
import Network.TLS.Compression
import Network.TLS.MAC
import qualified Data.ByteString as B
import Control.Applicative ((<$>))
import Control.Monad
import Control.Monad.State
import Control.Monad.Error
import Crypto.Random
import Data.Certificate.X509

assert :: Monad m => String -> [(String,Bool)] -> m ()
assert fctname list = forM_ list $ \ (name, assumption) -> do
        when assumption $ fail (fctname ++ ": assumption about " ++ name ++ " failed")

data TLSCryptState = TLSCryptState
        { cstKey        :: !Bytes
        , cstIV         :: !Bytes
        , cstMacSecret  :: !Bytes
        } deriving (Show)

data TLSMacState = TLSMacState
        { msSequence :: Word64
        } deriving (Show)

type ClientCertRequestData = ([CertificateType],
                              Maybe [(HashAlgorithm, SignatureAlgorithm)],
                              [DistinguishedName])
  
data TLSHandshakeState = TLSHandshakeState
        { hstClientVersion   :: !(Version)
        , hstClientRandom    :: !ClientRandom
        , hstServerRandom    :: !(Maybe ServerRandom)
        , hstMasterSecret    :: !(Maybe Bytes)
        , hstRSAPublicKey    :: !(Maybe PublicKey)
        , hstRSAPrivateKey   :: !(Maybe PrivateKey)
        , hstRSAClientPublicKey    :: !(Maybe PublicKey)
        , hstRSAClientPrivateKey   :: !(Maybe PrivateKey)
        , hstHandshakeDigest :: !HashCtx
        , hstClientCertRequest   :: !(Maybe ClientCertRequestData) -- ^ Set to Just-value when certificate request was received
        , hstClientCertSent      :: !Bool -- ^ Set to true when a client certificate chain was sent
        , hstClientCertChain :: !(Maybe [X509])
        } deriving (Show)

data StateRNG = forall g . CryptoRandomGen g => StateRNG g

instance Show StateRNG where
        show _ = "rng[..]"

data TLSState = TLSState
        { stClientContext       :: Bool
        , stVersion             :: !Version
        , stHandshake           :: !(Maybe TLSHandshakeState)
        , stSession             :: Session
        , stSessionResuming     :: Bool
        , stTxEncrypted         :: Bool
        , stRxEncrypted         :: Bool
        , stTxCryptState        :: !(Maybe TLSCryptState)
        , stRxCryptState        :: !(Maybe TLSCryptState)
        , stTxMacState          :: !(Maybe TLSMacState)
        , stRxMacState          :: !(Maybe TLSMacState)
        , stCipher              :: Maybe Cipher
        , stCompression         :: Compression
        , stRandomGen           :: StateRNG
        , stSecureRenegotiation :: Bool  -- RFC 5746
        , stClientVerifiedData  :: Bytes -- RFC 5746
        , stServerVerifiedData  :: Bytes -- RFC 5746
        , stExtensionNPN        :: Bool  -- NPN draft extension
        , stNegotiatedProtocol  :: Maybe B.ByteString -- NPN protocol
        , stServerNextProtocolSuggest :: Maybe [B.ByteString]
        , stClientCertificateChain :: Maybe [X509]
        } deriving (Show)

newtype TLSSt a = TLSSt { runTLSSt :: ErrorT TLSError (State TLSState) a }
        deriving (Monad, MonadError TLSError)

instance Functor TLSSt where
        fmap f = TLSSt . fmap f . runTLSSt

instance MonadState TLSState TLSSt where
        put x = TLSSt (lift $ put x)
        get   = TLSSt (lift get)
#if MIN_VERSION_mtl(2,1,0)
        state f = TLSSt (lift $ state f)
#endif

runTLSState :: TLSSt a -> TLSState -> (Either TLSError a, TLSState)
runTLSState f st = runState (runErrorT (runTLSSt f)) st

newTLSState :: CryptoRandomGen g => g -> TLSState
newTLSState rng = TLSState
        { stClientContext       = False
        , stVersion             = TLS10
        , stHandshake           = Nothing
        , stSession             = Session Nothing
        , stSessionResuming     = False
        , stTxEncrypted         = False
        , stRxEncrypted         = False
        , stTxCryptState        = Nothing
        , stRxCryptState        = Nothing
        , stTxMacState          = Nothing
        , stRxMacState          = Nothing
        , stCipher              = Nothing
        , stCompression         = nullCompression
        , stRandomGen           = StateRNG rng
        , stSecureRenegotiation = False
        , stClientVerifiedData  = B.empty
        , stServerVerifiedData  = B.empty
        , stExtensionNPN        = False
        , stNegotiatedProtocol  = Nothing
        , stServerNextProtocolSuggest = Nothing
        , stClientCertificateChain = Nothing
        }

withTLSRNG :: StateRNG -> (forall g . CryptoRandomGen g => g -> Either e (a,g)) -> Either e (a, StateRNG)
withTLSRNG (StateRNG rng) f = case f rng of
        Left err        -> Left err
        Right (a, rng') -> Right (a, StateRNG rng')

withCompression :: (Compression -> (Compression, a)) -> TLSSt a
withCompression f = do
        compression <- stCompression <$> get
        let (nc, a) = f compression
        modify (\st -> st { stCompression = nc })
        return a

genTLSRandom :: (MonadState TLSState m, MonadError TLSError m) => Int -> m Bytes
genTLSRandom n = do
        st <- get
        case withTLSRNG (stRandomGen st) (genBytes n) of
                Left err            -> throwError $ Error_Random $ show err
                Right (bytes, rng') -> put (st { stRandomGen = rng' }) >> return bytes

makeDigest :: MonadState TLSState m => Bool -> Header -> Bytes -> m Bytes
makeDigest w hdr content = do
        st <- get
        let ver = stVersion st
        let cst = fromJust "crypt state" $ if w then stTxCryptState st else stRxCryptState st
        let ms = fromJust "mac state" $ if w then stTxMacState st else stRxMacState st
        let cipher = fromJust "cipher" $ stCipher st
        let hashf = hashF $ cipherHash cipher

        let (macF, msg) =
                if ver < TLS10
                        then (macSSL hashf, B.concat [ encodeWord64 $ msSequence ms, encodeHeaderNoVer hdr, content ])
                        else (hmac hashf 64, B.concat [ encodeWord64 $ msSequence ms, encodeHeader hdr, content ])
        let digest = macF (cstMacSecret cst) msg

        let newms = ms { msSequence = (msSequence ms) + 1 }

        modify (\_ -> if w then st { stTxMacState = Just newms } else st { stRxMacState = Just newms })
        return digest

updateVerifiedData :: MonadState TLSState m => Bool -> Bytes -> m ()
updateVerifiedData sending bs = do
        cc <- isClientContext
        if cc /= sending
                then modify (\st -> st { stServerVerifiedData = bs })
                else modify (\st -> st { stClientVerifiedData = bs })

finishHandshakeTypeMaterial :: HandshakeType -> Bool
finishHandshakeTypeMaterial HandshakeType_ClientHello     = True
finishHandshakeTypeMaterial HandshakeType_ServerHello     = True
finishHandshakeTypeMaterial HandshakeType_Certificate     = True
finishHandshakeTypeMaterial HandshakeType_HelloRequest    = False
finishHandshakeTypeMaterial HandshakeType_ServerHelloDone = True
finishHandshakeTypeMaterial HandshakeType_ClientKeyXchg   = True
finishHandshakeTypeMaterial HandshakeType_ServerKeyXchg   = True
finishHandshakeTypeMaterial HandshakeType_CertRequest     = True
finishHandshakeTypeMaterial HandshakeType_CertVerify      = False
finishHandshakeTypeMaterial HandshakeType_Finished        = True
finishHandshakeTypeMaterial HandshakeType_NPN             = True

finishHandshakeMaterial :: Handshake -> Bool
finishHandshakeMaterial = finishHandshakeTypeMaterial . typeOfHandshake

switchTxEncryption, switchRxEncryption :: MonadState TLSState m => m ()
switchTxEncryption = modify (\st -> st { stTxEncrypted = True })
switchRxEncryption = modify (\st -> st { stRxEncrypted = True })

setServerRandom :: MonadState TLSState m => ServerRandom -> m ()
setServerRandom ran = updateHandshake "srand" (\hst -> hst { hstServerRandom = Just ran })

setMasterSecret :: MonadState TLSState m => Bytes -> m ()
setMasterSecret masterSecret = do
        hasValidHandshake "master secret"

        updateHandshake "master secret" (\hst -> hst { hstMasterSecret = Just masterSecret } )
        setKeyBlock
        return ()

setMasterSecretFromPre :: MonadState TLSState m => Bytes -> m ()
setMasterSecretFromPre premasterSecret = do
        hasValidHandshake "generate master secret"
        st <- get
        setMasterSecret $ genSecret st
        where
                genSecret st =
                        let hst = fromJust "handshake" $ stHandshake st in
                        generateMasterSecret (stVersion st)
                                             premasterSecret
                                             (hstClientRandom hst)
                                             (fromJust "server random" $ hstServerRandom hst)

setPublicKey :: MonadState TLSState m => PublicKey -> m ()
setPublicKey pk = updateHandshake "publickey" (\hst -> hst { hstRSAPublicKey = Just pk })

setPrivateKey :: MonadState TLSState m => PrivateKey -> m ()
setPrivateKey pk = updateHandshake "privatekey" (\hst -> hst { hstRSAPrivateKey = Just pk })

setClientPublicKey :: MonadState TLSState m => PublicKey -> m ()
setClientPublicKey pk = updateHandshake "client publickey" (\hst -> hst { hstRSAClientPublicKey = Just pk })

setClientPrivateKey :: MonadState TLSState m => PrivateKey -> m ()
setClientPrivateKey pk = updateHandshake "client privatekey" (\hst -> hst { hstRSAClientPrivateKey = Just pk })

setClientCertSent :: MonadState TLSState m => Bool -> m ()
setClientCertSent b = updateHandshake "client cert sent" (\hst -> hst { hstClientCertSent = b })

getClientCertSent :: MonadState TLSState m => m (Maybe Bool)
getClientCertSent = do
        st <- get
        return (stHandshake st >>= Just . hstClientCertSent)

setClientCertChain :: MonadState TLSState m => [X509] -> m ()
setClientCertChain b = updateHandshake "client certificate chain" (\hst -> hst { hstClientCertChain = Just b })

getClientCertChain :: MonadState TLSState m => m (Maybe [X509])
getClientCertChain = do
        st <- get
        return (stHandshake st >>= hstClientCertChain)

setClientCertRequest :: MonadState TLSState m => ClientCertRequestData -> m ()
setClientCertRequest d = updateHandshake "client cert data" (\hst -> hst { hstClientCertRequest = Just d })

getClientCertRequest :: MonadState TLSState m => m (Maybe ClientCertRequestData)
getClientCertRequest = do
        st <- get
        return (stHandshake st >>= hstClientCertRequest)

getSessionData :: MonadState TLSState m => m (Maybe SessionData)
getSessionData = do
        st <- get
        return (stHandshake st >>= hstMasterSecret >>= wrapSessionData st)
        where wrapSessionData st masterSecret = do
                return $ SessionData
                        { sessionVersion = stVersion st
                        , sessionCipher  = cipherID $ fromJust "cipher" $ stCipher st
                        , sessionSecret  = masterSecret
                        }

setSession :: MonadState TLSState m => Session -> Bool -> m ()
setSession session resuming = modify (\st -> st { stSession = session, stSessionResuming = resuming })

getSession :: MonadState TLSState m => m Session
getSession = gets stSession

isSessionResuming :: MonadState TLSState m => m Bool
isSessionResuming = gets stSessionResuming

setKeyBlock :: MonadState TLSState m => m ()
setKeyBlock = do
        st <- get

        let hst = fromJust "handshake" $ stHandshake st

        let cc = stClientContext st
        let cipher = fromJust "cipher" $ stCipher st
        let keyblockSize = cipherKeyBlockSize cipher

        let bulk = cipherBulk cipher
        let digestSize   = hashSize $ cipherHash cipher
        let keySize      = bulkKeySize bulk
        let ivSize       = bulkIVSize bulk
        let kb = generateKeyBlock (stVersion st) (hstClientRandom hst)
                                  (fromJust "server random" $ hstServerRandom hst)
                                  (fromJust "master secret" $ hstMasterSecret hst) keyblockSize

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

        let cstClient = TLSCryptState
                { cstKey        = cWriteKey
                , cstIV         = cWriteIV
                , cstMacSecret  = cMACSecret }
        let cstServer = TLSCryptState
                { cstKey        = sWriteKey
                , cstIV         = sWriteIV
                , cstMacSecret  = sMACSecret }
        let msClient = TLSMacState { msSequence = 0 }
        let msServer = TLSMacState { msSequence = 0 }
        put $ st
                { stTxCryptState = Just $ if cc then cstClient else cstServer
                , stRxCryptState = Just $ if cc then cstServer else cstClient
                , stTxMacState   = Just $ if cc then msClient else msServer
                , stRxMacState   = Just $ if cc then msServer else msClient
                }

setCipher :: MonadState TLSState m => Cipher -> m ()
setCipher cipher = modify (\st -> st { stCipher = Just cipher })

setVersion :: MonadState TLSState m => Version -> m ()
setVersion ver = modify (\st -> st { stVersion = ver })

setSecureRenegotiation :: MonadState TLSState m => Bool -> m ()
setSecureRenegotiation b = modify (\st -> st { stSecureRenegotiation = b })

getSecureRenegotiation :: MonadState TLSState m => m Bool
getSecureRenegotiation = get >>= return . stSecureRenegotiation

setExtensionNPN :: MonadState TLSState m => Bool -> m ()
setExtensionNPN b = modify (\st -> st { stExtensionNPN = b })

getExtensionNPN :: MonadState TLSState m => m Bool
getExtensionNPN = get >>= return . stExtensionNPN

setNegotiatedProtocol :: MonadState TLSState m => B.ByteString -> m ()
setNegotiatedProtocol s = modify (\st -> st { stNegotiatedProtocol = Just s })

getNegotiatedProtocol :: MonadState TLSState m => m (Maybe B.ByteString)
getNegotiatedProtocol = get >>= return . stNegotiatedProtocol

setServerNextProtocolSuggest :: MonadState TLSState m => [B.ByteString] -> m ()
setServerNextProtocolSuggest ps = modify (\st -> st { stServerNextProtocolSuggest = Just ps})

getServerNextProtocolSuggest :: MonadState TLSState m => m (Maybe [B.ByteString])
getServerNextProtocolSuggest = get >>= return . stServerNextProtocolSuggest

setClientCertificateChain :: MonadState TLSState m => [X509] -> m ()
setClientCertificateChain s = modify (\st -> st { stClientCertificateChain = Just s })

getClientCertificateChain :: MonadState TLSState m => m (Maybe [X509])
getClientCertificateChain = get >>= return . stClientCertificateChain

getCipherKeyExchangeType :: MonadState TLSState m => m (Maybe CipherKeyExchangeType)
getCipherKeyExchangeType = get >>= return . (maybe Nothing (Just . cipherKeyExchange) . stCipher)

getVerifiedData :: MonadState TLSState m => Bool -> m Bytes
getVerifiedData client = get >>= return . (if client then stClientVerifiedData else stServerVerifiedData)

isClientContext :: MonadState TLSState m => m Bool
isClientContext = get >>= return . stClientContext

-- create a new empty handshake state
newEmptyHandshake :: Version -> ClientRandom -> HashCtx -> TLSHandshakeState
newEmptyHandshake ver crand digestInit = TLSHandshakeState
        { hstClientVersion   = ver
        , hstClientRandom    = crand
        , hstServerRandom    = Nothing
        , hstMasterSecret    = Nothing
        , hstRSAPublicKey    = Nothing
        , hstRSAPrivateKey   = Nothing
        , hstRSAClientPublicKey    = Nothing
        , hstRSAClientPrivateKey   = Nothing
        , hstHandshakeDigest = digestInit
        , hstClientCertRequest = Nothing
        , hstClientCertSent  = False
        , hstClientCertChain = Nothing
        }

startHandshakeClient :: MonadState TLSState m => Version -> ClientRandom -> m ()
startHandshakeClient ver crand = do
        -- FIXME check if handshake is already not null
        let initCtx = if ver < TLS12 then hashMD5SHA1 else hashSHA256
        chs <- get >>= return . stHandshake
        when (isNothing chs) $
                modify (\st -> st { stHandshake = Just $ newEmptyHandshake ver crand initCtx })

hasValidHandshake :: MonadState TLSState m => String -> m ()
hasValidHandshake name = get >>= \st -> assert name [ ("valid handshake", isNothing $ stHandshake st) ]

updateHandshake :: MonadState TLSState m => String -> (TLSHandshakeState -> TLSHandshakeState) -> m ()
updateHandshake n f = do
        hasValidHandshake n
        modify (\st -> st { stHandshake = f <$> stHandshake st })

updateHandshakeDigest :: MonadState TLSState m => Bytes -> m ()
updateHandshakeDigest content = updateHandshake "update digest" $ \hs ->
        hs { hstHandshakeDigest = hashUpdate (hstHandshakeDigest hs) content }

getHandshakeDigest :: MonadState TLSState m => Bool -> m Bytes
getHandshakeDigest client = do
        st <- get
        let hst = fromJust "handshake" $ stHandshake st
        let hashctx = hstHandshakeDigest hst
        let msecret = fromJust "master secret" $ hstMasterSecret hst
        return $ (if client then generateClientFinished else generateServerFinished) (stVersion st) msecret hashctx

getCertVerifyDigest :: MonadState TLSState m => m Bytes
getCertVerifyDigest = do
        st <- get
        let hst = fromJust "handshake" $ stHandshake st
        let hashctx = hstHandshakeDigest hst
        return (hashFinal hashctx)

endHandshake :: MonadState TLSState m => m ()
endHandshake = modify (\st -> st { stHandshake = Nothing })