OSDN Git Service

Abstract p2p security module
[bytom/bytom.git] / 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         "github.com/tendermint/tmlibs/flowrate"
11
12         "github.com/bytom/consensus"
13         "github.com/bytom/protocol/bc"
14         "github.com/bytom/protocol/bc/types"
15         "github.com/bytom/test/mock"
16 )
17
18 type P2PPeer struct {
19         id   string
20         ip   *net.IPAddr
21         flag consensus.ServiceFlag
22
23         srcPeer    *P2PPeer
24         remoteNode *SyncManager
25         msgCh      chan []byte
26         async      bool
27 }
28
29 func NewP2PPeer(addr, id string, flag consensus.ServiceFlag) *P2PPeer {
30         return &P2PPeer{
31                 id:    id,
32                 ip:    &net.IPAddr{IP: net.ParseIP(addr)},
33                 flag:  flag,
34                 msgCh: make(chan []byte),
35                 async: false,
36         }
37 }
38
39 func (p *P2PPeer) Addr() net.Addr {
40         return p.ip
41 }
42
43 func (p *P2PPeer) ID() string {
44         return p.id
45 }
46
47 func (p *P2PPeer) IsLAN() bool {
48         return false
49 }
50
51 func (p *P2PPeer) ServiceFlag() consensus.ServiceFlag {
52         return p.flag
53 }
54
55 func (p *P2PPeer) SetConnection(srcPeer *P2PPeer, node *SyncManager) {
56         p.srcPeer = srcPeer
57         p.remoteNode = node
58 }
59
60 func (p *P2PPeer) TrafficStatus() (*flowrate.Status, *flowrate.Status) {
61         return nil, nil
62 }
63
64 func (p *P2PPeer) TrySend(b byte, msg interface{}) bool {
65         msgBytes := wire.BinaryBytes(msg)
66         if p.async {
67                 p.msgCh <- msgBytes
68         } else {
69                 msgType, msg, _ := DecodeMessage(msgBytes)
70                 p.remoteNode.processMsg(p.srcPeer, msgType, msg)
71         }
72         return true
73 }
74
75 func (p *P2PPeer) setAsync(b bool) {
76         p.async = b
77 }
78
79 func (p *P2PPeer) postMan() {
80         for msgBytes := range p.msgCh {
81                 msgType, msg, _ := DecodeMessage(msgBytes)
82                 p.remoteNode.processMsg(p.srcPeer, msgType, msg)
83         }
84 }
85
86 type PeerSet struct{}
87
88 func NewPeerSet() *PeerSet {
89         return &PeerSet{}
90 }
91
92 func (ps *PeerSet) IsBanned(peerID string, level byte, reason string) bool {
93         return false
94 }
95
96 func (ps *PeerSet) StopPeerGracefully(string) {}
97
98 type NetWork struct {
99         nodes map[*SyncManager]P2PPeer
100 }
101
102 func NewNetWork() *NetWork {
103         return &NetWork{map[*SyncManager]P2PPeer{}}
104 }
105
106 func (nw *NetWork) Register(node *SyncManager, addr, id string, flag consensus.ServiceFlag) {
107         peer := NewP2PPeer(addr, id, flag)
108         nw.nodes[node] = *peer
109 }
110
111 func (nw *NetWork) HandsShake(nodeA, nodeB *SyncManager) (*P2PPeer, *P2PPeer, error) {
112         B2A, ok := nw.nodes[nodeA]
113         if !ok {
114                 return nil, nil, errors.New("can't find nodeA's p2p peer on network")
115         }
116         A2B, ok := nw.nodes[nodeB]
117         if !ok {
118                 return nil, nil, errors.New("can't find nodeB's p2p peer on network")
119         }
120
121         A2B.SetConnection(&B2A, nodeB)
122         B2A.SetConnection(&A2B, nodeA)
123
124         nodeA.handleStatusRequestMsg(&A2B)
125         nodeB.handleStatusRequestMsg(&B2A)
126
127         A2B.setAsync(true)
128         B2A.setAsync(true)
129         return &B2A, &A2B, nil
130 }
131
132 func mockBlocks(startBlock *types.Block, height uint64) []*types.Block {
133         blocks := []*types.Block{}
134         indexBlock := &types.Block{}
135         if startBlock == nil {
136                 indexBlock = &types.Block{BlockHeader: types.BlockHeader{Nonce: uint64(rand.Uint32())}}
137                 blocks = append(blocks, indexBlock)
138         } else {
139                 indexBlock = startBlock
140         }
141
142         for indexBlock.Height < height {
143                 block := &types.Block{
144                         BlockHeader: types.BlockHeader{
145                                 Height:            indexBlock.Height + 1,
146                                 PreviousBlockHash: indexBlock.Hash(),
147                                 Nonce:             uint64(rand.Uint32()),
148                         },
149                 }
150                 blocks = append(blocks, block)
151                 indexBlock = block
152         }
153         return blocks
154 }
155
156 func mockSync(blocks []*types.Block) *SyncManager {
157         chain := mock.NewChain()
158         peers := newPeerSet(NewPeerSet())
159         chain.SetBestBlockHeader(&blocks[len(blocks)-1].BlockHeader)
160         for _, block := range blocks {
161                 chain.SetBlockByHeight(block.Height, block)
162         }
163
164         genesis, _ := chain.GetHeaderByHeight(0)
165         return &SyncManager{
166                 genesisHash: genesis.Hash(),
167                 chain:       chain,
168                 blockKeeper: newBlockKeeper(chain, peers),
169                 peers:       peers,
170         }
171 }
172
173 func mockTxs(txCount int) ([]*types.Tx, []*bc.Tx) {
174         var txs []*types.Tx
175         var bcTxs []*bc.Tx
176         for i := 0; i < txCount; i++ {
177                 trueProg := mockControlProgram(60)
178                 assetID := bc.ComputeAssetID(trueProg, 1, &bc.EmptyStringHash)
179                 now := []byte(time.Now().String())
180                 issuanceInp := types.NewIssuanceInput(now, 1, trueProg, nil, nil)
181                 tx := types.NewTx(types.TxData{
182                         Version: 1,
183                         Inputs:  []*types.TxInput{issuanceInp},
184                         Outputs: []*types.TxOutput{types.NewTxOutput(assetID, 1, trueProg)},
185                 })
186                 txs = append(txs, tx)
187                 bcTxs = append(bcTxs, tx.Tx)
188         }
189         return txs, bcTxs
190 }
191
192 func mockControlProgram(length int) []byte {
193         var cp []byte
194         for i := 0; i < length; i++ {
195                 cp = append(cp, byte(rand.Intn(1<<8)))
196         }
197         return cp
198 }