netsync/tool_test.go

   1 package netsync
   2
   3 import (
   4         "errors"
   5         "math/rand"
   6         "net"
   7         "time"
   8
   9         wire "github.com/tendermint/go-wire"
  10
  11         "github.com/bytom/consensus"
  12         "github.com/bytom/protocol/bc"
  13         "github.com/bytom/protocol/bc/types"
  14         "github.com/bytom/test/mock"
  15 )
  16
  17 type P2PPeer struct {
  18         id   string
  19         ip   *net.IPAddr
  20         flag consensus.ServiceFlag
  21
  22         srcPeer    *P2PPeer
  23         remoteNode *SyncManager
  24         msgCh      chan []byte
  25         async      bool
  26 }
  27
  28 func NewP2PPeer(addr, id string, flag consensus.ServiceFlag) *P2PPeer {
  29         return &P2PPeer{
  30                 id:    id,
  31                 ip:    &net.IPAddr{IP: net.ParseIP(addr)},
  32                 flag:  flag,
  33                 msgCh: make(chan []byte),
  34                 async: false,
  35         }
  36 }
  37
  38 func (p *P2PPeer) Addr() net.Addr {
  39         return p.ip
  40 }
  41
  42 func (p *P2PPeer) ID() string {
  43         return p.id
  44 }
  45
  46 func (p *P2PPeer) ServiceFlag() consensus.ServiceFlag {
  47         return p.flag
  48 }
  49
  50 func (p *P2PPeer) SetConnection(srcPeer *P2PPeer, node *SyncManager) {
  51         p.srcPeer = srcPeer
  52         p.remoteNode = node
  53 }
  54
  55 func (p *P2PPeer) TrySend(b byte, msg interface{}) bool {
  56         msgBytes := wire.BinaryBytes(msg)
  57         if p.async {
  58                 p.msgCh <- msgBytes
  59         } else {
  60                 msgType, msg, _ := DecodeMessage(msgBytes)
  61                 p.remoteNode.processMsg(p.srcPeer, msgType, msg)
  62         }
  63         return true
  64 }
  65
  66 func (p *P2PPeer) setAsync(b bool) {
  67         p.async = b
  68 }
  69
  70 func (p *P2PPeer) postMan() {
  71         for msgBytes := range p.msgCh {
  72                 msgType, msg, _ := DecodeMessage(msgBytes)
  73                 p.remoteNode.processMsg(p.srcPeer, msgType, msg)
  74         }
  75 }
  76
  77 type PeerSet struct{}
  78
  79 func NewPeerSet() *PeerSet {
  80         return &PeerSet{}
  81 }
  82
  83 func (ps *PeerSet) AddBannedPeer(string) error { return nil }
  84 func (ps *PeerSet) StopPeerGracefully(string)  {}
  85
  86 type NetWork struct {
  87         nodes map[*SyncManager]P2PPeer
  88 }
  89
  90 func NewNetWork() *NetWork {
  91         return &NetWork{map[*SyncManager]P2PPeer{}}
  92 }
  93
  94 func (nw *NetWork) Register(node *SyncManager, addr, id string, flag consensus.ServiceFlag) {
  95         peer := NewP2PPeer(addr, id, flag)
  96         nw.nodes[node] = *peer
  97 }
  98
  99 func (nw *NetWork) HandsShake(nodeA, nodeB *SyncManager) (*P2PPeer, *P2PPeer, error) {
 100         B2A, ok := nw.nodes[nodeA]
 101         if !ok {
 102                 return nil, nil, errors.New("can't find nodeA's p2p peer on network")
 103         }
 104         A2B, ok := nw.nodes[nodeB]
 105         if !ok {
 106                 return nil, nil, errors.New("can't find nodeB's p2p peer on network")
 107         }
 108
 109         A2B.SetConnection(&B2A, nodeB)
 110         B2A.SetConnection(&A2B, nodeA)
 111
 112         nodeA.handleStatusRequestMsg(&A2B)
 113         nodeB.handleStatusRequestMsg(&B2A)
 114
 115         A2B.setAsync(true)
 116         B2A.setAsync(true)
 117         return &B2A, &A2B, 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 }
 160
 161 func mockTxs(txCount int) ([]*types.Tx, []*bc.Tx) {
 162         var txs []*types.Tx
 163         var bcTxs []*bc.Tx
 164         for i := 0; i < txCount; i++ {
 165                 trueProg := mockControlProgram(60)
 166                 assetID := bc.ComputeAssetID(trueProg, 1, &bc.EmptyStringHash)
 167                 now := []byte(time.Now().String())
 168                 issuanceInp := types.NewIssuanceInput(now, 1, trueProg, nil, nil)
 169                 tx := types.NewTx(types.TxData{
 170                         Version: 1,
 171                         Inputs:  []*types.TxInput{issuanceInp},
 172                         Outputs: []*types.TxOutput{types.NewTxOutput(assetID, 1, trueProg)},
 173                 })
 174                 txs = append(txs, tx)
 175                 bcTxs = append(bcTxs, tx.Tx)
 176         }
 177         return txs, bcTxs
 178 }
 179
 180 func mockControlProgram(length int) []byte {
 181         var cp []byte
 182         for i := 0; i < length; i++ {
 183                 cp = append(cp, byte(rand.Intn(1<<8)))
 184         }
 185         return cp
 186 }