Continue to rekey with DH by sharing part of the code with the initia…

…l KEX, and don't recurse into the rekeying. We get farther now, but still die

src/Ssh/KeyExchange.hs

@@ -80,6 +80,12 @@ doKex clientVersionString serverVersionString clientKEXAlgos clientHostKeys clie
     sPutPacket clientKex
     serverKex <- sGetPacket -- TODO assert this is a KEXInit packet
 
+    continueKex clientVersionString serverVersionString clientKexInitPayload serverKex clientKEXAlgos clientHostKeys clientCryptos serverCryptos clientHashMacs serverHashMacs
+
+-- | The Kex can be done multiple times, at the moment we have a split between the first and the later ones. But both share the
+--   actual computations, which are located in this function
+--continueKex :: 
+continueKex clientVersionString serverVersionString clientKexInitPayload serverKex clientKEXAlgos clientHostKeys clientCryptos serverCryptos clientHashMacs serverHashMacs = do
     printDebugLifted logLowLevelDebug "ServerKEX before filtering:"
     printDebugLifted logLowLevelDebug $ show serverKex
 
@@ -118,18 +124,27 @@ doKex clientVersionString serverVersionString clientKEXAlgos clientHostKeys clie
                           serverVector = server2ClientIV connectiondata,
                           client2server = SshTransport c2sfun s2cmacfun,
                           clientVector = client2ServerIV connectiondata,
-                          maybeConnectionData = Just connectiondata
+                          maybeConnectionData = Just connectiondata,
+                          isRekeying = False -- In case we were rekeying, this has been finished
                          }
 
     printDebugLifted logLowLevelDebug "KEX DONE?"
 
     return connectiondata
 
-
 -- | Start a new key exchange from an existing connection. It returns a new packet handler!
 startRekey :: [KeyExchangeAlgorithm] -> [HostKeyAlgorithm] -> [CryptionAlgorithm] -> [CryptionAlgorithm] -> [HashMac] -> [HashMac] -> SshConnection (Packet -> SshConnection Bool)
 startRekey clientKEXAlgos clientHostKeys clientCryptos serverCryptos clientHashMacs serverHashMacs  = do
     cookie <- MS.liftIO $ BS.unpack `liftM` randBytes 16
-    sPutPacket $ KEXInit B.empty cookie (map kexName clientKEXAlgos) (map hostKeyAlgorithmName clientHostKeys) (map cryptoName clientCryptos) (map cryptoName serverCryptos) (map hashName clientHashMacs) (map hashName serverHashMacs)
+    let clientKex = KEXInit B.empty cookie (map kexName clientKEXAlgos) (map hostKeyAlgorithmName clientHostKeys) (map cryptoName clientCryptos) (map cryptoName serverCryptos) (map hashName clientHashMacs) (map hashName serverHashMacs)
+        clientKexInitPayload = runPut $ putPacket clientKex
+
+    -- We are currently rekeying!
+    MS.modify $ \s -> s { isRekeying = True }
+
+    sPutPacket clientKex
 
-    return undefined
+    previousHandler <- handlePacket `liftM` MS.get
+    return $ \p -> printDebugLifted logLowLevelDebug "WE SHOULD BE REKEYING NOW" >> case p of
+        (KEXInit _ _ _ _ _ _ _ _) -> continueKex undefined undefined clientKexInitPayload p clientKEXAlgos clientHostKeys clientCryptos serverCryptos clientHashMacs serverHashMacs >> return True
+        otherwise                 -> previousHandler p

src/Ssh/Transport.hs

@@ -86,16 +86,18 @@ data SshTransportInfo = SshTransportInfo {
     , c2sStats :: TrafficStats
     , s2cStats :: TrafficStats
 
+    , isRekeying :: Bool -- ^ Keeps track of whether or not we are currently performing a rekey. When True, we shouldn't rekey AGAIN
+
     , handlePacket :: Packet -> SshConnection Bool -- ^ Handle a packet, returns True if it was handled, false if it didn't handle it
 } deriving Show
 
 -- | Convencience method: most data can correctly assume that maybeConnectionData is actually a Just ConnectionData. Unwrap that automatically with a decent name
 connectionData = fromJust . maybeConnectionData
 
 
--- | Provide a convenient wrapper constructor that automatically initiates empty traffic
+-- | Provide a convenient wrapper constructor that automatically initiates empty traffic and isRekeying to False
 mkTransportInfo s hn hka c2s cv cs s2c sv ss cd hp =
-    SshTransportInfo s hn hka c2s cv cs s2c sv ss cd emptyTraffic emptyTraffic hp
+    SshTransportInfo s hn hka c2s cv cs s2c sv ss cd emptyTraffic emptyTraffic False hp
 
 -- | We keep around the SSH Transport State when interacting with the server (it changes for every packet sent/received)
 type SshConnection = MS.StateT SshTransportInfo IO

src/SshClient.hs

@@ -209,7 +209,8 @@ clientLoop username hostname options mCommand cd = do
                         else handleChannel packet >> return ()
 
                     -- Maybe it is time to start rekeying!
-                    MS.lift $ checkToRekey connection
+                    newHandler <- MS.lift $ checkToRekey connection
+                    MS.lift $ MS.modify $ \c -> c { handlePacket = newHandler }
 
                     -- Put back the changed state
                     globalInfo' <- MS.get
@@ -227,14 +228,15 @@ clientLoop username hostname options mCommand cd = do
 
       -- A check to see if we should rekey, and if we should do it: actually send a message to start it
       checkToRekey connection =
-        if bytes > 750 -- For now: a pretty low number to debug it. TODO: use the right values from the RFC
+        if bytes > 750 && canRekey -- For now: a pretty low number to debug it. TODO: use the right values from the RFC
           then do
             printDebugLifted logDebug $ "Already " ++ show bytes ++ " bytes sent, starting a rekey"
             startRekey clientKEXAlgos serverHostKeyAlgos clientCryptos clientCryptos clientHashMacs clientHashMacs
           else
             return $ handlePacket connection
         where
-          bytes = totalBytes . c2sStats $ connection
+          bytes = totalBytes . c2sStats $ connection -- TODO: since last (re)key instead of since beginning
+          canRekey = not $ isRekeying connection -- We don't want to initiate a rekey when we are currently already rekeying!
 
 handlePackets :: Packet -> SshConnection Bool
 handlePackets (Ignore s) = printDebugLifted logDebug "Got a packet 'Ignore', and we print this message" >> return True