lnwire: cap decoded short channel IDs

In this commit, we cap each decompressed short channel ID set at 100,000
entries, matching the aggregate range reply budget. The old zlib reader
bounded compressed input rather than decoded output, so the two working-set
limits could drift apart.

We retain compatibility with protocol-valid compressed replies, reject
truncated or corrupt zlib streams, and close the reader on every exit.
Boundary, compatibility, corruption, and property tests cover the
decoder.

(cherry picked from commit d162291941)

lnwire/query_short_chan_ids.go

@@ -3,6 +3,7 @@ package lnwire
 import (
 	"bytes"
 	"compress/zlib"
+	"errors"
 	"fmt"
 	"io"
 	"sort"
@@ -12,10 +13,10 @@ import (
 )
 
 const (
-	// maxZlibBufSize is the max number of bytes that we'll accept from a
-	// zlib decoding instance. We do this in order to limit the total
-	// amount of memory allocated during a decoding instance.
-	maxZlibBufSize = 67413630
+	// maxDecodedShortChanIDs is the maximum number of short channel IDs
+	// accepted from a single message. The plain encoding is also bounded
+	// by the wire size, so its check is defense in depth.
+	maxDecodedShortChanIDs = 100_000
 )
 
 // ErrUnsortedSIDs is returned when decoding a QueryShortChannelID request whose
@@ -164,6 +165,12 @@ func decodeShortChanIDs(r io.Reader) (QueryEncoding, []ShortChannelID, error) {
 		// compute the number of bytes encoded based on the size of the
 		// query body.
 		numShortChanIDs := len(queryBody) / 8
+		if numShortChanIDs > maxDecodedShortChanIDs {
+			return 0, nil, fmt.Errorf(
+				"too many short channel IDs: max=%v, got=%v",
+				maxDecodedShortChanIDs, numShortChanIDs,
+			)
+		}
 		if numShortChanIDs == 0 {
 			return encodingType, nil, nil
 		}
@@ -210,61 +217,28 @@ func decodeShortChanIDs(r io.Reader) (QueryEncoding, []ShortChannelID, error) {
 			return encodingType, nil, nil
 		}
 
-		// Before we start to decode, we'll create a limit reader over
-		// the current reader. This will ensure that we can control how
-		// much memory we're allocating during the decoding process.
-		limitedDecompressor, err := zlib.NewReader(&io.LimitedReader{
-			R: bytes.NewReader(queryBody),
-			N: maxZlibBufSize,
-		})
+		decompressor, err := zlib.NewReader(bytes.NewReader(queryBody))
 		if err != nil {
 			return 0, nil, fmt.Errorf("unable to create zlib "+
 				"reader: %w", err)
 		}
 
-		var (
-			shortChanIDs []ShortChannelID
-			lastChanID   ShortChannelID
-			i            int
+		shortChanIDs, decodeErr := decodeCompressedShortChanIDs(
+			decompressor,
 		)
-		for {
-			// We'll now attempt to read the next short channel ID
-			// encoded in the payload.
-			var cid ShortChannelID
-			err := ReadElements(limitedDecompressor, &cid)
+		closeErr := decompressor.Close()
 
-			switch {
-			// If we get an EOF error, then that either means we've
-			// read all that's contained in the buffer, or have hit
-			// our limit on the number of bytes we'll read. In
-			// either case, we'll return what we have so far.
-			case err == io.ErrUnexpectedEOF || err == io.EOF:
-				return encodingType, shortChanIDs, nil
+		switch {
+		case decodeErr != nil:
+			return 0, nil, decodeErr
 
-			// Otherwise, we hit some other sort of error, possibly
-			// an invalid payload, so we'll exit early with the
-			// error.
-			case err != nil:
-				return 0, nil, fmt.Errorf("unable to "+
-					"deflate next short chan "+
-					"ID: %v", err)
-			}
+		case closeErr != nil:
+			return 0, nil, fmt.Errorf(
+				"unable to close zlib reader: %w", closeErr,
+			)
 
-			// We successfully read the next ID, so we'll collect
-			// that in the set of final ID's to return.
-			shortChanIDs = append(shortChanIDs, cid)
-
-			// Finally, we'll ensure that this short chan ID is
-			// greater than the last one. This is a requirement
-			// within the encoding, and if violated can aide us in
-			// detecting malicious payloads. This can only be true
-			// starting at the second chanID.
-			if i > 0 && cid.ToUint64() <= lastChanID.ToUint64() {
-				return 0, nil, ErrUnsortedSIDs{lastChanID, cid}
-			}
-
-			lastChanID = cid
-			i++
+		default:
+			return encodingType, shortChanIDs, nil
 		}
 
 	default:
@@ -275,6 +249,45 @@ func decodeShortChanIDs(r io.Reader) (QueryEncoding, []ShortChannelID, error) {
 	}
 }
 
+// decodeCompressedShortChanIDs decodes and validates the decompressed short
+// channel ID stream.
+func decodeCompressedShortChanIDs(r io.Reader) ([]ShortChannelID, error) {
+	var (
+		shortChanIDs []ShortChannelID
+		lastChanID   ShortChannelID
+	)
+
+	for {
+		var cid ShortChannelID
+		err := ReadElements(r, &cid)
+
+		switch {
+		// Only a clean EOF terminates the stream. A partial final ID
+		// returns io.ErrUnexpectedEOF and remains an error.
+		case errors.Is(err, io.EOF):
+			return shortChanIDs, nil
+
+		case err != nil:
+			return nil, fmt.Errorf("unable to deflate next short "+
+				"chan ID: %w", err)
+		}
+
+		if len(shortChanIDs) == maxDecodedShortChanIDs {
+			return nil, fmt.Errorf("too many short channel IDs: "+
+				"max=%v", maxDecodedShortChanIDs)
+		}
+
+		if len(shortChanIDs) > 0 &&
+			cid.ToUint64() <= lastChanID.ToUint64() {
+
+			return nil, ErrUnsortedSIDs{lastChanID, cid}
+		}
+
+		shortChanIDs = append(shortChanIDs, cid)
+		lastChanID = cid
+	}
+}
+
 // Encode serializes the target QueryShortChanIDs into the passed io.Writer
 // observing the protocol version specified.
 //

lnwire/query_short_chan_ids_test.go

@@ -3,6 +3,9 @@ package lnwire
 import (
 	"bytes"
 	"testing"
+
+	"github.com/stretchr/testify/require"
+	"pgregory.net/rapid"
 )
 
 type unsortedSidTest struct {
@@ -118,3 +121,208 @@ func TestQueryShortChanIDsZero(t *testing.T) {
 		})
 	}
 }
+
+// TestQueryShortChanIDsRoundTrip uses property-based testing to ensure both
+// supported encodings preserve sorted short channel ID sets.
+func TestQueryShortChanIDsRoundTrip(t *testing.T) {
+	t.Parallel()
+
+	rapid.Check(t, func(t *rapid.T) {
+		encoding := rapid.SampledFrom([]QueryEncoding{
+			EncodingSortedPlain,
+			EncodingSortedZlib,
+		}).Draw(t, "encoding")
+
+		numSCIDs := rapid.IntRange(0, 512).Draw(t, "num-scids")
+		var scids []ShortChannelID
+		if numSCIDs > 0 {
+			scids = make([]ShortChannelID, numSCIDs)
+		}
+
+		offset := rapid.IntRange(0, 1_000_000).Draw(t, "offset")
+		step := rapid.IntRange(1, 1_000_000).Draw(t, "step")
+		for i := range scids {
+			scid := uint64(offset + i*step)
+			scids[i] = NewShortChanIDFromInt(scid)
+		}
+
+		var b bytes.Buffer
+		require.NoError(t, encodeShortChanIDs(
+			&b, encoding, scids,
+		))
+
+		decodedEncoding, decoded, err := decodeShortChanIDs(
+			bytes.NewReader(b.Bytes()),
+		)
+		require.NoError(t, err)
+		require.Equal(t, encoding, decodedEncoding)
+		require.Equal(t, scids, decoded)
+	})
+}
+
+// TestQueryShortChanIDsDecodeLimit ensures that a decompressed short channel
+// ID stream cannot exceed its resource limit.
+func TestQueryShortChanIDsDecodeLimit(t *testing.T) {
+	t.Parallel()
+
+	var stream bytes.Buffer
+	for i := 0; i <= maxDecodedShortChanIDs; i++ {
+		require.NoError(t, WriteElements(
+			&stream, NewShortChanIDFromInt(uint64(i)),
+		))
+	}
+
+	decoded, err := decodeCompressedShortChanIDs(bytes.NewReader(
+		stream.Bytes()[:maxDecodedShortChanIDs*8],
+	))
+	require.NoError(t, err)
+	require.Len(t, decoded, maxDecodedShortChanIDs)
+
+	_, err = decodeCompressedShortChanIDs(
+		bytes.NewReader(stream.Bytes()),
+	)
+	require.ErrorContains(t, err, "too many short channel IDs")
+}
+
+// TestQueryShortChanIDsZlibCompatibility ensures that a protocol-valid
+// compressed reply can contain far more short channel IDs than a plain reply.
+// The plain encoding is bounded by the wire size at maxPlainReplySCIDs, so it
+// is the compressed encoding that determines how much headroom a single reply
+// actually has.
+func TestQueryShortChanIDsZlibCompatibility(t *testing.T) {
+	t.Parallel()
+
+	const (
+		// maxWireMsgSize is the largest a message may be on the wire,
+		// including its type prefix.
+		maxWireMsgSize = MaxMsgBody + MessageTypeSize
+
+		// maxPlainReplySCIDs is the number of SCIDs that saturate a
+		// ReplyChannelRange under the plain encoding. The message
+		// carries 41 bytes of fixed fields, and the SCID blob adds a
+		// 2-byte length prefix plus a 1-byte encoding type, leaving
+		// (65533 - 44) / 8 SCIDs.
+		maxPlainReplySCIDs = 8186
+
+		// maxZlibReplySCIDs is the number of consecutive SCIDs that
+		// saturate the same message under the zlib encoding. Runs of
+		// consecutive SCIDs are the best case for the compressor, so
+		// this is an upper bound rather than a figure real peers hit.
+		maxZlibReplySCIDs = 30_794
+	)
+
+	// A reply full of consecutive SCIDs is what we'll size both encodings
+	// against.
+	newReply := func(enc QueryEncoding, n int) *ReplyChannelRange {
+		scids := make([]ShortChannelID, n)
+		for i := range scids {
+			scids[i] = NewShortChanIDFromInt(uint64(i))
+		}
+
+		return &ReplyChannelRange{
+			Complete:     1,
+			EncodingType: enc,
+			ShortChanIDs: scids,
+			ExtraData:    make([]byte, 0),
+		}
+	}
+
+	// The plain encoding tops out at maxPlainReplySCIDs: that many SCIDs
+	// fit, and one more overflows the message.
+	plain := newReply(EncodingSortedPlain, maxPlainReplySCIDs)
+	size, err := plain.SerializedSize()
+	require.NoError(t, err)
+	require.LessOrEqual(t, size, uint32(maxWireMsgSize))
+
+	plain = newReply(EncodingSortedPlain, maxPlainReplySCIDs+1)
+	size, err = plain.SerializedSize()
+	require.NoError(t, err)
+	require.Greater(t, size, uint32(maxWireMsgSize))
+
+	// The zlib encoding fits far more SCIDs into the very same message,
+	// which is the compatibility property we care about: a compressed
+	// reply can carry a much larger slice of the graph than a plain one.
+	zlib := newReply(EncodingSortedZlib, maxZlibReplySCIDs)
+	size, err = zlib.SerializedSize()
+	require.NoError(t, err)
+	require.LessOrEqual(t, size, uint32(maxWireMsgSize))
+	require.Greater(t, maxZlibReplySCIDs, maxPlainReplySCIDs)
+
+	// One more SCID pushes the compressed reply over the wire limit, so
+	// maxZlibReplySCIDs really is the ceiling.
+	over := newReply(EncodingSortedZlib, maxZlibReplySCIDs+1)
+	size, err = over.SerializedSize()
+	require.NoError(t, err)
+	require.Greater(t, size, uint32(maxWireMsgSize))
+
+	// Finally, the saturated compressed reply must still round trip
+	// cleanly through the decoder.
+	var b bytes.Buffer
+	require.NoError(t, encodeShortChanIDs(
+		&b, EncodingSortedZlib, zlib.ShortChanIDs,
+	))
+
+	encoding, decoded, err := decodeShortChanIDs(
+		bytes.NewReader(b.Bytes()),
+	)
+	require.NoError(t, err)
+	require.Equal(t, EncodingSortedZlib, encoding)
+	require.Equal(t, zlib.ShortChanIDs, decoded)
+}
+
+// TestQueryShortChanIDsRejectsCorruptZlib ensures that truncated or corrupt
+// compressed streams are not accepted as valid partial replies.
+func TestQueryShortChanIDsRejectsCorruptZlib(t *testing.T) {
+	t.Parallel()
+
+	scids := []ShortChannelID{
+		NewShortChanIDFromInt(1),
+		NewShortChanIDFromInt(2),
+		NewShortChanIDFromInt(3),
+	}
+
+	var encoded bytes.Buffer
+	require.NoError(t, encodeShortChanIDs(
+		&encoded, EncodingSortedZlib, scids,
+	))
+
+	body := encoded.Bytes()[2:]
+	corruptChecksum := append([]byte(nil), body...)
+	corruptChecksum[len(corruptChecksum)-1] ^= 1
+
+	tests := []struct {
+		name string
+		body []byte
+	}{
+		{
+			name: "truncated header",
+			body: body[:2],
+		},
+		{
+			name: "truncated checksum",
+			body: body[:len(body)-1],
+		},
+		{
+			name: "corrupt checksum",
+			body: corruptChecksum,
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			t.Parallel()
+
+			var message bytes.Buffer
+			require.NoError(t, WriteElements(
+				&message, uint16(len(test.body)),
+			))
+			_, err := message.Write(test.body)
+			require.NoError(t, err)
+
+			_, _, err = decodeShortChanIDs(
+				bytes.NewReader(message.Bytes()),
+			)
+			require.Error(t, err)
+		})
+	}
+}