peer: Limit peer queued output bytes.

peer/peer.go

@@ -36,6 +36,12 @@ const (
 	// outputBufferSize is the number of elements the output channels use.
 	outputBufferSize = 5000
 
+	// maxQueuedOutputBytes is the maximum number of bytes queued for delivery
+	// to a single peer. A peer whose queued bytes exceed this will be
+	// disconnected. The value is chosen to comfortably exceed the largest
+	// single message.
+	maxQueuedOutputBytes = 40 * 1024 * 1024 // 40 MiB
+
 	// maxInvTrickleSize is the maximum amount of inventory to send in a single
 	// message when trickling inventory to remote peers.
 	maxInvTrickleSize = 1000
@@ -370,6 +376,18 @@ func newNetAddress(addr net.Addr, services wire.ServiceFlag) (*wire.NetAddressV2
 type outMsg struct {
 	msg      wire.Message
 	doneChan chan<- struct{}
+
+	// size is the serialized size of the message as calculated by [msgSize] at
+	// the time the message is queued.  It is stored so the exact same value is
+	// used both when adding to and deducting from the total bytes queued for
+	// the peer.
+	size int64
+}
+
+// msgSize returns the total serialized size of the passed message including the
+// message header.
+func msgSize(msg wire.Message) int64 {
+	return int64(wire.MessageHeaderSize + msg.SerializeSize())
 }
 
 // stallControlCmd represents the command of a stall control message.
@@ -468,6 +486,7 @@ type Peer struct {
 	// The following variables must only be used atomically.
 	bytesReceived uint64
 	bytesSent     uint64
+	bytesQueued   int64 // total serialized bytes currently queued for send
 	lastRecv      int64
 	lastSend      int64
 	disconnect    int32
@@ -982,8 +1001,10 @@ func (p *Peer) readMessage() (wire.Message, []byte, error) {
 	return msg, buf, nil
 }
 
-// writeMessage sends a wire message to the peer with logging.
-func (p *Peer) writeMessage(msg wire.Message) error {
+// writeMessage sends a wire message to the peer with logging.  The size must be
+// the total serialized size of the message, including the message header, as
+// returned by [msgSize].
+func (p *Peer) writeMessage(msg wire.Message, size int64) error {
 	// Don't do anything if we're disconnecting.
 	if atomic.LoadInt32(&p.disconnect) != 0 {
 		return nil
@@ -1011,8 +1032,7 @@ func (p *Peer) writeMessage(msg wire.Message) error {
 	// Limit how long the write is allowed to block.  A base timeout plus an
 	// allowance proportional to the message size affords larger messages
 	// proportionally more time.
-	msgSize := wire.MessageHeaderSize + msg.SerializeSize()
-	allowance := time.Duration(msgSize/writeStallBytesPerSec) * time.Second
+	allowance := time.Duration(size/writeStallBytesPerSec) * time.Second
 	deadline := p.nowFn().Add(writeStallTimeout + allowance)
 	err := p.conn.SetWriteDeadline(deadline)
 	if err != nil {
@@ -1663,9 +1683,11 @@ out:
 
 				invMsg.AddInvVect(iv)
 				if len(invMsg.InvList) >= maxInvTrickleSize {
-					waiting = queuePacket(
-						outMsg{msg: invMsg},
-						&pendingMsgs, waiting)
+					om, ok := p.queueOutMsg(invMsg, nil)
+					if !ok {
+						break out
+					}
+					waiting = queuePacket(om, &pendingMsgs, waiting)
 					invMsg = wire.NewMsgInvSizeHint(uint(len(invSendQueue)))
 				}
 
@@ -1674,8 +1696,11 @@ out:
 				p.AddKnownInventory(iv)
 			}
 			if len(invMsg.InvList) > 0 {
-				waiting = queuePacket(outMsg{msg: invMsg},
-					&pendingMsgs, waiting)
+				om, ok := p.queueOutMsg(invMsg, nil)
+				if !ok {
+					break out
+				}
+				waiting = queuePacket(om, &pendingMsgs, waiting)
 			}
 			invSendQueue = nil
 
@@ -1758,7 +1783,10 @@ out:
 			case <-p.quit:
 				break out
 			}
-			if err := p.writeMessage(msg.msg); err != nil {
+
+			err := p.writeMessage(msg.msg, msg.size)
+			atomic.AddInt64(&p.bytesQueued, -msg.size)
+			if err != nil {
 				p.Disconnect()
 				if p.shouldLogWriteError(err) {
 					log.Errorf("Failed to send message to "+
@@ -1812,6 +1840,24 @@ cleanup:
 	log.Tracef("Peer output handler done for %s", p)
 }
 
+// queueOutMsg adds the passed wire message to the delivery queue for the peer.
+// If adding this message would exceed the limit of queued bytes, the peer is instead
+// disconnected and false is returned.
+//
+// This function is safe for concurrent access.
+func (p *Peer) queueOutMsg(msg wire.Message, doneChan chan<- struct{}) (outMsg, bool) {
+	size := msgSize(msg)
+	queued := atomic.AddInt64(&p.bytesQueued, size)
+	if queued > maxQueuedOutputBytes {
+		log.Debugf("%s exceeded max allowed queued output bytes (queued %d, "+
+			"max %d) -- disconnecting", p, queued, maxQueuedOutputBytes)
+		p.Disconnect()
+		return outMsg{}, false
+	}
+
+	return outMsg{msg: msg, doneChan: doneChan, size: size}, true
+}
+
 // QueueMessage adds the passed wire message to the peer send queue.
 //
 // This function is safe for concurrent access.
@@ -1844,7 +1890,18 @@ func (p *Peer) QueueMessage(msg wire.Message, doneChan chan<- struct{}) {
 		}
 		return
 	}
-	p.outputQueue <- outMsg{msg: msg, doneChan: doneChan}
+
+	om, ok := p.queueOutMsg(msg, doneChan)
+	if !ok {
+		if doneChan != nil {
+			go func() {
+				doneChan <- struct{}{}
+			}()
+		}
+		return
+	}
+
+	p.outputQueue <- om
 }
 
 // QueueInventory adds the passed inventory to the inventory send queue which
@@ -1893,7 +1950,12 @@ func (p *Peer) QueueInventoryImmediate(invVect *wire.InvVect) {
 	invMsg := wire.NewMsgInvSizeHint(1)
 	invMsg.AddInvVect(invVect)
 	p.AddKnownInventory(invVect)
-	p.outputQueue <- outMsg{msg: invMsg, doneChan: nil}
+	om, ok := p.queueOutMsg(invMsg, nil)
+	if !ok {
+		return
+	}
+
+	p.outputQueue <- om
 }
 
 // Connected returns whether or not the peer is currently connected.
@@ -2148,7 +2210,7 @@ func (p *Peer) writeLocalVersionMsg() error {
 		return err
 	}
 
-	if err := p.writeMessage(localVerMsg); err != nil {
+	if err := p.writeMessage(localVerMsg, msgSize(localVerMsg)); err != nil {
 		return err
 	}
 
@@ -2173,7 +2235,8 @@ func (p *Peer) inboundHandshake(onVersion OnVersionCallback) error {
 		return err
 	}
 
-	if err := p.writeMessage(wire.NewMsgVerAck()); err != nil {
+	verAckMsg := wire.NewMsgVerAck()
+	if err := p.writeMessage(verAckMsg, msgSize(verAckMsg)); err != nil {
 		return err
 	}
 
@@ -2214,7 +2277,8 @@ func (p *Peer) outboundHandshake(onVersion OnVersionCallback) error {
 		return err
 	}
 
-	if err := p.writeMessage(wire.NewMsgVerAck()); err != nil {
+	verAckMsg := wire.NewMsgVerAck()
+	if err := p.writeMessage(verAckMsg, msgSize(verAckMsg)); err != nil {
 		return err
 	}
 

peer/peer_test.go

@@ -15,6 +15,7 @@ import (
 	"reflect"
 	"strconv"
 	"sync"
+	"sync/atomic"
 	"testing"
 	"time"
 
@@ -981,7 +982,8 @@ func TestWriteMessageDeadline(t *testing.T) {
 	p.nowFn = mockNowFn
 
 	for _, test := range tests {
-		if err := p.writeMessage(&sizedMsg{size: test.payloadSize}); err != nil {
+		msg := &sizedMsg{size: test.payloadSize}
+		if err := p.writeMessage(msg, msgSize(msg)); err != nil {
 			t.Fatalf("%s: unexpected error writing message: %v", test.name, err)
 		}
 
@@ -1005,12 +1007,83 @@ func TestWriteMessageDeadlineError(t *testing.T) {
 	wantErr := errors.New("set write deadline failed")
 	inConn.setWriteDeadlineErr = wantErr
 
-	got := p.writeMessage(&sizedMsg{})
+	got := p.writeMessage(&sizedMsg{}, 0)
 	if !errors.Is(got, wantErr) {
 		t.Fatalf("wrong error - got %v, want %v", got, wantErr)
 	}
 }
 
+// TestQueuedOutputLimit ensures a peer is disconnected when the total bytes
+// queued for it exceeds the allowed limit.
+func TestQueuedOutputLimit(t *testing.T) {
+	// Create a peer with a mock connection.  Note the peer is intentionally not
+	// run so nothing drains the output queue and thus the queued bytes
+	// accumulate.
+	inConn, outConn := pipe("10.0.0.1:9108", "10.0.0.2:9108")
+	defer inConn.Close()
+	defer outConn.Close()
+	p := NewInboundPeer(mockPeerConfig(), inConn)
+
+	// Queue messages that total exactly the limit and ensure the peer remains
+	// connected.
+	const numMsgs = 4
+	const msgSize = maxQueuedOutputBytes/numMsgs - wire.MessageHeaderSize
+	for i := 0; i < numMsgs; i++ {
+		p.QueueMessage(&sizedMsg{size: msgSize}, nil)
+	}
+	if got := atomic.LoadInt64(&p.bytesQueued); got != maxQueuedOutputBytes {
+		t.Fatalf("wrong bytes queued - got %d, want %d", got,
+			maxQueuedOutputBytes)
+	}
+	if !p.Connected() {
+		t.Fatal("peer disconnected without exceeding queued output limit")
+	}
+
+	// Queue one more message to exceed the limit and ensure the peer is
+	// disconnected and the done channel is signalled.
+	done := make(chan struct{}, 1)
+	p.QueueMessage(&sizedMsg{size: 1}, done)
+	select {
+	case <-done:
+	case <-time.After(time.Second):
+		t.Fatal("done channel not signalled")
+	}
+	if p.Connected() {
+		t.Fatal("peer not disconnected after exceeding queued output limit")
+	}
+}
+
+// TestQueuedOutputBytesDeducted ensures the bytes queued for a peer are
+// deducted once the associated message has been sent.
+func TestQueuedOutputBytesDeducted(t *testing.T) {
+	// Create a pair of peers that are connected to each other using a fake
+	// connection.
+	peerCfg := mockPeerConfig()
+	inConn, outConn := pipe("10.0.0.1:9108", "10.0.0.2:9108")
+	outPeer := NewOutboundPeer(peerCfg, outConn.RemoteAddr(), outConn)
+	inPeer := NewInboundPeer(peerCfg, inConn)
+	if err := runHandshakes(inPeer, outPeer); err != nil {
+		t.Fatalf("failed to perform handshake: %v", err)
+	}
+	cancel, wg := runPeersAsync(inPeer, outPeer)
+	defer wg.Wait()
+	defer cancel()
+
+	// Queue a message and wait until it has been sent.
+	done := make(chan struct{})
+	outPeer.QueueMessage(wire.NewMsgPing(0), done)
+	select {
+	case <-done:
+	case <-time.After(time.Second):
+		t.Fatal("done channel not signalled")
+	}
+
+	// Ensure the bytes queued for the message were deducted once it was sent.
+	if got := atomic.LoadInt64(&outPeer.bytesQueued); got != 0 {
+		t.Fatalf("wrong bytes queued - got %d, want 0", got)
+	}
+}
+
 func init() {
 	// Allow self connection when running the tests.
 	allowSelfConns = true