Integrate KES agent functionality into ouroboros-consensus (#1620)

This PR supercedes #1487

includes the following squashed commit messages:

- Update to use newest cardano-crypto-class with unsound pure KES
implementation

- Use mlocked KES

- Add KES agent connectivity

- Rebase cleanup

- Handle drop-key messages from KES Agent

- Provide KESAgentClientTrace to BlockForging

- Revert change to MockCrypto and require DSIGN only when running the
KES agent

- Bump kes-agent SRP to remove SerDoc dependency

# Description

Please include a meaningful description of the PR and link the relevant
issues
this PR might resolve.

Also note that:

- New code should be properly tested (even if it does not add new
features).
- The fix for a regression should include a test that reproduces said
regression.

Author: fraser-iohk

cabal.project

@@ -14,7 +14,7 @@ repository cardano-haskell-packages
 -- update either of these.
 index-state:
   -- Bump this if you need newer packages from Hackage
-  , hackage.haskell.org 2025-07-22T09:13:54Z
+  , hackage.haskell.org 2025-08-05T11:23:47Z
   -- Bump this if you need newer packages from CHaP
   , cardano-haskell-packages 2025-08-21T09:41:03Z
 

ouroboros-consensus-cardano/changelog.d/20250130_093803_tdammers_mlocked_kes_rebase.md

@@ -0,0 +1,8 @@
+### Breaking
+
+- Use new mlocked KES API for all internal KES sign key handling.
+- Add finalizers to all block forgings (required by `ouroboros-consensus`).
+- Change `ShelleyLeaderCredentials` to not contain the KES sign key itself
+  anymore. Instead, the `CanBeLeader` data structure now contains a
+  `praosCanBeLeaderCredentialsSource` field, which specifies how to obtain the
+  actual credentials (OpCert and KES SignKey).

ouroboros-consensus-cardano/ouroboros-consensus-cardano.cabal

@@ -156,6 +156,7 @@ library
     cardano-strict-containers,
     cborg ^>=0.2.2,
     containers >=0.5 && <0.8,
+    contra-tracer,
     crypton,
     deepseq,
     formatting >=6.3 && <7.3,
@@ -325,6 +326,8 @@ library unstable-shelley-testlib
     cardano-slotting,
     cardano-strict-containers,
     containers,
+    contra-tracer,
+    kes-agent,
     microlens,
     mtl,
     ouroboros-consensus:{ouroboros-consensus, unstable-consensus-testlib},
@@ -364,6 +367,7 @@ test-suite shelley-test
     cborg,
     constraints,
     containers,
+    contra-tracer,
     filepath,
     measures,
     microlens,
@@ -415,6 +419,7 @@ library unstable-cardano-testlib
     cardano-strict-containers,
     cborg,
     containers,
+    contra-tracer,
     mempack,
     microlens,
     mtl,

ouroboros-consensus-cardano/src/byron/Ouroboros/Consensus/Byron/Node.hs

@@ -143,6 +143,7 @@ byronBlockForging creds =
           slot
           tickedPBftState
     , forgeBlock = \cfg -> return ....: forgeByronBlock cfg
+    , finalize = pure ()
     }
  where
   canBeLeader = mkPBftCanBeLeader creds

ouroboros-consensus-cardano/src/ouroboros-consensus-cardano/Ouroboros/Consensus/Cardano/Node.hs

@@ -55,14 +55,12 @@ import qualified Cardano.Ledger.Api.Transition as L
 import qualified Cardano.Ledger.BaseTypes as SL
 import qualified Cardano.Ledger.Shelley.API as SL
 import Cardano.Prelude (cborError)
-import qualified Cardano.Protocol.TPraos.OCert as Absolute
-  ( KESPeriod (..)
-  , ocertKESPeriod
-  )
+import qualified Cardano.Protocol.TPraos.OCert as Absolute (KESPeriod (..))
 import qualified Codec.CBOR.Decoding as CBOR
 import Codec.CBOR.Encoding (Encoding)
 import qualified Codec.CBOR.Encoding as CBOR
 import Control.Exception (assert)
