Handle drop-key messages from KES Agent

Author: tdammers

ouroboros-consensus-protocol/src/ouroboros-consensus-protocol/Ouroboros/Consensus/Protocol/Ledger/HotKey.hs

@@ -264,17 +264,21 @@ mkKESState maxKESEvolutions newOCert newKey evolution startPeriod@(Absolute.KESP
       , kesStateKey = KESKey newOCert newKey
     }
 
+type KeyProducer c m =
+  (OCert.OCert c -> KES.SignKeyKES (KES c) -> Word -> Absolute.KESPeriod -> m ())
+  -- ^ Callback that will be invoked when a new key has been received
+  -> m ()
+  -- ^ Callback that will be invoked when a key deletion has been received
+  -> m ()
+
 -- | Create a new 'HotKey' that runs a key-producer action on a separate thread.
 -- The key producer action will receive a callback that can be used to pass
 -- keys into the HotKey; the HotKey will dynamically update its internal state
 -- to reflect new keys as they arrive.
 mkDynamicHotKey ::
      forall m c. (Crypto c, IOLike m)
   => Word64              -- ^ Max KES evolutions
-  -> Maybe (
-        (OCert.OCert c -> KES.SignKeyKES (KES c) -> Word -> Absolute.KESPeriod -> m ())
-        -> m ()
-     )
+  -> Maybe (KeyProducer c m)
   -> m ()
   -> m (HotKey c m)
 mkDynamicHotKey = mkHotKeyWith Nothing
@@ -285,10 +289,7 @@ mkHotKeyWith ::
      forall m c. (Crypto c, IOLike m)
   => Maybe (OCert.OCert c, KES.SignKeyKES (KES c), Word, Absolute.KESPeriod)
   -> Word64              -- ^ Max KES evolutions
-  -> Maybe (
-        (OCert.OCert c -> KES.SignKeyKES (KES c) -> Word -> Absolute.KESPeriod -> m ())
-        -> m ()
-     )
+  -> Maybe (KeyProducer c m)
   -> m ()
   -> m (HotKey c m)
 mkHotKeyWith initialStateMay maxKESEvolutions keyThreadMay finalizer = do
@@ -298,6 +299,8 @@ mkHotKeyWith initialStateMay maxKESEvolutions keyThreadMay finalizer = do
           modifyMVar_ varKESState $ \oldState -> do
             _ <- poisonState oldState
             return $ mkKESState maxKESEvolutions newOCert newKey evolution startPeriod
+        unset =
+          modifyMVar_ varKESState $ poisonState
 
     forM_ initialStateMay $ \(newOCert, newKey, evolution, startPeriod) ->
             set newOCert newKey evolution startPeriod
@@ -306,7 +309,7 @@ mkHotKeyWith initialStateMay maxKESEvolutions keyThreadMay finalizer = do
       Just keyThread -> do
         keyThreadAsync <- async $ do
                             labelThisThread "HotKey receiver"
-                            keyThread set
+                            keyThread set unset
         return (cancel keyThreadAsync >> finalizer)
       Nothing ->
         return finalizer
@@ -327,8 +330,7 @@ mkHotKeyWith initialStateMay maxKESEvolutions keyThreadMay finalizer = do
               KESKey _ key -> do
                 let evolution = kesEvolution kesStateInfo
                 KES.signedKES () evolution toSign key
-      , forget = do
-          modifyMVar_ varKESState $ poisonState
+      , forget = unset
       , finalize_ = finalizer'
       }
   where

ouroboros-consensus-protocol/src/ouroboros-consensus-protocol/Ouroboros/Consensus/Protocol/Praos/AgentClient.hs

@@ -82,6 +82,9 @@ class ( Crypto c
 instance AgentCrypto StandardCrypto where
   type ACrypto StandardCrypto = Agent.StandardCrypto
 
+convertPeriod :: Agent.KESPeriod -> OCert.KESPeriod
+convertPeriod (Agent.KESPeriod p) = OCert.KESPeriod p
+
 convertOCert :: ( AgentCrypto c
                 )
               => Agent.OCert (ACrypto c) -> OCert.OCert c
@@ -130,9 +133,10 @@ runKESAgentClient :: forall m c.
                      )
                   => Tracer m KESAgentClientTrace
                   -> FilePath
-                  -> (OCert.OCert c -> SignKeyKES (KES c) -> Word -> m ())
+                  -> (OCert.OCert c -> SignKeyKES (KES c) -> Word -> OCert.KESPeriod -> m ())
+                  -> m ()
                   -> m ()
-runKESAgentClient tracer path handleKey = do
+runKESAgentClient tracer path handleKey handleDropKey = do
   withAgentContext $ \snocket -> do
     forever $ do
       Agent.runServiceClient
@@ -151,9 +155,11 @@ runKESAgentClient tracer path handleKey = do
                 -- finishes.
                 _ <- acquireCRef skpRef
                 withCRefValue skpRef $ \(SignKeyWithPeriodKES sk p) ->
-                  handleKey (convertOCert ocert) sk p
+                  handleKey (convertOCert ocert) sk p (convertPeriod $ Agent.ocertKESPeriod ocert)
+                return Agent.RecvOK
+              _ -> do
+                handleDropKey
                 return Agent.RecvOK
-              _ -> undefined -- TODO
         )
         (contramap KESAgentClientTrace tracer)
         `catch` ( \(_e :: AsyncCancelled) ->

ouroboros-consensus-protocol/src/ouroboros-consensus-protocol/Ouroboros/Consensus/Protocol/Praos/Common.hs

@@ -300,10 +300,8 @@ instantiatePraosCredentials maxKESEvolutions (PraosCredentialsUnsound ocert skUn
 instantiatePraosCredentials maxKESEvolutions (PraosCredentialsAgent path) = do
   HotKey.mkDynamicHotKey
       maxKESEvolutions
-      (Just $ \send -> do
-        let handleKey ocert sk p = do
-              send ocert sk p (OCert.ocertKESPeriod ocert)
-        runKESAgentClient nullTracer path handleKey
+      (Just $ \handleKey handleDrop -> do
+        runKESAgentClient nullTracer path handleKey handleDrop
       )
       (pure ())