netsync/tool_test.go

   1 package netsync
   2
   3 import (
   4         "errors"
   5         "math/rand"
   6         "net"
   7
   8         wire "github.com/tendermint/go-wire"
   9
  10         "github.com/bytom/consensus"
  11         "github.com/bytom/protocol/bc/types"
  12         "github.com/bytom/test/mock"
  13 )
  14
  15 type P2PPeer struct {
  16         id   string
  17         ip   *net.IPAddr
  18         flag consensus.ServiceFlag
  19
  20         srcPeer    *P2PPeer
  21         remoteNode *SyncManager
  22         msgCh      chan []byte
  23         async      bool
  24 }
  25
  26 func NewP2PPeer(addr, id string, flag consensus.ServiceFlag) *P2PPeer {
  27         return &P2PPeer{
  28                 id:    id,
  29                 ip:    &net.IPAddr{IP: net.ParseIP(addr)},
  30                 flag:  flag,
  31                 msgCh: make(chan []byte),
  32                 async: false,
  33         }
  34 }
  35
  36 func (p *P2PPeer) Addr() net.Addr {
  37         return p.ip
  38 }
  39
  40 func (p *P2PPeer) ID() string {
  41         return p.id
  42 }
  43
  44 func (p *P2PPeer) ServiceFlag() consensus.ServiceFlag {
  45         return p.flag
  46 }
  47
  48 func (p *P2PPeer) SetConnection(srcPeer *P2PPeer, node *SyncManager) {
  49         p.srcPeer = srcPeer
  50         p.remoteNode = node
  51 }
  52
  53 func (p *P2PPeer) TrySend(b byte, msg interface{}) bool {
  54         msgBytes := wire.BinaryBytes(msg)
  55         if p.async {
  56                 p.msgCh <- msgBytes
  57         } else {
  58                 msgType, msg, _ := DecodeMessage(msgBytes)
  59                 p.remoteNode.processMsg(p.srcPeer, msgType, msg)
  60         }
  61         return true
  62 }
  63
  64 func (p *P2PPeer) setAsync(b bool) {
  65         p.async = b
  66 }
  67
  68 func (p *P2PPeer) postMan() {
  69         for msgBytes := range p.msgCh {
  70                 msgType, msg, _ := DecodeMessage(msgBytes)
  71                 p.remoteNode.processMsg(p.srcPeer, msgType, msg)
  72         }
  73 }
  74
  75 type PeerSet struct{}
  76
  77 func NewPeerSet() *PeerSet {
  78         return &PeerSet{}
  79 }
  80
  81 func (ps *PeerSet) AddBannedPeer(string) error { return nil }
  82 func (ps *PeerSet) StopPeerGracefully(string)  {}
  83
  84 type NetWork struct {
  85         nodes map[*SyncManager]P2PPeer
  86 }
  87
  88 func NewNetWork() *NetWork {
  89         return &NetWork{map[*SyncManager]P2PPeer{}}
  90 }
  91
  92 func (nw *NetWork) Register(node *SyncManager, addr, id string, flag consensus.ServiceFlag) {
  93         peer := NewP2PPeer(addr, id, flag)
  94         nw.nodes[node] = *peer
  95 }
  96
  97 func (nw *NetWork) HandsShake(nodeA, nodeB *SyncManager) error {
  98         B2A, ok := nw.nodes[nodeA]
  99         if !ok {
 100                 return errors.New("can't find nodeA's p2p peer on network")
 101         }
 102         A2B, ok := nw.nodes[nodeB]
 103         if !ok {
 104                 return errors.New("can't find nodeB's p2p peer on network")
 105         }
 106
 107         A2B.SetConnection(&B2A, nodeB)
 108         B2A.SetConnection(&A2B, nodeA)
 109         go A2B.postMan()
 110         go B2A.postMan()
 111
 112         nodeA.handleStatusRequestMsg(&A2B)
 113         nodeB.handleStatusRequestMsg(&B2A)
 114
 115         A2B.setAsync(true)
 116         B2A.setAsync(true)
 117         return nil
 118 }
 119
 120 func mockBlocks(startBlock *types.Block, height uint64) []*types.Block {
 121         blocks := []*types.Block{}
 122         indexBlock := &types.Block{}
 123         if startBlock == nil {
 124                 indexBlock = &types.Block{BlockHeader: types.BlockHeader{Nonce: uint64(rand.Uint32())}}
 125                 blocks = append(blocks, indexBlock)
 126         } else {
 127                 indexBlock = startBlock
 128         }
 129
 130         for indexBlock.Height < height {
 131                 block := &types.Block{
 132                         BlockHeader: types.BlockHeader{
 133                                 Height:            indexBlock.Height + 1,
 134                                 PreviousBlockHash: indexBlock.Hash(),
 135                                 Nonce:             uint64(rand.Uint32()),
 136                         },
 137                 }
 138                 blocks = append(blocks, block)
 139                 indexBlock = block
 140         }
 141         return blocks
 142 }
 143
 144 func mockSync(blocks []*types.Block) *SyncManager {
 145         chain := mock.NewChain()
 146         peers := newPeerSet(NewPeerSet())
 147         chain.SetBestBlockHeader(&blocks[len(blocks)-1].BlockHeader)
 148         for _, block := range blocks {
 149                 chain.SetBlockByHeight(block.Height, block)
 150         }
 151
 152         genesis, _ := chain.GetHeaderByHeight(0)
 153         return &SyncManager{
 154                 genesisHash: genesis.Hash(),
 155                 chain:       chain,
 156                 blockKeeper: newBlockKeeper(chain, peers),
 157                 peers:       peers,
 158         }
 159 }