simplify the handshake key state.

we don't need to differentiate client/server, as a remote key will always be public and
the private key always local.

core/Network/TLS/Handshake/Client.hs

@@ -118,7 +118,7 @@ sendClientData cparams ctx = sendCertificate >> sendClientKeyXchg >> sendCertifi
                             case certPubKey $ getCertificate c of
                                 PubKeyRSA _ -> return ()
                                 _           -> throwCore $ Error_Protocol ("no supported certificate type", True, HandshakeFailure)
-                            usingHState ctx $ setClientPrivateKey pk
+                            usingHState ctx $ setPrivateKey pk
                             usingHState ctx $ setClientCertSent True
                             sendPacket ctx $ Handshake [Certificates cc]
 

core/Network/TLS/Handshake/Key.hs

@@ -15,12 +15,9 @@ module Network.TLS.Handshake.Key
     , generateDHE
     ) where
 
-import Control.Monad.State
-
 import Data.ByteString (ByteString)
 import qualified Data.ByteString as B
 
-import Network.TLS.Util
 import Network.TLS.Handshake.State
 import Network.TLS.State (withRNG, getVersion)
 import Network.TLS.Crypto
@@ -32,39 +29,34 @@ import Network.TLS.Context
  -}
 encryptRSA :: Context -> ByteString -> IO ByteString
 encryptRSA ctx content = do
-    rsakey <- return . fromJust "rsa public key" =<< handshakeGet ctx hstRSAPublicKey
+    publicKey <- usingHState ctx getRemotePublicKey
     usingState_ ctx $ do
-        v      <- withRNG (\g -> kxEncrypt g rsakey content)
+        v      <- withRNG (\g -> kxEncrypt g publicKey content)
         case v of
             Left err       -> fail ("rsa encrypt failed: " ++ show err)
             Right econtent -> return econtent
 
 signRSA :: Context -> Role -> HashDescr -> ByteString -> IO ByteString
-signRSA ctx role hsh content = do
-    rsakey <- return . fromJust "rsa client private key" =<< access
+signRSA ctx _ hsh content = do
+    privateKey <- usingHState ctx getLocalPrivateKey
     usingState_ ctx $ do
-        r      <- withRNG (\g -> kxSign g rsakey hsh content)
+        r      <- withRNG (\g -> kxSign g privateKey hsh content)
         case r of
             Left err       -> fail ("rsa sign failed: " ++ show err)
             Right econtent -> return econtent
-  where access = handshakeGet ctx $ (if role == ClientRole then hstRSAClientPrivateKey else hstRSAPrivateKey)
 
 decryptRSA :: Context -> ByteString -> IO (Either KxError ByteString)
 decryptRSA ctx econtent = do
-    rsapriv <- return . fromJust "rsa private key" =<< handshakeGet ctx hstRSAPrivateKey
+    privateKey <- usingHState ctx getLocalPrivateKey
     usingState_ ctx $ do
         ver     <- getVersion
         let cipher = if ver < TLS10 then econtent else B.drop 2 econtent
-        withRNG (\g -> kxDecrypt g rsapriv cipher)
+        withRNG (\g -> kxDecrypt g privateKey cipher)
 
 verifyRSA :: Context -> Role -> HashDescr -> ByteString -> ByteString -> IO Bool
-verifyRSA ctx role hsh econtent sign = do
-    rsapriv <- return . fromJust "rsa client public key" =<< access
-    return $ kxVerify rsapriv hsh econtent sign
-  where access = handshakeGet ctx $ (if role == ClientRole then hstRSAPublicKey else hstRSAClientPublicKey)
-
-handshakeGet :: Context -> (HandshakeState -> a) -> IO a
-handshakeGet ctx f = usingHState ctx (gets f)
+verifyRSA ctx _ hsh econtent sign = do
+    publicKey <- usingHState ctx getRemotePublicKey
+    return $ kxVerify publicKey hsh econtent sign
 
 generateDHE :: Context -> DHParams -> IO (DHPrivate, DHPublic)
 generateDHE ctx dhp = usingState_ ctx $ withRNG $ \rng -> dhGenerateKeyPair rng dhp

core/Network/TLS/Handshake/Process.hs

@@ -61,9 +61,8 @@ processHandshake ctx hs = do
         processCertificates ServerRole (CertificateChain []) = return ()
         processCertificates ClientRole (CertificateChain []) =
             throwCore $ Error_Protocol ("server certificate missing", True, HandshakeFailure)
