diff options
| -rw-r--r-- | whisper/peer_test.go | 50 | ||||
| -rw-r--r-- | whisper/whisper.go | 58 |
2 files changed, 86 insertions, 22 deletions
diff --git a/whisper/peer_test.go b/whisper/peer_test.go index de67b2463..9008cdc59 100644 --- a/whisper/peer_test.go +++ b/whisper/peer_test.go @@ -190,3 +190,53 @@ func TestPeerDeliver(t *testing.T) { t.Fatalf("repeating message arrived") } } + +func TestPeerMessageExpiration(t *testing.T) { + // Start a tester and execute the handshake + tester, err := startTestPeerInited() + if err != nil { + t.Fatalf("failed to start initialized peer: %v", err) + } + defer tester.stream.Close() + + // Fetch the peer instance for later inspection + tester.client.peerMu.RLock() + if peers := len(tester.client.peers); peers != 1 { + t.Fatalf("peer pool size mismatch: have %v, want %v", peers, 1) + } + var peer *peer + for peer, _ = range tester.client.peers { + break + } + tester.client.peerMu.RUnlock() + + // Construct a message and pass it through the tester + message := NewMessage([]byte("peer test message")) + envelope, err := message.Wrap(DefaultPoW, Options{ + TTL: time.Second, + }) + if err != nil { + t.Fatalf("failed to wrap message: %v", err) + } + if err := tester.client.Send(envelope); err != nil { + t.Fatalf("failed to send message: %v", err) + } + payload := []interface{}{envelope} + if err := p2p.ExpectMsg(tester.stream, messagesCode, payload); err != nil { + t.Fatalf("message mismatch: %v", err) + } + // Check that the message is inside the cache + if !peer.known.Has(envelope.Hash()) { + t.Fatalf("message not found in cache") + } + // Discard messages until expiration and check cache again + exp := time.Now().Add(time.Second + expirationCycle) + for time.Now().Before(exp) { + if err := p2p.ExpectMsg(tester.stream, messagesCode, []interface{}{}); err != nil { + t.Fatalf("message mismatch: %v", err) + } + } + if peer.known.Has(envelope.Hash()) { + t.Fatalf("message not expired from cache") + } +} diff --git a/whisper/whisper.go b/whisper/whisper.go index f04075e1f..48efff622 100644 --- a/whisper/whisper.go +++ b/whisper/whisper.go @@ -46,22 +46,26 @@ type Whisper struct { protocol p2p.Protocol filters *filter.Filters - mmu sync.RWMutex // Message mutex to sync the below pool - messages map[common.Hash]*Envelope // Pool of messages currently tracked by this node - expiry map[uint32]*set.SetNonTS // Message expiration pool (TODO: something lighter) + keys map[string]*ecdsa.PrivateKey - quit chan struct{} + messages map[common.Hash]*Envelope // Pool of messages currently tracked by this node + expirations map[uint32]*set.SetNonTS // Message expiration pool (TODO: something lighter) + poolMu sync.RWMutex // Mutex to sync the message and expiration pools - keys map[string]*ecdsa.PrivateKey + peers map[*peer]struct{} // Set of currently active peers + peerMu sync.RWMutex // Mutex to sync the active peer set + + quit chan struct{} } func New() *Whisper { whisper := &Whisper{ - messages: make(map[common.Hash]*Envelope), - filters: filter.New(), - expiry: make(map[uint32]*set.SetNonTS), - quit: make(chan struct{}), - keys: make(map[string]*ecdsa.PrivateKey), + filters: filter.New(), + keys: make(map[string]*ecdsa.PrivateKey), + messages: make(map[common.Hash]*Envelope), + expirations: make(map[uint32]*set.SetNonTS), + peers: make(map[*peer]struct{}), + quit: make(chan struct{}), } whisper.filters.Start() @@ -179,6 +183,16 @@ func (self *Whisper) handlePeer(peer *p2p.Peer, rw p2p.MsgReadWriter) error { whisperPeer.start() defer whisperPeer.stop() + // Start tracking the active peer + self.peerMu.Lock() + self.peers[whisperPeer] = struct{}{} + self.peerMu.Unlock() + + defer func() { + self.peerMu.Lock() + delete(self.peers, whisperPeer) + self.peerMu.Unlock() + }() // Read and process inbound messages directly to merge into client-global state for { // Fetch the next packet and decode the contained envelopes @@ -206,8 +220,8 @@ func (self *Whisper) handlePeer(peer *p2p.Peer, rw p2p.MsgReadWriter) error { // whisper network. It also inserts the envelope into the expiration pool at the // appropriate time-stamp. func (self *Whisper) add(envelope *Envelope) error { - self.mmu.Lock() - defer self.mmu.Unlock() + self.poolMu.Lock() + defer self.poolMu.Unlock() // Insert the message into the tracked pool hash := envelope.Hash() @@ -218,11 +232,11 @@ func (self *Whisper) add(envelope *Envelope) error { self.messages[hash] = envelope // Insert the message into the expiration pool for later removal - if self.expiry[envelope.Expiry] == nil { - self.expiry[envelope.Expiry] = set.NewNonTS() + if self.expirations[envelope.Expiry] == nil { + self.expirations[envelope.Expiry] = set.NewNonTS() } - if !self.expiry[envelope.Expiry].Has(hash) { - self.expiry[envelope.Expiry].Add(hash) + if !self.expirations[envelope.Expiry].Has(hash) { + self.expirations[envelope.Expiry].Add(hash) // Notify the local node of a message arrival go self.postEvent(envelope) @@ -292,11 +306,11 @@ func (self *Whisper) update() { // expire iterates over all the expiration timestamps, removing all stale // messages from the pools. func (self *Whisper) expire() { - self.mmu.Lock() - defer self.mmu.Unlock() + self.poolMu.Lock() + defer self.poolMu.Unlock() now := uint32(time.Now().Unix()) - for then, hashSet := range self.expiry { + for then, hashSet := range self.expirations { // Short circuit if a future time if then > now { continue @@ -306,14 +320,14 @@ func (self *Whisper) expire() { delete(self.messages, v.(common.Hash)) return true }) - self.expiry[then].Clear() + self.expirations[then].Clear() } } // envelopes retrieves all the messages currently pooled by the node. func (self *Whisper) envelopes() []*Envelope { - self.mmu.RLock() - defer self.mmu.RUnlock() + self.poolMu.RLock() + defer self.poolMu.RUnlock() envelopes := make([]*Envelope, 0, len(self.messages)) for _, envelope := range self.messages { |