+import qualified Control.Tracer as Tracer
 import qualified Data.ByteString.Short as Short
 import Data.Functor.These (These1 (..))
 import qualified Data.Map.Strict as Map
@@ -97,10 +95,11 @@ import Ouroboros.Consensus.Ledger.Tables.Utils (forgetLedgerTables)
 import Ouroboros.Consensus.Node.NetworkProtocolVersion
 import Ouroboros.Consensus.Node.ProtocolInfo
 import Ouroboros.Consensus.Node.Run
-import qualified Ouroboros.Consensus.Protocol.Ledger.HotKey as HotKey
 import Ouroboros.Consensus.Protocol.Praos (Praos, PraosParams (..))
+import Ouroboros.Consensus.Protocol.Praos.AgentClient
 import Ouroboros.Consensus.Protocol.Praos.Common
-  ( praosCanBeLeaderOpCert
+  ( PraosCanBeLeader (..)
+  , instantiatePraosCredentials
   )
 import Ouroboros.Consensus.Protocol.TPraos (TPraos, TPraosParams (..))
 import qualified Ouroboros.Consensus.Protocol.TPraos as Shelley
@@ -122,7 +121,6 @@ import qualified Ouroboros.Consensus.Shelley.Node.TPraos as TPraos
 import Ouroboros.Consensus.Storage.Serialisation
 import Ouroboros.Consensus.TypeFamilyWrappers
 import Ouroboros.Consensus.Util.Assert
-import Ouroboros.Consensus.Util.IOLike
 
 {-------------------------------------------------------------------------------
   SerialiseHFC
@@ -569,10 +567,12 @@ data CardanoProtocolParams c = CardanoProtocolParams
 -- for mainnet (check against @'SL.gNetworkId' == 'SL.Mainnet'@).
 protocolInfoCardano ::
   forall c m.
-  (IOLike m, CardanoHardForkConstraints c) =>
+  ( CardanoHardForkConstraints c
+  , KESAgentContext c m
+  ) =>
   CardanoProtocolParams c ->
   ( ProtocolInfo (CardanoBlock c)
-  , m [BlockForging m (CardanoBlock c)]
+  , Tracer.Tracer m KESAgentClientTrace -> m [MkBlockForging m (CardanoBlock c)]
   )
 protocolInfoCardano paramsCardano
   | SL.Mainnet <- SL.sgNetworkId genesisShelley
@@ -585,7 +585,7 @@ protocolInfoCardano paramsCardano
             { pInfoConfig = cfg
             , pInfoInitLedger = initExtLedgerStateCardano
             }
-        , blockForging
+        , pure . mkBlockForgings
         )
  where
   CardanoProtocolParams
@@ -980,46 +980,41 @@ protocolInfoCardano paramsCardano
   -- credentials. If there are multiple Shelley credentials, we merge the
   -- Byron credentials with the first Shelley one but still have separate
   -- threads for the remaining Shelley ones.
-  blockForging :: m [BlockForging m (CardanoBlock c)]
-  blockForging = do
-    shelleyBased <- traverse blockForgingShelleyBased credssShelleyBased
-    let blockForgings :: [NonEmptyOptNP (BlockForging m) (CardanoEras c)]
+  mkBlockForgings :: Tracer.Tracer m KESAgentClientTrace -> [MkBlockForging m (CardanoBlock c)]
+  mkBlockForgings tr = do
+    let shelleyBased = blockForgingShelleyBased tr <$> credssShelleyBased
+        blockForgings :: [m (NonEmptyOptNP (BlockForging m) (CardanoEras c))]
         blockForgings = case (mBlockForgingByron, shelleyBased) of
           (Nothing, shelleys) -> shelleys
           (Just byron, []) -> [byron]
           (Just byron, shelley : shelleys) ->
-            OptNP.zipWith merge byron shelley : shelleys
+            (OptNP.zipWith merge <$> byron <*> shelley) : shelleys
            where
             -- When merging Byron with Shelley-based eras, we should never
             -- merge two from the same era.
             merge (These1 _ _) = error "forgings of the same era"
             merge (This1 x) = x
             merge (That1 y) = y
 
-    return $ hardForkBlockForging "Cardano" <$> blockForgings
+    let mkHardForkBlockForgings ::
+          m (NonEmptyOptNP (BlockForging m) (CardanoEras c)) -> MkBlockForging m (CardanoBlock c)
+        mkHardForkBlockForgings mbfs = MkBlockForging $ do
+          bfs <- mbfs
+          mkBlockForging $ hardForkBlockForging (const "Cardano") (hmap (MkBlockForging . pure) bfs)
+
+    fmap mkHardForkBlockForgings blockForgings
 
-  mBlockForgingByron :: Maybe (NonEmptyOptNP (BlockForging m) (CardanoEras c))
+  mBlockForgingByron :: Maybe (m (NonEmptyOptNP (BlockForging m) (CardanoEras c)))
   mBlockForgingByron = do
     creds <- mCredsByron
-    return $ byronBlockForging creds `OptNP.at` IZ
+    return $ pure $ byronBlockForging creds `OptNP.at` IZ
 
   blockForgingShelleyBased ::
+    Tracer.Tracer m KESAgentClientTrace ->
     ShelleyLeaderCredentials c ->
     m (NonEmptyOptNP (BlockForging m) (CardanoEras c))
-  blockForgingShelleyBased credentials = do
-    let ShelleyLeaderCredentials
-          { shelleyLeaderCredentialsInitSignKey = initSignKey
-          , shelleyLeaderCredentialsCanBeLeader = canBeLeader
-          } = credentials
-
-    hotKey <- do
-      let maxKESEvo :: Word64
-          maxKESEvo = assert (tpraosMaxKESEvo == praosMaxKESEvo) praosMaxKESEvo
-
-          startPeriod :: Absolute.KESPeriod
-          startPeriod = Absolute.ocertKESPeriod $ praosCanBeLeaderOpCert canBeLeader
-
-      HotKey.mkHotKey @m @c initSignKey startPeriod maxKESEvo
+  blockForgingShelleyBased tr credentials = do
+    let canBeLeader = shelleyLeaderCredentialsCanBeLeader credentials
 
     let slotToPeriod :: SlotNo -> Absolute.KESPeriod
         slotToPeriod (SlotNo slot) =
@@ -1028,6 +1023,15 @@ protocolInfoCardano paramsCardano
               fromIntegral $
                 slot `div` praosSlotsPerKESPeriod
 
+        maxKESEvo :: Word64
+        maxKESEvo = assert (tpraosMaxKESEvo == praosMaxKESEvo) praosMaxKESEvo
+
+    hotKey <-
+      instantiatePraosCredentials
+        maxKESEvo
+        tr
+        (praosCanBeLeaderCredentialsSource canBeLeader)
+
     let tpraos ::
           forall era.
           ShelleyEraWithCrypto c (TPraos c) era =>

ouroboros-consensus-cardano/src/shelley/Ouroboros/Consensus/Shelley/HFEras.hs

@@ -18,9 +18,10 @@ module Ouroboros.Consensus.Shelley.HFEras
   , StandardShelleyBlock
   ) where
 