-        processCertificates role (CertificateChain (c:_))
-            | role == ClientRole = usingHState ctx $ setPublicKey pubkey
-            | otherwise          = usingHState ctx $ setClientPublicKey pubkey
+        processCertificates _ (CertificateChain (c:_)) =
+            usingHState ctx $ setPublicKey pubkey
           where pubkey = certPubKey $ getCertificate c
 
 -- process the client key exchange message. the protocol expects the initial

core/Network/TLS/Handshake/State.hs

@@ -17,8 +17,8 @@ module Network.TLS.Handshake.State
     -- * key accessors
     , setPublicKey
     , setPrivateKey
-    , setClientPublicKey
-    , setClientPrivateKey
+    , getLocalPrivateKey
+    , getRemotePublicKey
     , setServerDHParams
     -- * cert accessors
     , setClientCertSent
@@ -54,15 +54,17 @@ import Control.Applicative (Applicative, (<$>))
 import Control.Monad.State
 import Data.X509 (CertificateChain)
 
+data HandshakeKeyState = HandshakeKeyState
+    { hksRemotePublicKey :: !(Maybe PubKey)
+    , hksLocalPrivateKey :: !(Maybe PrivKey)
+    } deriving (Show)
+
 data HandshakeState = HandshakeState
     { hstClientVersion       :: !(Version)
     , hstClientRandom        :: !ClientRandom
     , hstServerRandom        :: !(Maybe ServerRandom)
     , hstMasterSecret        :: !(Maybe Bytes)
-    , hstRSAPublicKey        :: !(Maybe PubKey)
-    , hstRSAPrivateKey       :: !(Maybe PrivKey)
-    , hstRSAClientPublicKey  :: !(Maybe PubKey)
-    , hstRSAClientPrivateKey :: !(Maybe PrivKey)
+    , hstKeyState            :: !HandshakeKeyState
     , hstServerDHParams      :: !(Maybe ServerDHParams)
     , hstDHPrivate           :: !(Maybe DHPrivate)
     , hstHandshakeDigest     :: !(Either [Bytes] HashCtx)
@@ -98,10 +100,7 @@ newEmptyHandshake ver crand = HandshakeState
     , hstClientRandom        = crand
     , hstServerRandom        = Nothing
     , hstMasterSecret        = Nothing
-    , hstRSAPublicKey        = Nothing
-    , hstRSAPrivateKey       = Nothing
-    , hstRSAClientPublicKey  = Nothing
-    , hstRSAClientPrivateKey = Nothing
+    , hstKeyState            = HandshakeKeyState Nothing Nothing
     , hstServerDHParams      = Nothing
     , hstDHPrivate           = Nothing
     , hstHandshakeDigest     = Left []
@@ -120,20 +119,22 @@ runHandshake :: HandshakeState -> HandshakeM a -> (a, HandshakeState)
 runHandshake hst f = runState (runHandshakeM f) hst
 
 setPublicKey :: PubKey -> HandshakeM ()
-setPublicKey pk = modify (\hst -> hst { hstRSAPublicKey = Just pk })
+setPublicKey pk = modify (\hst -> hst { hstKeyState = setPK (hstKeyState hst) })
+  where setPK hks = hks { hksRemotePublicKey = Just pk }
 
 setPrivateKey :: PrivKey -> HandshakeM ()
-setPrivateKey pk = modify (\hst -> hst { hstRSAPrivateKey = Just pk })
+setPrivateKey pk = modify (\hst -> hst { hstKeyState = setPK (hstKeyState hst) })
+  where setPK hks = hks { hksLocalPrivateKey = Just pk }
+
+getRemotePublicKey :: HandshakeM PubKey
+getRemotePublicKey = fromJust "remote public key" <$> gets (hksRemotePublicKey . hstKeyState)
+
+getLocalPrivateKey :: HandshakeM PrivKey
+getLocalPrivateKey = fromJust "local private key" <$> gets (hksLocalPrivateKey . hstKeyState)
 
 setServerDHParams :: ServerDHParams -> HandshakeM ()
 setServerDHParams shp = modify (\hst -> hst { hstServerDHParams = Just shp })
 
-setClientPublicKey :: PubKey -> HandshakeM ()
-setClientPublicKey pk = modify (\hst -> hst { hstRSAClientPublicKey = Just pk })
-
-setClientPrivateKey :: PrivKey -> HandshakeM ()
-setClientPrivateKey pk = modify (\hst -> hst { hstRSAClientPrivateKey = Just pk })
-
 setCertReqSent :: Bool -> HandshakeM ()
 setCertReqSent b = modify (\hst -> hst { hstCertReqSent = b })