fix: cache slot to header data while checking BABE equivocation (#3364)

dot/services.go

@@ -456,7 +456,7 @@ func (nodeBuilder) createGRANDPAService(config *cfg.Config, st *state.Service, k
 }
 
 func (nodeBuilder) createBlockVerifier(st *state.Service) *babe.VerificationManager {
-	return babe.NewVerificationManager(st.Block, st.Epoch)
+	return babe.NewVerificationManager(st.Block, st.Slot, st.Epoch)
 }
 
 func (nodeBuilder) newSyncService(config *cfg.Config, st *state.Service, fg BlockJustificationVerifier,

dot/state/initialize.go

@@ -100,6 +100,7 @@ func (s *Service) Initialise(gen *genesis.Genesis, header *types.Header, t *trie
 		s.Block = blockState
 		s.Epoch = epochState
 		s.Grandpa = grandpaState
+		s.Slot = NewSlotState(db)
 	} else if err = db.Close(); err != nil {
 		return fmt.Errorf("failed to close database: %s", err)
 	}

dot/state/service.go

@@ -34,6 +34,7 @@ type Service struct {
 	Transaction *TransactionState
 	Epoch       *EpochState
 	Grandpa     *GrandpaState
+	Slot        *SlotState
 	closeCh     chan interface{}
 
 	PrunerCfg pruner.Config
@@ -157,6 +158,7 @@ func (s *Service) Start() (err error) {
 		"created state service with head %s, highest number %d and genesis hash %s",
 		s.Block.BestBlockHash(), num, s.Block.genesisHash.String())
 
+	s.Slot = NewSlotState(s.db)
 	return nil
 }
 

dot/state/slot.go

@@ -0,0 +1,193 @@
+// Copyright 2023 ChainSafe Systems (ON)
+// SPDX-License-Identifier: LGPL-3.0-only
+
+package state
+
+import (
+	"bytes"
+	"encoding/binary"
+	"errors"
+	"fmt"
+
+	"github.com/ChainSafe/chaindb"
+	"github.com/ChainSafe/gossamer/dot/types"
+	"github.com/ChainSafe/gossamer/pkg/scale"
+)
+
+const slotTablePrefix = "slot"
+
+// We keep at least this number of slots in database.
+const maxSlotCapacity uint64 = 1000
+
+// We prune slots when they reach this number.
+const pruningBound = 2 * maxSlotCapacity
+
+var (
+	slotHeaderMapKey   = []byte("slot_header_map")
+	slotHeaderStartKey = []byte("slot_header_start")
+)
+
+type SlotState struct {
+	db chaindb.Database
+}
+
+func NewSlotState(db *chaindb.BadgerDB) *SlotState {
+	slotStateDB := chaindb.NewTable(db, slotTablePrefix)
+
+	return &SlotState{
+		db: slotStateDB,
+	}
+}
+
+type headerAndSigner struct {
+	Header *types.Header     `scale:"1"`
+	Signer types.AuthorityID `scale:"2"`
+}
+
+func (s *SlotState) CheckEquivocation(slotNow, slot uint64, header *types.Header,
+	signer types.AuthorityID) (*types.BabeEquivocationProof, error) {
+	// We don't check equivocations for old headers out of our capacity.
+	// checking slotNow is greater than slot to avoid overflow, same as saturating_sub
+	if saturatingSub(slotNow, slot) > maxSlotCapacity {
+		return nil, nil //nolint:nilnil
+	}
+
+	slotEncoded := make([]byte, 8)
+	binary.LittleEndian.PutUint64(slotEncoded, slot)
+
+	currentSlotKey := bytes.Join([][]byte{slotHeaderMapKey, slotEncoded[:]}, nil)
+	encodedHeadersWithSigners, err := s.db.Get(currentSlotKey)
+	if err != nil && !errors.Is(err, chaindb.ErrKeyNotFound) {
+		return nil, fmt.Errorf("getting key slot header map key %d: %w", slot, err)
+	}
+
+	headersWithSigners := make([]headerAndSigner, 0)
+	if len(encodedHeadersWithSigners) > 0 {
+		encodedSliceHeadersWithSigners := make([][]byte, 0)
+
+		err = scale.Unmarshal(encodedHeadersWithSigners, &encodedSliceHeadersWithSigners)
+		if err != nil {
+			return nil, fmt.Errorf("unmarshaling encoded headers with signers: %w", err)
+		}
+
+		for _, encodedHeaderAndSigner := range encodedSliceHeadersWithSigners {
+			// each header and signer instance should have an empty header
+			// so we will be able to scale decode the whole byte stream with
+			// the digests correctly in place
+			decodedHeaderAndSigner := headerAndSigner{
+				Header: types.NewEmptyHeader(),
+			}
+
+			err := scale.Unmarshal(encodedHeaderAndSigner, &decodedHeaderAndSigner)
+			if err != nil {
+				return nil, fmt.Errorf("unmarshaling header with signer: %w", err)
+			}
+
+			headersWithSigners = append(headersWithSigners, decodedHeaderAndSigner)
+		}
+	}
+
+	firstSavedSlot := slot
+	firstSavedSlotEncoded, err := s.db.Get(slotHeaderStartKey)
+	if err != nil && !errors.Is(err, chaindb.ErrKeyNotFound) {
+		return nil, fmt.Errorf("getting key slot header start key: %w", err)
+	}
+
+	if len(firstSavedSlotEncoded) > 0 {
+		firstSavedSlot = binary.LittleEndian.Uint64(firstSavedSlotEncoded)
+	}
+
+	if slotNow < firstSavedSlot {
+		// The code below assumes that slots will be visited sequentially.
+		return nil, nil //nolint:nilnil
+	}
+
+	for _, headerAndSigner := range headersWithSigners {
+		// A proof of equivocation consists of two headers:
+		// 1) signed by the same voter,
+		if headerAndSigner.Signer == signer {
+			// 2) with different hash
+			if headerAndSigner.Header.Hash() != header.Hash() {
+				return &types.BabeEquivocationProof{
+					Slot:         slot,
+					Offender:     signer,
+					FirstHeader:  *headerAndSigner.Header,
+					SecondHeader: *header,
+				}, nil
+			} else {
+				// We don't need to continue in case of duplicated header,
+				// since it's already saved and a possible equivocation
+				// would have been detected before.
+				return nil, nil //nolint:nilnil
+			}
+		}
+	}
+
+	keysToDelete := make([][]byte, 0)
+	newFirstSavedSlot := firstSavedSlot
+
+	if slotNow-firstSavedSlot >= pruningBound {
+		newFirstSavedSlot = saturatingSub(slotNow, maxSlotCapacity)
+
+		for s := firstSavedSlot; s < newFirstSavedSlot; s++ {
+			slotEncoded := make([]byte, 8)
+			binary.LittleEndian.PutUint64(slotEncoded, s)
+
+			toDelete := bytes.Join([][]byte{slotHeaderMapKey, slotEncoded[:]}, nil)
+			keysToDelete = append(keysToDelete, toDelete)
+		}
+	}
+
+	headersWithSigners = append(headersWithSigners, headerAndSigner{Header: header, Signer: signer})
+	encodedHeaderAndSigner := make([][]byte, len(headersWithSigners))
+
+	// encode each header and signer and push to a slice of bytes
+	// that will be scale encoded and stored in the database
+	for idx, headerAndSigner := range headersWithSigners {
+		encoded, err := scale.Marshal(headerAndSigner)
+		if err != nil {
+			return nil, fmt.Errorf("marshalling header and signer: %w", err)
+		}
+
+		encodedHeaderAndSigner[idx] = encoded
+	}
+
+	encodedHeadersWithSigners, err = scale.Marshal(encodedHeaderAndSigner)
+	if err != nil {
+		return nil, fmt.Errorf("marshalling: %w", err)
+	}
+
+	batch := s.db.NewBatch()
+	err = batch.Put(currentSlotKey, encodedHeadersWithSigners)
+	if err != nil {
+		return nil, fmt.Errorf("while batch putting encoded headers with signers: %w", err)
+	}
+
+	newFirstSavedSlotEncoded := make([]byte, 8)
+	binary.LittleEndian.PutUint64(newFirstSavedSlotEncoded, newFirstSavedSlot)
+	err = batch.Put(slotHeaderStartKey, newFirstSavedSlotEncoded)
+	if err != nil {
+		return nil, fmt.Errorf("while batch putting encoded new first saved slot: %w", err)
+	}
+
+	for _, toDelete := range keysToDelete {
+		err := batch.Del(toDelete)
+		if err != nil {
+			return nil, fmt.Errorf("while batch deleting key %s: %w", string(toDelete), err)
+		}
+	}
+
+	err = batch.Flush()
+	if err != nil {
+		return nil, fmt.Errorf("failed to flush batch operations: %w", err)
+	}
+
+	return nil, nil //nolint:nilnil
+}
+
+func saturatingSub(a, b uint64) uint64 {
+	if a > b {
+		return a - b
+	}
+	return 0
+}

dot/state/slot_test.go

@@ -0,0 +1,138 @@
+// Copyright 2023 ChainSafe Systems (ON)
+// SPDX-License-Identifier: LGPL-3.0-only
+
+package state
+
+import (
+	"bytes"
+	"crypto/rand"
+	"encoding/binary"
+	"errors"
+	"io"
+	"testing"
+
+	"github.com/ChainSafe/chaindb"
+	"github.com/ChainSafe/gossamer/dot/types"
+	"github.com/ChainSafe/gossamer/lib/common"
+	"github.com/ChainSafe/gossamer/lib/crypto/sr25519"
+	"github.com/ChainSafe/gossamer/lib/keystore"
+	"github.com/minio/sha256-simd"
+	"github.com/stretchr/testify/require"
+)
+
+func createHeader(t *testing.T, n uint) (header *types.Header) {
+	t.Helper()
+
+	randomBytes := make([]byte, 32)
+	_, err := io.ReadFull(rand.Reader, randomBytes)
+	require.NoError(t, err)
+
+	hasher := sha256.New()
+	_, err = hasher.Write(randomBytes)
+	require.NoError(t, err)
+
+	header = types.NewEmptyHeader()
+	header.Number = n
+
+	// so that different headers for the same number get different hashes
+	header.ParentHash = common.NewHash(hasher.Sum(nil))
+
+	header.Hash()
+	return header
+}
+
+func checkSlotToMapKeyExists(t *testing.T, db chaindb.Database, slotNumber uint64) bool {
+	t.Helper()
+
+	slotEncoded := make([]byte, 8)
+	binary.LittleEndian.PutUint64(slotEncoded, slotNumber)
+
+	slotToHeaderKey := bytes.Join([][]byte{slotHeaderMapKey, slotEncoded[:]}, nil)
+
+	_, err := db.Get(slotToHeaderKey)
+	if err != nil {
+		if errors.Is(err, chaindb.ErrKeyNotFound) {
+			return false
+		}
+
+		t.Fatalf("unexpected error while getting key: %s", err)
+	}
+
+	return true
+}
+
+func Test_checkEquivocation(t *testing.T) {
+	inMemoryDB, err := chaindb.NewBadgerDB(&chaindb.Config{
+		DataDir:  t.TempDir(),
+		InMemory: true,
+	})
+	require.NoError(t, err)
+
+	kr, err := keystore.NewSr25519Keyring()
+	require.NoError(t, err)
+
+	alicePublicKey := kr.KeyAlice.Public().(*sr25519.PublicKey)
+	aliceAuthorityID := types.AuthorityID(alicePublicKey.AsBytes())
+
+	header1 := createHeader(t, 1) // @ slot 2
+	header2 := createHeader(t, 2) // @ slot 2
+	header3 := createHeader(t, 2) // @ slot 4
+	header4 := createHeader(t, 3) // @ slot MAX_SLOT_CAPACITY + 4
+	header5 := createHeader(t, 4) // @ slot MAX_SLOT_CAPACITY + 4
+	header6 := createHeader(t, 3) // @ slot 4
+
+	slotState := NewSlotState(inMemoryDB)
+
+	// It's ok to sign same headers.
+	equivProf, err := slotState.CheckEquivocation(2, 2, header1, aliceAuthorityID)
+	require.NoError(t, err)
+	require.Nil(t, equivProf)
+
+	equivProf, err = slotState.CheckEquivocation(3, 2, header1, aliceAuthorityID)
+	require.NoError(t, err)
+	require.Nil(t, equivProf)
+
+	// But not two different headers at the same slot.
+	equivProf, err = slotState.CheckEquivocation(4, 2, header2, aliceAuthorityID)
+	require.NoError(t, err)
+	require.NotNil(t, equivProf)
+	require.Equal(t, &types.BabeEquivocationProof{
+		Slot:         2,
+		Offender:     aliceAuthorityID,
+		FirstHeader:  *header1,
+		SecondHeader: *header2,
+	}, equivProf)
+
+	// Different slot is ok.
+	equivProf, err = slotState.CheckEquivocation(5, 4, header3, aliceAuthorityID)
+	require.NoError(t, err)
+	require.Nil(t, equivProf)
+
+	// Here we trigger pruning and save header 4.
+	equivProf, err = slotState.CheckEquivocation(
+		pruningBound+2, maxSlotCapacity+4, header4, aliceAuthorityID)
+	require.NoError(t, err)
+	require.Nil(t, equivProf)
+
+	require.False(t, checkSlotToMapKeyExists(t, slotState.db, 2))
+	require.False(t, checkSlotToMapKeyExists(t, slotState.db, 4))
+
+	// This fails because header 5 is an equivocation of header 4.
+	equivProf, err = slotState.CheckEquivocation(
+		pruningBound+3, maxSlotCapacity+4, header5, aliceAuthorityID)
+	require.NoError(t, err)
+	require.NotNil(t, equivProf)
+
+	require.Equal(t, &types.BabeEquivocationProof{
+		Slot:         maxSlotCapacity + 4,
+		Offender:     aliceAuthorityID,
+		FirstHeader:  *header4,
+		SecondHeader: *header5,
+	}, equivProf)
+
+	// This is ok because we pruned the corresponding header. Shows that we are pruning.
+	equivProf, err = slotState.CheckEquivocation(
+		pruningBound+4, 4, header6, aliceAuthorityID)
+	require.NoError(t, err)
+	require.Nil(t, equivProf)
+}