diff --git a/swarm/pss/pss.go b/swarm/pss/pss.go index b788fd050a..d2eeff5189 100644 --- a/swarm/pss/pss.go +++ b/swarm/pss/pss.go @@ -33,7 +33,7 @@ const ( defaultMaxMsgSize = 1024 * 1024 defaultCleanInterval = time.Second * 60 * 10 defaultDequeueInterval = time.Millisecond * 10 - defaultOutboxQueueSize = 10000 + defaultOutboxCapacity = 10000 pssProtocolName = "pss" pssVersion = 1 hasherCount = 8 @@ -104,6 +104,8 @@ type Pss struct { paddingByteSize int capstring string outbox chan *PssMsg + outboxCapacity int + outboxCounter int // keys and peers pubKeyPool map[string]map[Topic]*pssPeer // mapping of hex public keys to peer address by topic. @@ -149,7 +151,8 @@ func NewPss(k network.Overlay, params *PssParams) *Pss { msgTTL: params.MsgTTL, paddingByteSize: defaultPaddingByteSize, capstring: cap.String(), - outbox: make(chan *PssMsg, defaultOutboxQueueSize), + outbox: make(chan *PssMsg, defaultOutboxCapacity), + outboxCapacity: defaultOutboxCapacity, pubKeyPool: make(map[string]map[Topic]*pssPeer), symKeyPool: make(map[string]map[Topic]*pssPeer), @@ -328,12 +331,16 @@ func (self *Pss) handlePssMsg(msg interface{}) error { var err error if !self.isSelfPossibleRecipient(pssmsg) { log.Trace("pss was for someone else :'( ... forwarding", "pss", common.ToHex(self.BaseAddr())) - self.outbox <- pssmsg + if !self.enqueue(pssmsg) { + return errors.New("outbox full!") + } } log.Trace("pss for us, yay! ... let's process!", "pss", common.ToHex(self.BaseAddr())) if !self.process(pssmsg) { - self.outbox <- pssmsg + if !self.enqueue(pssmsg) { + return errors.New("outbox full!") + } } return err } @@ -370,9 +377,10 @@ func (self *Pss) process(pssmsg *PssMsg) bool { } if len(pssmsg.To) < addressLength { - go func() { - self.outbox <- pssmsg - }() + if !self.enqueue(pssmsg) { + log.Error("outbox full!") + return false + } } if psstopic == rawTopic { return false @@ -600,6 +608,17 @@ func (self *Pss) cleanKeys() (count int) { // SECTION: Message sending ///////////////////////////////////////////////////////////////////// +func (self *Pss) enqueue(msg *PssMsg) bool { + self.fwdPoolMu.Lock() + defer self.fwdPoolMu.Unlock() + if self.outboxCounter == self.outboxCapacity { + return false + } + self.outboxCounter++ + self.outbox <- msg + return true +} + // Send a raw message (any encryption is responsibility of calling client) // // Will fail if raw messages are disallowed @@ -616,7 +635,9 @@ func (self *Pss) SendRaw(msg []byte, address PssAddress) error { }, } self.addFwdCache(pssmsg) - self.outbox <- pssmsg + if !self.enqueue(pssmsg) { + return errors.New("outbox full!") + } return nil } @@ -710,14 +731,16 @@ func (self *Pss) send(to []byte, topic Topic, msg []byte, asymmetric bool, key [ Expire: uint32(time.Now().Add(self.msgTTL).Unix()), Payload: envelope, } - self.outbox <- pssmsg + if !self.enqueue(pssmsg) { + return errors.New("outbox full!") + } return nil } // Forwards a pss message to the peer(s) closest to the to recipient address in the PssMsg struct // The recipient address can be of any length, and the byte slice will be matched to the MSB slice // of the peer address of the equivalent length. -func (self *Pss) forward(msg *PssMsg) { +func (self *Pss) forward(msg *PssMsg) error { to := make([]byte, addressLength) copy(to[:len(msg.To)], msg.To) @@ -782,11 +805,19 @@ func (self *Pss) forward(msg *PssMsg) { if sent == 0 { log.Debug("unable to forward to any peers") time.Sleep(time.Millisecond) - self.outbox <- msg + if !self.enqueue(msg) { + return errors.New("outbox full!") + } } + // remove from queue + self.fwdPoolMu.Lock() + self.outboxCounter-- + self.fwdPoolMu.Unlock() + // cache the message self.addFwdCache(msg) + return nil } ///////////////////////////////////////////////////////////////////// diff --git a/swarm/pss/pss_test.go b/swarm/pss/pss_test.go index 6d06c9f264..9970881bd6 100644 --- a/swarm/pss/pss_test.go +++ b/swarm/pss/pss_test.go @@ -419,13 +419,13 @@ func TestMismatch(t *testing.T) { } -func TestRawSend(t *testing.T) { - t.Run("32", testRawSend) - t.Run("8", testRawSend) - t.Run("0", testRawSend) +func TestSendRaw(t *testing.T) { + t.Run("32", testSendRaw) + t.Run("8", testSendRaw) + t.Run("0", testSendRaw) } -func testRawSend(t *testing.T) { +func testSendRaw(t *testing.T) { var addrsize int64 var err error @@ -502,13 +502,13 @@ func testRawSend(t *testing.T) { } // send symmetrically encrypted message between two directly connected peers -func TestSymSend(t *testing.T) { - t.Run("32", testSymSend) - t.Run("8", testSymSend) - t.Run("0", testSymSend) +func TestSendSym(t *testing.T) { + t.Run("32", testSendSym) + t.Run("8", testSendSym) + t.Run("0", testSendSym) } -func testSymSend(t *testing.T) { +func testSendSym(t *testing.T) { // address hint size var addrsize int64 @@ -617,13 +617,13 @@ func testSymSend(t *testing.T) { } // send asymmetrically encrypted message between two directly connected peers -func TestAsymSend(t *testing.T) { - t.Run("32", testAsymSend) - t.Run("8", testAsymSend) - t.Run("0", testAsymSend) +func TestSendAsym(t *testing.T) { + t.Run("32", testSendAsym) + t.Run("8", testSendAsym) + t.Run("0", testSendAsym) } -func testAsymSend(t *testing.T) { +func testSendAsym(t *testing.T) { // address hint size var addrsize int64