+import Cardano.Protocol.Crypto
 import Ouroboros.Consensus.Protocol.Praos (Praos)
 import qualified Ouroboros.Consensus.Protocol.Praos as Praos
-import Ouroboros.Consensus.Protocol.TPraos (StandardCrypto, TPraos)
+import Ouroboros.Consensus.Protocol.TPraos (TPraos)
 import qualified Ouroboros.Consensus.Protocol.TPraos as TPraos
 import Ouroboros.Consensus.Shelley.Eras
   ( AllegraEra

ouroboros-consensus-cardano/src/shelley/Ouroboros/Consensus/Shelley/Node/Common.hs

@@ -19,12 +19,10 @@ module Ouroboros.Consensus.Shelley.Node.Common
   , shelleyBlockIssuerVKey
   ) where
 
-import Cardano.Crypto.KES (UnsoundPureSignKeyKES)
 import Cardano.Ledger.BaseTypes (unNonZero)
 import qualified Cardano.Ledger.Keys as SL
 import qualified Cardano.Ledger.Shelley.API as SL
 import Cardano.Ledger.Slot
-import Cardano.Protocol.Crypto
 import Data.Text (Text)
 import Ouroboros.Consensus.Block
   ( CannotForge
@@ -59,12 +57,7 @@ import Ouroboros.Consensus.Storage.ImmutableDB
 -------------------------------------------------------------------------------}
 
 data ShelleyLeaderCredentials c = ShelleyLeaderCredentials
