From aa3004b08301e20e4d9f98b6afcda457cc703c90 Mon Sep 17 00:00:00 2001
From: Evgeny Danienko <6655321@bk.ru>
Date: Thu, 5 May 2022 18:08:31 +0300
Subject: [PATCH] fix types

---
 eth/backend.go                      | 10 +---------
 eth/downloader/downloader.go        |  1 +
 eth/downloader/whitelist/service.go | 21 +++++++++++++++++++--
 eth/handler.go                      |  4 +++-
 4 files changed, 24 insertions(+), 12 deletions(-)

diff --git a/eth/backend.go b/eth/backend.go
index 994a7d473f..727525da55 100644
--- a/eth/backend.go
+++ b/eth/backend.go
@@ -616,8 +616,6 @@ Loop:
 
 // handleWhitelistCheckpoint handles the checkpoint whitelist mechanism.
 func (s *Ethereum) handleWhitelistCheckpoint() error {
-	var m sync.Mutex
-
 	ethHandler := (*ethHandler)(s.handler)
 
 	if !ethHandler.chain.Engine().(*bor.Bor).WithoutHeimdall {
@@ -626,14 +624,8 @@ func (s *Ethereum) handleWhitelistCheckpoint() error {
 			return err
 		}
 
-		m.Lock()
 		// Update the checkpoint whitelist map.
-		ethHandler.downloader.EnqueueCheckpointWhitelist(endBlockNum, endBlockHash)
-		// If size of checkpoint whitelist map is greater than 10, remove the oldest entry.
-		if len(ethHandler.downloader.GetCheckpointWhitelist()) > 10 {
-			ethHandler.downloader.DequeueCheckpointWhitelist()
-		}
-		m.Unlock()
+		ethHandler.downloader.ProcessCheckpoint(endBlockNum, endBlockHash)
 	}
 
 	return nil
diff --git a/eth/downloader/downloader.go b/eth/downloader/downloader.go
index 2ee61c3372..ae00350464 100644
--- a/eth/downloader/downloader.go
+++ b/eth/downloader/downloader.go
@@ -156,6 +156,7 @@ type Downloader struct {
 // interface for whitelist service
 type ChainValidator interface {
 	IsValidChain(remoteHeader *types.Header, fetchHeadersByNumber func(number uint64, amount int, skip int, reverse bool) ([]*types.Header, []common.Hash, error)) (bool, error)
+	ProcessCheckpoint(endBlockNum uint64, endBlockHash common.Hash)
 }
 
 // LightChain encapsulates functions required to synchronise a light chain.
diff --git a/eth/downloader/whitelist/service.go b/eth/downloader/whitelist/service.go
index 538b77988f..63851f3058 100644
--- a/eth/downloader/whitelist/service.go
+++ b/eth/downloader/whitelist/service.go
@@ -2,6 +2,7 @@ package whitelist
 
 import (
 	"errors"
+	"sync"
 
 	"github.com/ethereum/go-ethereum/common"
 	"github.com/ethereum/go-ethereum/core/types"
@@ -12,16 +13,19 @@ var (
 	ErrCheckpointMismatch = errors.New("checkpoint mismatch")
 )
 
+// Checkpoint whitelist
 type Service struct {
-	// Checkpoint whitelist
+	m 					sync.RWMutex
 	checkpointWhitelist map[uint64]common.Hash // Checkpoint whitelist, populated by reaching out to heimdall
 	checkpointOrder     []uint64               // Checkpoint order, populated by reaching out to heimdall
+	maxCapacity 		uint
 }
 
-func NewService() *Service {
+func NewService(maxCapacity uint) *Service {
 	return &Service{
 		checkpointWhitelist: make(map[uint64]common.Hash),
 		checkpointOrder:     []uint64{},
+		maxCapacity: 		 maxCapacity,
 	}
 }
 
@@ -61,6 +65,18 @@ func (w *Service) IsValidChain(remoteHeader *types.Header, fetchHeadersByNumber
 	return false, ErrCheckpointMismatch
 }
 
+func (w *Service) ProcessCheckpoint(endBlockNum uint64, endBlockHash common.Hash) {
+	w.m.Lock()
+	defer w.m.Unlock()
+
+	w.EnqueueCheckpointWhitelist(endBlockNum, endBlockHash)
+	// If size of checkpoint whitelist map is greater than 10, remove the oldest entry.
+
+	if len(w.GetCheckpointWhitelist()) > int(w.maxCapacity) {
+		w.DequeueCheckpointWhitelist()
+	}
+}
+
 // PurgeWhitelistMap purges data from checkpoint whitelist map
 func (w *Service) PurgeWhitelistMap() error {
 	for k := range w.checkpointWhitelist {
@@ -73,6 +89,7 @@ func (w *Service) PurgeWhitelistMap() error {
 func (w *Service) EnqueueCheckpointWhitelist(key uint64, val common.Hash) {
 	if _, ok := w.checkpointWhitelist[key]; !ok {
 		log.Debug("Enqueing new checkpoint whitelist", "block number", key, "block hash", val)
+
 		w.checkpointWhitelist[key] = val
 		w.checkpointOrder = append(w.checkpointOrder, key)
 	}
diff --git a/eth/handler.go b/eth/handler.go
index fa0852d69e..96eb3dc7be 100644
--- a/eth/handler.go
+++ b/eth/handler.go
@@ -31,6 +31,7 @@ import (
 	"github.com/ethereum/go-ethereum/core/forkid"
 	"github.com/ethereum/go-ethereum/core/types"
 	"github.com/ethereum/go-ethereum/eth/downloader"
+	"github.com/ethereum/go-ethereum/eth/downloader/whitelist"
 	"github.com/ethereum/go-ethereum/eth/fetcher"
 	"github.com/ethereum/go-ethereum/eth/protocols/eth"
 	"github.com/ethereum/go-ethereum/eth/protocols/snap"
@@ -200,7 +201,8 @@ func newHandler(config *handlerConfig) (*handler, error) {
 	// Construct the downloader (long sync) and its backing state bloom if snap
 	// sync is requested. The downloader is responsible for deallocating the state
 	// bloom when it's done.
-	h.downloader = downloader.New(h.checkpointNumber, config.Database, h.eventMux, h.chain, nil, h.removePeer, success)
+	// todo: it'd better to extract maxCapacity into config
+	h.downloader = downloader.New(h.checkpointNumber, config.Database, h.eventMux, h.chain, nil, h.removePeer, success, whitelist.NewService(10))
 
 	// Construct the fetcher (short sync)
 	validator := func(header *types.Header) error {