From 45da3e17e2ada226d7c82db667d263f8e004cbfc Mon Sep 17 00:00:00 2001
From: obscuren <geffobscura@gmail.com>
Date: Mon, 20 Apr 2015 12:29:02 +0200
Subject: [PATCH] core: added chain head reset to known block

---
 cmd/geth/admin.go            | 49 ++++++++++++++++++++++++++++--------
 core/chain_manager.go        | 24 ++++++++++++++++++
 eth/downloader/downloader.go |  2 ++
 eth/handler.go               |  3 ++-
 4 files changed, 66 insertions(+), 12 deletions(-)

diff --git a/cmd/geth/admin.go b/cmd/geth/admin.go
index c42e91615..1cd0aa2a9 100644
--- a/cmd/geth/admin.go
+++ b/cmd/geth/admin.go
@@ -1,6 +1,7 @@
 package main
 
 import (
+	"errors"
 	"fmt"
 	"os"
 	"time"
@@ -50,15 +51,10 @@ func (js *jsre) adminBindings() {
 	debug.Set("printBlock", js.printBlock)
 	debug.Set("dumpBlock", js.dumpBlock)
 	debug.Set("getBlockRlp", js.getBlockRlp)
+	debug.Set("setHead", js.setHead)
 }
 
-func (js *jsre) downloadProgress(call otto.FunctionCall) otto.Value {
-	current, max := js.ethereum.Downloader().Stats()
-
-	return js.re.ToVal(fmt.Sprintf("%d/%d", current, max))
-}
-
-func (js *jsre) getBlockRlp(call otto.FunctionCall) otto.Value {
+func (js *jsre) getBlock(call otto.FunctionCall) (*types.Block, error) {
 	var block *types.Block
 	if len(call.ArgumentList) > 0 {
 		if call.Argument(0).IsNumber() {
@@ -68,12 +64,43 @@ func (js *jsre) getBlockRlp(call otto.FunctionCall) otto.Value {
 			hash, _ := call.Argument(0).ToString()
 			block = js.ethereum.ChainManager().GetBlock(common.HexToHash(hash))
 		} else {
-			fmt.Println("invalid argument for dump. Either hex string or number")
+			return nil, errors.New("invalid argument for dump. Either hex string or number")
 		}
-
-	} else {
-		block = js.ethereum.ChainManager().CurrentBlock()
+		return block, nil
 	}
+
+	return nil, errors.New("requires block number or block hash as argument")
+}
+
+func (js *jsre) setHead(call otto.FunctionCall) otto.Value {
+	block, err := js.getBlock(call)
+	if err != nil {
+		fmt.Println(err)
+		return otto.UndefinedValue()
+	}
+
+	if block == nil {
+		fmt.Println("block not found")
+		return otto.UndefinedValue()
+	}
+
+	js.ethereum.ChainManager().SetHead(block)
+	return otto.UndefinedValue()
+}
+
+func (js *jsre) downloadProgress(call otto.FunctionCall) otto.Value {
+	current, max := js.ethereum.Downloader().Stats()
+
+	return js.re.ToVal(fmt.Sprintf("%d/%d", current, max))
+}
+
+func (js *jsre) getBlockRlp(call otto.FunctionCall) otto.Value {
+	block, err := js.getBlock(call)
+	if err != nil {
+		fmt.Println(err)
+		return otto.UndefinedValue()
+	}
+
 	if block == nil {
 		fmt.Println("block not found")
 		return otto.UndefinedValue()
diff --git a/core/chain_manager.go b/core/chain_manager.go
index 76fa3e1ea..d312f2495 100644
--- a/core/chain_manager.go
+++ b/core/chain_manager.go
@@ -109,6 +109,30 @@ func NewChainManager(blockDb, stateDb common.Database, mux *event.TypeMux) *Chai
 	return bc
 }
 
+func (bc *ChainManager) SetHead(block *types.Block) {
+	bc.mu.Lock()
+	defer bc.mu.Unlock()
+
+	for block := bc.currentBlock; block != nil && block.Hash() != block.Hash(); block = bc.GetBlock(block.Header().ParentHash) {
+		bc.removeBlock(block)
+	}
+
+	if bc.cache == nil {
+		bc.cache = NewBlockCache(blockCacheLimit)
+	}
+
+	bc.currentBlock = block
+	bc.makeCache()
+
+	statedb := state.New(block.Root(), bc.stateDb)
+	bc.txState = state.ManageState(statedb)
+	bc.transState = statedb.Copy()
+	bc.setTotalDifficulty(block.Td)
+	bc.setLastBlock()
+	bc.insert(block)
+	bc.setLastBlock()
+}
+
 func (self *ChainManager) Td() *big.Int {
 	self.mu.RLock()
 	defer self.mu.RUnlock()
diff --git a/eth/downloader/downloader.go b/eth/downloader/downloader.go
index addcbcc44..cfc494b2f 100644
--- a/eth/downloader/downloader.go
+++ b/eth/downloader/downloader.go
@@ -472,6 +472,8 @@ func (d *Downloader) process() error {
 			}
 			break
 		} else if err != nil {
+			// immediatly unregister the false peer but do not disconnect
+			d.UnregisterPeer(d.activePeer)
 			// Reset chain completely. This needs much, much improvement.
 			// instead: check all blocks leading down to this block false block and remove it
 			blocks = nil
diff --git a/eth/handler.go b/eth/handler.go
index 5c0660d84..622f22132 100644
--- a/eth/handler.go
+++ b/eth/handler.go
@@ -36,6 +36,7 @@ pm.chainman.InsertChain(blocks)
 
 import (
 	"fmt"
+	"math"
 	"math/big"
 	"sync"
 
@@ -326,7 +327,7 @@ func (pm *ProtocolManager) BroadcastBlock(hash common.Hash, block *types.Block)
 	}
 	// Broadcast block to peer set
 	// XXX due to the current shit state of the network disable the limit
-	//peers = peers[:int(math.Sqrt(float64(len(peers))))]
+	peers = peers[:int(math.Sqrt(float64(len(peers))))]
 	for _, peer := range peers {
 		peer.sendNewBlock(block)
 	}