fix CalculateTD so shorter forks don't take over

chain/block_manager.go

@@ -3,7 +3,7 @@ package chain
 import (
 	"bytes"
 	"container/list"
-	"errors"
+	//	"errors"
 	"fmt"
 	"math/big"
 	"sync"
@@ -266,11 +266,15 @@ func (sm *BlockManager) ProcessWithParent(block, parent *types.Block) (td *big.I
 
 		sm.transState = state.Copy()
 
+		// XXX: the only reference to the eth object
+		// if this could be moved to BlockManager, testing wouldn't require
+		//  the eth object
 		sm.eth.TxPool().RemoveSet(block.Transactions())
 
 		return td, messages, nil
 	} else {
-		return nil, nil, errors.New("total diff failed")
+		// this is a fork, so return the td
+		return td, nil, nil
 	}
 }
 
@@ -288,16 +292,12 @@ func (sm *BlockManager) ApplyDiff(state *state.State, parent, block *types.Block
 }
 
 func (sm *BlockManager) CalculateTD(block *types.Block) (*big.Int, bool) {
-	uncleDiff := new(big.Int)
+	td, err := sm.bc.CalcTotalDiff(block)
-	for _, uncle := range block.Uncles {
+	if err != nil {
-		uncleDiff = uncleDiff.Add(uncleDiff, uncle.Difficulty)
+		fmt.Println(err)
+		return nil, false
 	}
 
-	// TD(genesis_block) = 0 and TD(B) = TD(B.parent) + sum(u.difficulty for u in B.uncles) + B.difficulty
-	td := new(big.Int)
-	td = td.Add(sm.bc.TD, uncleDiff)
-	td = td.Add(td, block.Difficulty)
-
 	// The new TD will only be accepted if the new difficulty is
 	// is greater than the previous.
 	if td.Cmp(sm.bc.TD) > 0 {
@@ -307,7 +307,7 @@ func (sm *BlockManager) CalculateTD(block *types.Block) (*big.Int, bool) {
 		//sm.bc.SetTotalDifficulty(td)
 	}
 
-	return nil, false
+	return td, false
 }
 
 // Validates the current block. Returns an error if the block was invalid,

chain/chain_manager_test.go

@@ -1,19 +1,19 @@
 package chain
 
 import (
+	"container/list"
 	"fmt"
 	"math/big"
 	"testing"
 	"time"
-    "container/list"
 
 	"github.com/ethereum/go-ethereum/chain/types"
+	"github.com/ethereum/go-ethereum/crypto"
 	"github.com/ethereum/go-ethereum/ethdb"
 	"github.com/ethereum/go-ethereum/ethutil"
+	"github.com/ethereum/go-ethereum/event"
 	"github.com/ethereum/go-ethereum/state"
 	"github.com/ethereum/go-ethereum/wire"
-	"github.com/ethereum/go-ethereum/event"
-	"github.com/ethereum/go-ethereum/crypto"
 )
 
 // In these tests, TD = block.Number
@@ -34,6 +34,7 @@ func (self fakeproc) ProcessWithParent(a, b *types.Block) (*big.Int, state.Messa
 
 // So we can generate blocks easily
 type fakePow struct{}
+
 func (f fakePow) Search(block *types.Block, stop <-chan struct{}) []byte { return nil }
 func (f fakePow) Verify(hash []byte, diff *big.Int, nonce []byte) bool   { return true }
 func (f fakePow) GetHashrate() int64                                     { return 0 }
@@ -41,6 +42,7 @@ func (f fakePow) Turbo(bool){}
 
 // We need this guy because ProcessWithParent clears txs from the pool
 type fakeEth struct{}
+
 func (e *fakeEth) BlockManager() *BlockManager                        { return nil }
 func (e *fakeEth) ChainManager() *ChainManager                        { return nil }
 func (e *fakeEth) TxPool() *TxPool                                    { return &TxPool{} }
@@ -54,7 +56,6 @@ func (e *fakeEth) ClientIdentity() wire.ClientIdentity { return nil }
 func (e *fakeEth) Db() ethutil.Database                               { return nil }
 func (e *fakeEth) EventMux() *event.TypeMux                           { return nil }
 
-
 func makechain(cman *ChainManager, max int) *BlockChain {
 	blocks := make(types.Blocks, max)
 	for i := 0; i < max; i++ {