fix(lib/babe): add pre-runtime digest before calling initialize_block (…

…ChainSafe#1581)
timwu20 · Dec 6, 2021 · f0060ae · f0060ae
1 parent d2197b0
commit f0060ae
Show file tree

Hide file tree

Showing 5 changed files with 20 additions and 18 deletions.
diff --git a/dot/types/digest.go b/dot/types/digest.go
@@ -34,6 +34,11 @@ func NewEmptyDigest() Digest {
 	return []DigestItem{}
 }
 
+// NewDigest returns a new Digest from the given DigestItems
+func NewDigest(items ...DigestItem) Digest {
+	return items
+}
+
 // Encode returns the SCALE encoded digest
 func (d *Digest) Encode() ([]byte, error) {
 	enc, err := scale.Encode(big.NewInt(int64(len(*d))))

diff --git a/lib/babe/build.go b/lib/babe/build.go
@@ -49,7 +49,7 @@ func (b *Service) buildBlock(parent *types.Header, slot Slot) (*types.Block, err
 
 	// create new block header
 	number := big.NewInt(0).Add(parent.Number, big.NewInt(1))
-	header, err := types.NewHeader(parent.Hash(), common.Hash{}, common.Hash{}, number, types.NewEmptyDigest())
+	header, err := types.NewHeader(parent.Hash(), common.Hash{}, common.Hash{}, number, types.NewDigest(preDigest))
 	if err != nil {
 		return nil, err
 	}
@@ -87,9 +87,6 @@ func (b *Service) buildBlock(parent *types.Header, slot Slot) (*types.Block, err
 	header.ParentHash = parent.Hash()
 	header.Number.Add(parent.Number, big.NewInt(1))
 
-	// add BABE header to digest
-	header.Digest = append(header.Digest, preDigest)
-
 	// create seal and add to digest
 	seal, err := b.buildBlockSeal(header)
 	if err != nil {

diff --git a/lib/babe/build_test.go b/lib/babe/build_test.go
@@ -130,13 +130,16 @@ func TestBuildBlock_ok(t *testing.T) {
 		Digest:     types.Digest{preDigest},
 	}
 
-	// remove seal from built block, since we can't predict the signature
-	block.Header.Digest = block.Header.Digest[:1]
 	require.Equal(t, expectedBlockHeader.ParentHash, block.Header.ParentHash)
 	require.Equal(t, expectedBlockHeader.Number, block.Header.Number)
 	require.NotEqual(t, block.Header.StateRoot, emptyHash)
 	require.NotEqual(t, block.Header.ExtrinsicsRoot, emptyHash)
-	require.Equal(t, expectedBlockHeader.Digest, block.Header.Digest)
+	require.Equal(t, 3, len(block.Header.Digest))
+	require.Equal(t, preDigest, block.Header.Digest[0])
+	require.Equal(t, types.PreRuntimeDigestType, block.Header.Digest[0].Type())
+	require.Equal(t, types.ConsensusDigestType, block.Header.Digest[1].Type())
+	require.Equal(t, types.SealDigestType, block.Header.Digest[2].Type())
+	require.Equal(t, types.NextEpochDataType, block.Header.Digest[1].(*types.ConsensusDigest).DataType())
 
 	// confirm block body is correct
 	extsRes, err := block.Body.AsExtrinsics()

diff --git a/lib/runtime/wasmer/exports.go b/lib/runtime/wasmer/exports.go
@@ -148,7 +148,6 @@ func (in *Instance) FinalizeBlock() (*types.Header, error) {
 func (in *Instance) ExecuteBlock(block *types.Block) ([]byte, error) {
 	// copy block since we're going to modify it
 	b := block.DeepCopy()
-	b.Header.Digest = types.NewEmptyDigest()
 
 	if in.version == nil {
 		var err error
@@ -158,17 +157,14 @@ func (in *Instance) ExecuteBlock(block *types.Block) ([]byte, error) {
 		}
 	}
 
-	// TODO: hack since substrate node_runtime can't seem to handle BABE pre-runtime digests
-	// with type prefix (ie Primary, Secondary...)
-	if bytes.Equal(in.version.SpecName(), []byte("kusama")) || bytes.Equal(in.version.SpecName(), []byte("polkadot")) {
-		// remove seal digest only
-		for _, d := range block.Header.Digest {
-			if d.Type() == types.SealDigestType {
-				continue
-			}
-
-			b.Header.Digest = append(b.Header.Digest, d)
+	// remove seal digest only
+	b.Header.Digest = types.NewEmptyDigest()
+	for _, d := range block.Header.Digest {
+		if d.Type() == types.SealDigestType {
+			continue
 		}
+
+		b.Header.Digest = append(b.Header.Digest, d)
 	}
 
 	bdEnc, err := b.Encode()

diff --git a/lib/runtime/wasmer/exports_test.go b/lib/runtime/wasmer/exports_test.go
@@ -453,6 +453,7 @@ func TestInstance_ExecuteBlock_NodeRuntime(t *testing.T) {
 	require.NoError(t, err)
 	instance.SetContextStorage(parentState)
 
+	block.Header.Digest = types.NewEmptyDigest()
 	_, err = instance.ExecuteBlock(block)
 	require.NoError(t, err)
 }