lnwire: cap decoded short channel IDs

diff --git a/lnwire/query_short_chan_ids.go b/lnwire/query_short_chan_ids.go
index 0460b2e..3e98a2b 100644
--- a/lnwire/query_short_chan_ids.go
+++ b/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)
-
-			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
-
-			// 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)
-			}
+		closeErr := decompressor.Close()
 
-			// We successfully read the next ID, so we'll collect
-			// that in the set of final ID's to return.
-			shortChanIDs = append(shortChanIDs, cid)
+		switch {
+		case decodeErr != nil:
+			return 0, nil, decodeErr
 
-			// 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}
-			}
+		case closeErr != nil:
+			return 0, nil, fmt.Errorf(
+				"unable to close zlib reader: %w", closeErr,
+			)
 
-			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.
 //
diff --git a/lnwire/query_short_chan_ids_test.go b/lnwire/query_short_chan_ids_test.go
index a147061..30ecbbe 100644
--- a/lnwire/query_short_chan_ids_test.go
+++ b/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 {
@@ -116,3 +119,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)
+		})
+	}
+}