-  { shelleyLeaderCredentialsInitSignKey :: UnsoundPureSignKeyKES (KES c)
-  -- ^ The unevolved signing KES key (at evolution 0).
-  --
-  -- Note that this is not inside 'ShelleyCanBeLeader' since it gets evolved
-  -- automatically, whereas 'ShelleyCanBeLeader' does not change.
-  , shelleyLeaderCredentialsCanBeLeader :: PraosCanBeLeader c
+  { shelleyLeaderCredentialsCanBeLeader :: PraosCanBeLeader c
   , shelleyLeaderCredentialsLabel :: Text
   -- ^ Identifier for this set of credentials.
   --

ouroboros-consensus-cardano/src/shelley/Ouroboros/Consensus/Shelley/Node/Praos.hs

@@ -29,9 +29,6 @@ import Ouroboros.Consensus.Protocol.Praos
   , PraosParams (..)
   , praosCheckCanForge
   )
-import Ouroboros.Consensus.Protocol.Praos.Common
-  ( PraosCanBeLeader (praosCanBeLeaderOpCert)
-  )
 import Ouroboros.Consensus.Shelley.Ledger
   ( ShelleyBlock
   , ShelleyCompatible
@@ -56,21 +53,13 @@ praosBlockForging ::
   , IOLike m
   ) =>
   PraosParams ->
+  HotKey.HotKey c m ->
   ShelleyLeaderCredentials c ->
-  m (BlockForging m (ShelleyBlock (Praos c) era))
-praosBlockForging praosParams credentials = do
-  hotKey <- HotKey.mkHotKey @m @c initSignKey startPeriod praosMaxKESEvo
-  pure $ praosSharedBlockForging hotKey slotToPeriod credentials
+  BlockForging m (ShelleyBlock (Praos c) era)
+praosBlockForging praosParams hotKey credentials =
+  praosSharedBlockForging hotKey slotToPeriod credentials
  where
-  PraosParams{praosMaxKESEvo, praosSlotsPerKESPeriod} = praosParams
-
-  ShelleyLeaderCredentials
-    { shelleyLeaderCredentialsInitSignKey = initSignKey
-    , shelleyLeaderCredentialsCanBeLeader = canBeLeader
-    } = credentials
-
-  startPeriod :: Absolute.KESPeriod
-  startPeriod = SL.ocertKESPeriod $ praosCanBeLeaderOpCert canBeLeader
+  PraosParams{praosSlotsPerKESPeriod} = praosParams
 
   slotToPeriod :: SlotNo -> Absolute.KESPeriod
   slotToPeriod (SlotNo slot) =
@@ -95,7 +84,7 @@ praosSharedBlockForging
   ShelleyLeaderCredentials
     { shelleyLeaderCredentialsCanBeLeader = canBeLeader
     , shelleyLeaderCredentialsLabel = label
-    } = do
+    } =
     BlockForging
       { forgeLabel = label <> "_" <> T.pack (L.eraName @era)
       , canBeLeader = canBeLeader
@@ -111,4 +100,5 @@ praosSharedBlockForging
             hotKey
             canBeLeader
             cfg
+      , finalize = HotKey.finalize hotKey
       }