From 4b12f40a48e580f3b34b7f43eb600b3aaed787b2 Mon Sep 17 00:00:00 2001
From: yyforyongyu <yy2452@columbia.edu>
Date: Thu, 17 Jun 2021 10:49:33 +0800
Subject: [PATCH 01/11] lnwire: remove unused MaxMessagePayload const

MaxMessagePayload and MaxSliceLength are duplicate variables. This
commit deletes MaxMessagePayload and keeps MaxSliceLength.
---
 buffer/read.go    | 2 +-
 buffer/write.go   | 2 +-
 lnwire/message.go | 4 ----
 3 files changed, 2 insertions(+), 6 deletions(-)

diff --git a/buffer/read.go b/buffer/read.go
index 57050d58f..f38ab971e 100644
--- a/buffer/read.go
+++ b/buffer/read.go
@@ -7,7 +7,7 @@ import (
 // ReadSize represents the size of the maximum message that can be read off the
 // wire by brontide. The buffer is used to hold the ciphertext while the
 // brontide state machine decrypts the message.
-const ReadSize = lnwire.MaxMessagePayload + 16
+const ReadSize = lnwire.MaxSliceLength + 16
 
 // Read is a static byte array sized to the maximum-allowed Lightning message
 // size, plus 16 bytes for the MAC.
diff --git a/buffer/write.go b/buffer/write.go
index 47010c29b..712571b97 100644
--- a/buffer/write.go
+++ b/buffer/write.go
@@ -7,7 +7,7 @@ import (
 // WriteSize represents the size of the maximum plaintext message than can be
 // sent using brontide. The buffer does not include extra space for the MAC, as
 // that is applied by the Noise protocol after encrypting the plaintext.
-const WriteSize = lnwire.MaxMessagePayload
+const WriteSize = lnwire.MaxSliceLength
 
 // Write is static byte array occupying to maximum-allowed plaintext-message
 // size.
diff --git a/lnwire/message.go b/lnwire/message.go
index 048122887..c4f2fd36a 100644
--- a/lnwire/message.go
+++ b/lnwire/message.go
@@ -12,10 +12,6 @@ import (
 	"io"
 )
 
-// MaxMessagePayload is the maximum bytes a message can be regardless of other
-// individual limits imposed by messages themselves.
-const MaxMessagePayload = 65535 // 65KB
-
 // MessageType is the unique 2 byte big-endian integer that indicates the type
 // of message on the wire. All messages have a very simple header which
 // consists simply of 2-byte message type. We omit a length field, and checksum

From 77862e45af7aa257c114fae29656a419ebd96474 Mon Sep 17 00:00:00 2001
From: yyforyongyu <yy2452@columbia.edu>
Date: Fri, 11 Dec 2020 14:36:07 +0800
Subject: [PATCH 02/11] lnwire: add benchmark tests to ReadMessage and
 WriteMessage

---
 lnwire/message_test.go | 867 +++++++++++++++++++++++++++++++++++++++++
 1 file changed, 867 insertions(+)
 create mode 100644 lnwire/message_test.go

diff --git a/lnwire/message_test.go b/lnwire/message_test.go
new file mode 100644
index 000000000..fbade5082
--- /dev/null
+++ b/lnwire/message_test.go
@@ -0,0 +1,867 @@
+package lnwire_test
+
+import (
+	"bytes"
+	"encoding/binary"
+	"image/color"
+	"io"
+	"math"
+	"math/big"
+	"math/rand"
+	"net"
+	"sync"
+	"testing"
+
+	"github.com/btcsuite/btcd/btcec"
+	"github.com/btcsuite/btcutil"
+	"github.com/lightningnetwork/lnd/lnwire"
+	"github.com/lightningnetwork/lnd/tor"
+	"github.com/stretchr/testify/require"
+)
+
+const deliveryAddressMaxSize = 34
+const letterBytes = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ"
+
+var (
+	testSig = &btcec.Signature{
+		R: new(big.Int),
+		S: new(big.Int),
+	}
+	testNodeSig, _ = lnwire.NewSigFromSignature(testSig)
+
+	testNumExtraBytes = 1000
+	testNumSigs       = 100
+	testNumChanIDs    = 1000
+	buffer            = make([]byte, 0, lnwire.MaxSliceLength)
+
+	bufPool = sync.Pool{
+		New: func() interface{} {
+			return bytes.NewBuffer(buffer)
+		},
+	}
+)
+
+// BenchmarkWriteMessage benchmarks the performance of lnwire.WriteMessage. It
+// generates a test message for each of the lnwire.Message, calls the
+// WriteMessage method and benchmark it.
+func BenchmarkWriteMessage(b *testing.B) {
+	// Create testing messages. We will use a constant seed to make sure
+	// the benchmark uses the same data every time.
+	r := rand.New(rand.NewSource(42))
+
+	msgAll := makeAllMessages(b, r)
+
+	// Iterate all messages and write each once.
+	for _, msg := range msgAll {
+		m := msg
+		// Run each message as a sub benchmark test.
+		b.Run(msg.MsgType().String(), func(b *testing.B) {
+
+			b.ResetTimer()
+			for i := 0; i < b.N; i++ {
+				// Fetch a buffer from the pool and reset it.
+				buf := bufPool.Get().(*bytes.Buffer)
+				buf.Reset()
+
+				_, err := lnwire.WriteMessage(buf, m, 0)
+				require.NoError(b, err, "unable to write msg")
+
+				// Put the buffer back when done.
+				bufPool.Put(buf)
+			}
+		})
+	}
+}
+
+// BenchmarkReadMessage benchmarks the performance of lnwire.ReadMessage. It
+// first creates a test message for each of the lnwire.Message, writes it to
+// the buffer, then later reads it from the buffer.
+func BenchmarkReadMessage(b *testing.B) {
+	// Create testing messages. We will use a constant seed to make sure
+	// the benchmark uses the same data every time.
+	r := rand.New(rand.NewSource(42))
+	msgAll := makeAllMessages(b, r)
+
+	// Write all the messages to the buffer.
+	for _, msg := range msgAll {
+		// Fetch a buffer from the pool and reset it.
+		buf := bufPool.Get().(*bytes.Buffer)
+		buf.Reset()
+
+		_, err := lnwire.WriteMessage(buf, msg, 0)
+		require.NoError(b, err, "unable to write msg")
+
+		// Run each message as a sub benchmark test.
+		m := msg
+		b.Run(m.MsgType().String(), func(b *testing.B) {
+			b.ResetTimer()
+
+			for i := 0; i < b.N; i++ {
+				r := bytes.NewBuffer(buf.Bytes())
+
+				// Read the message from the buffer.
+				_, err := lnwire.ReadMessage(r, 0)
+				require.NoError(b, err, "unable to read msg")
+			}
+		})
+
+		// Put the buffer back when done.
+		bufPool.Put(buf)
+	}
+}
+
+// makeAllMessages is used to create testing messages for each lnwire message
+// type.
+//
+// TODO(yy): the following testing messages are created somewhat arbitrary. We
+// should standardlize each of the testing messages so that a better baseline
+// can be used.
+func makeAllMessages(t testing.TB, r *rand.Rand) []lnwire.Message {
+	msgAll := []lnwire.Message{}
+
+	msgAll = append(msgAll, newMsgInit(t, r))
+	msgAll = append(msgAll, newMsgError(t, r))
+	msgAll = append(msgAll, newMsgPing(t, r))
+	msgAll = append(msgAll, newMsgPong(t, r))
+	msgAll = append(msgAll, newMsgOpenChannel(t, r))
+	msgAll = append(msgAll, newMsgAcceptChannel(t, r))
+	msgAll = append(msgAll, newMsgFundingCreated(t, r))
+	msgAll = append(msgAll, newMsgFundingSigned(t, r))
+	msgAll = append(msgAll, newMsgFundingLocked(t, r))
+	msgAll = append(msgAll, newMsgShutdown(t, r))
+	msgAll = append(msgAll, newMsgClosingSigned(t, r))
+	msgAll = append(msgAll, newMsgUpdateAddHTLC(t, r))
+	msgAll = append(msgAll, newMsgUpdateFulfillHTLC(t, r))
+	msgAll = append(msgAll, newMsgUpdateFailHTLC(t, r))
+	msgAll = append(msgAll, newMsgCommitSig(t, r))
+	msgAll = append(msgAll, newMsgRevokeAndAck(t, r))
+	msgAll = append(msgAll, newMsgUpdateFee(t, r))
+	msgAll = append(msgAll, newMsgUpdateFailMalformedHTLC(t, r))
+	msgAll = append(msgAll, newMsgChannelReestablish(t, r))
+	msgAll = append(msgAll, newMsgChannelAnnouncement(t, r))
+	msgAll = append(msgAll, newMsgNodeAnnouncement(t, r))
+	msgAll = append(msgAll, newMsgChannelUpdate(t, r))
+	msgAll = append(msgAll, newMsgAnnounceSignatures(t, r))
+	msgAll = append(msgAll, newMsgQueryShortChanIDs(t, r))
+	msgAll = append(msgAll, newMsgReplyShortChanIDsEnd(t, r))
+	msgAll = append(msgAll, newMsgQueryChannelRange(t, r))
+	msgAll = append(msgAll, newMsgReplyChannelRange(t, r))
+	msgAll = append(msgAll, newMsgGossipTimestampRange(t, r))
+	msgAll = append(msgAll, newMsgQueryShortChanIDsZlib(t, r))
+	msgAll = append(msgAll, newMsgReplyChannelRangeZlib(t, r))
+
+	return msgAll
+}
+
+func newMsgInit(t testing.TB, r io.Reader) *lnwire.Init {
+	t.Helper()
+
+	return &lnwire.Init{
+		GlobalFeatures: rawFeatureVector(),
+		Features:       rawFeatureVector(),
+		ExtraData:      createExtraData(t, r),
+	}
+}
+
+// newMsgOpenChannel creates a testing OpenChannel message.
+func newMsgOpenChannel(t testing.TB, r *rand.Rand) *lnwire.OpenChannel {
+	t.Helper()
+
+	msg := &lnwire.OpenChannel{
+		FundingAmount:        btcutil.Amount(r.Int63()),
+		PushAmount:           lnwire.MilliSatoshi(r.Int63()),
+		DustLimit:            btcutil.Amount(r.Int63()),
+		MaxValueInFlight:     lnwire.MilliSatoshi(r.Int63()),
+		ChannelReserve:       btcutil.Amount(r.Int63()),
+		HtlcMinimum:          lnwire.MilliSatoshi(r.Int63()),
+		FeePerKiloWeight:     uint32(r.Int31()),
+		CsvDelay:             uint16(r.Intn(1 << 16)),
+		MaxAcceptedHTLCs:     uint16(r.Intn(1 << 16)),
+		ChannelFlags:         lnwire.FundingFlag(uint8(r.Intn(1 << 8))),
+		FundingKey:           randPubKey(t),
+		RevocationPoint:      randPubKey(t),
+		PaymentPoint:         randPubKey(t),
+		DelayedPaymentPoint:  randPubKey(t),
+		HtlcPoint:            randPubKey(t),
+		FirstCommitmentPoint: randPubKey(t),
+		ExtraData:            createExtraData(t, r),
+	}
+
+	_, err := r.Read(msg.ChainHash[:])
+	require.NoError(t, err, "unable to read bytes for ChainHash")
+
+	_, err = r.Read(msg.PendingChannelID[:])
+	require.NoError(t, err, "unable to read bytes for PendingChannelID")
+
+	return msg
+}
+
+func newMsgAcceptChannel(t testing.TB, r *rand.Rand) *lnwire.AcceptChannel {
+	t.Helper()
+
+	msg := &lnwire.AcceptChannel{
+		DustLimit:             btcutil.Amount(r.Int63()),
+		MaxValueInFlight:      lnwire.MilliSatoshi(r.Int63()),
+		ChannelReserve:        btcutil.Amount(r.Int63()),
+		MinAcceptDepth:        uint32(r.Int31()),
+		HtlcMinimum:           lnwire.MilliSatoshi(r.Int63()),
+		CsvDelay:              uint16(r.Intn(1 << 16)),
+		MaxAcceptedHTLCs:      uint16(r.Intn(1 << 16)),
+		FundingKey:            randPubKey(t),
+		RevocationPoint:       randPubKey(t),
+		PaymentPoint:          randPubKey(t),
+		DelayedPaymentPoint:   randPubKey(t),
+		HtlcPoint:             randPubKey(t),
+		FirstCommitmentPoint:  randPubKey(t),
+		UpfrontShutdownScript: randDeliveryAddress(t, r),
+		ExtraData:             createExtraData(t, r),
+	}
+	_, err := r.Read(msg.PendingChannelID[:])
+	require.NoError(t, err, "unable to generate pending chan id")
+
+	return msg
+}
+
+func newMsgError(t testing.TB, r io.Reader) *lnwire.Error {
+	t.Helper()
+
+	msg := lnwire.NewError()
+
+	_, err := r.Read(msg.ChanID[:])
+	require.NoError(t, err, "unable to generate chan id")
+
+	msg.Data = createExtraData(t, r)
+
+	return msg
+}
+
+func newMsgPing(t testing.TB, r *rand.Rand) *lnwire.Ping {
+	t.Helper()
+
+	return &lnwire.Ping{
+		NumPongBytes: uint16(r.Intn(1 << 16)),
+		PaddingBytes: createExtraData(t, r),
+	}
+}
+
+func newMsgPong(t testing.TB, r io.Reader) *lnwire.Pong {
+	t.Helper()
+
+	return &lnwire.Pong{
+		PongBytes: createExtraData(t, r),
+	}
+}
+
+func newMsgFundingCreated(t testing.TB, r *rand.Rand) *lnwire.FundingCreated {
+	t.Helper()
+
+	msg := &lnwire.FundingCreated{
+		CommitSig: testNodeSig,
+		ExtraData: createExtraData(t, r),
+	}
+
+	_, err := r.Read(msg.PendingChannelID[:])
+	require.NoError(t, err, "unable to generate pending chan id")
+
+	_, err = r.Read(msg.FundingPoint.Hash[:])
+	require.NoError(t, err, "unable to generate hash")
+
+	msg.FundingPoint.Index = uint32(r.Int31()) % math.MaxUint16
+
+	return msg
+}
+
+func newMsgFundingSigned(t testing.TB, r io.Reader) *lnwire.FundingSigned {
+	t.Helper()
+
+	var c [32]byte
+
+	_, err := r.Read(c[:])
+	require.NoError(t, err, "unable to generate chan id")
+
+	msg := &lnwire.FundingSigned{
+		ChanID:    lnwire.ChannelID(c),
+		CommitSig: testNodeSig,
+		ExtraData: createExtraData(t, r),
+	}
+
+	return msg
+}
+
+func newMsgFundingLocked(t testing.TB, r io.Reader) *lnwire.FundingLocked {
+	t.Helper()
+
+	var c [32]byte
+
+	_, err := r.Read(c[:])
+	require.NoError(t, err, "unable to generate chan id")
+
+	pubKey := randPubKey(t)
+
+	msg := lnwire.NewFundingLocked(lnwire.ChannelID(c), pubKey)
+	msg.ExtraData = createExtraData(t, r)
+
+	return msg
+}
+
+func newMsgShutdown(t testing.TB, r *rand.Rand) *lnwire.Shutdown {
+	t.Helper()
+
+	msg := &lnwire.Shutdown{
+		Address:   randDeliveryAddress(t, r),
+		ExtraData: createExtraData(t, r),
+	}
+
+	_, err := r.Read(msg.ChannelID[:])
+	require.NoError(t, err, "unable to generate channel id")
+
+	return msg
+}
+
+func newMsgClosingSigned(t testing.TB, r *rand.Rand) *lnwire.ClosingSigned {
+	t.Helper()
+
+	msg := &lnwire.ClosingSigned{
+		FeeSatoshis: btcutil.Amount(r.Int63()),
+		Signature:   testNodeSig,
+		ExtraData:   createExtraData(t, r),
+	}
+
+	_, err := r.Read(msg.ChannelID[:])
+	require.NoError(t, err, "unable to generate chan id")
+
+	return msg
+}
+
+func newMsgUpdateAddHTLC(t testing.TB, r *rand.Rand) *lnwire.UpdateAddHTLC {
+	t.Helper()
+
+	msg := &lnwire.UpdateAddHTLC{
+		ID:        r.Uint64(),
+		Amount:    lnwire.MilliSatoshi(r.Int63()),
+		Expiry:    r.Uint32(),
+		ExtraData: createExtraData(t, r),
+	}
+
+	_, err := r.Read(msg.ChanID[:])
+	require.NoError(t, err, "unable to generate chan id")
+
+	_, err = r.Read(msg.PaymentHash[:])
+	require.NoError(t, err, "unable to generate paymenthash")
+
+	_, err = r.Read(msg.OnionBlob[:])
+	require.NoError(t, err, "unable to generate onion blob")
+
+	return msg
+}
+
+func newMsgUpdateFulfillHTLC(t testing.TB,
+	r *rand.Rand) *lnwire.UpdateFulfillHTLC {
+
+	t.Helper()
+
+	msg := &lnwire.UpdateFulfillHTLC{
+		ID:        r.Uint64(),
+		ExtraData: createExtraData(t, r),
+	}
+
+	_, err := r.Read(msg.ChanID[:])
+	require.NoError(t, err, "unable to generate chan id")
+
+	_, err = r.Read(msg.PaymentPreimage[:])
+	require.NoError(t, err, "unable to generate payment preimage")
+
+	return msg
+}
+
+func newMsgUpdateFailHTLC(t testing.TB, r *rand.Rand) *lnwire.UpdateFailHTLC {
+	t.Helper()
+
+	msg := &lnwire.UpdateFailHTLC{
+		ID:        r.Uint64(),
+		ExtraData: createExtraData(t, r),
+	}
+
+	_, err := r.Read(msg.ChanID[:])
+	require.NoError(t, err, "unable to generate chan id")
+
+	return msg
+}
+
+func newMsgCommitSig(t testing.TB, r io.Reader) *lnwire.CommitSig {
+	t.Helper()
+
+	msg := lnwire.NewCommitSig()
+
+	_, err := r.Read(msg.ChanID[:])
+	require.NoError(t, err, "unable to generate chan id")
+
+	msg.CommitSig = testNodeSig
+	msg.ExtraData = createExtraData(t, r)
+
+	msg.HtlcSigs = make([]lnwire.Sig, testNumSigs)
+	for i := 0; i < testNumSigs; i++ {
+		msg.HtlcSigs[i] = testNodeSig
+	}
+
+	return msg
+}
+
+func newMsgRevokeAndAck(t testing.TB, r io.Reader) *lnwire.RevokeAndAck {
+	t.Helper()
+
+	msg := lnwire.NewRevokeAndAck()
+
+	_, err := r.Read(msg.ChanID[:])
+	require.NoError(t, err, "unable to generate chan id")
+
+	_, err = r.Read(msg.Revocation[:])
+	require.NoError(t, err, "unable to generate bytes")
+
+	msg.NextRevocationKey = randPubKey(t)
+	require.NoError(t, err, "unable to generate key")
+
+	msg.ExtraData = createExtraData(t, r)
+
+	return msg
+}
+
+func newMsgUpdateFee(t testing.TB, r *rand.Rand) *lnwire.UpdateFee {
+	t.Helper()
+
+	msg := &lnwire.UpdateFee{
+		FeePerKw:  uint32(r.Int31()),
+		ExtraData: createExtraData(t, r),
+	}
+
+	_, err := r.Read(msg.ChanID[:])
+	require.NoError(t, err, "unable to generate chan id")
+
+	return msg
+}
+
+func newMsgUpdateFailMalformedHTLC(t testing.TB,
+	r *rand.Rand) *lnwire.UpdateFailMalformedHTLC {
+
+	t.Helper()
+
+	msg := &lnwire.UpdateFailMalformedHTLC{
+		ID:          r.Uint64(),
+		FailureCode: lnwire.FailCode(r.Intn(1 << 16)),
+		ExtraData:   createExtraData(t, r),
+	}
+
+	_, err := r.Read(msg.ChanID[:])
+	require.NoError(t, err, "unable to generate chan id")
+
+	_, err = r.Read(msg.ShaOnionBlob[:])
+	require.NoError(t, err, "unable to generate sha256 onion blob")
+
+	return msg
+}
+
+func newMsgChannelReestablish(t testing.TB,
+	r *rand.Rand) *lnwire.ChannelReestablish {
+
+	t.Helper()
+
+	msg := &lnwire.ChannelReestablish{
+		NextLocalCommitHeight:     uint64(r.Int63()),
+		RemoteCommitTailHeight:    uint64(r.Int63()),
+		LocalUnrevokedCommitPoint: randPubKey(t),
+		ExtraData:                 createExtraData(t, r),
+	}
+
+	_, err := r.Read(msg.LastRemoteCommitSecret[:])
+	require.NoError(t, err, "unable to read commit secret")
+
+	return msg
+}
+
+func newMsgChannelAnnouncement(t testing.TB,
+	r *rand.Rand) *lnwire.ChannelAnnouncement {
+
+	t.Helper()
+
+	msg := &lnwire.ChannelAnnouncement{
+		ShortChannelID:  lnwire.NewShortChanIDFromInt(uint64(r.Int63())),
+		Features:        rawFeatureVector(),
+		NodeID1:         randRawKey(t),
+		NodeID2:         randRawKey(t),
+		BitcoinKey1:     randRawKey(t),
+		BitcoinKey2:     randRawKey(t),
+		ExtraOpaqueData: createExtraData(t, r),
+		NodeSig1:        testNodeSig,
+		NodeSig2:        testNodeSig,
+		BitcoinSig1:     testNodeSig,
+		BitcoinSig2:     testNodeSig,
+	}
+
+	_, err := r.Read(msg.ChainHash[:])
+	require.NoError(t, err, "unable to generate chain hash")
+
+	return msg
+}
+
+func newMsgNodeAnnouncement(t testing.TB,
+	r *rand.Rand) *lnwire.NodeAnnouncement {
+
+	t.Helper()
+
+	msg := &lnwire.NodeAnnouncement{
+		Features:  rawFeatureVector(),
+		Timestamp: uint32(r.Int31()),
+		Alias:     randAlias(r),
+		RGBColor: color.RGBA{
+			R: uint8(r.Intn(1 << 8)),
+			G: uint8(r.Intn(1 << 8)),
+			B: uint8(r.Intn(1 << 8)),
+		},
+		NodeID:          randRawKey(t),
+		Addresses:       randAddrs(t, r),
+		ExtraOpaqueData: createExtraData(t, r),
+		Signature:       testNodeSig,
+	}
+
+	return msg
+}
+
+func newMsgChannelUpdate(t testing.TB, r *rand.Rand) *lnwire.ChannelUpdate {
+	t.Helper()
+
+	msgFlags := lnwire.ChanUpdateMsgFlags(r.Int31())
+	maxHtlc := lnwire.MilliSatoshi(r.Int63())
+
+	// We make the max_htlc field zero if it is not flagged
+	// as being part of the ChannelUpdate, to pass
+	// serialization tests, as it will be ignored if the bit
+	// is not set.
+	if msgFlags&lnwire.ChanUpdateOptionMaxHtlc == 0 {
+		maxHtlc = 0
+	}
+
+	msg := &lnwire.ChannelUpdate{
+		ShortChannelID:  lnwire.NewShortChanIDFromInt(r.Uint64()),
+		Timestamp:       uint32(r.Int31()),
+		MessageFlags:    msgFlags,
+		ChannelFlags:    lnwire.ChanUpdateChanFlags(r.Int31()),
+		TimeLockDelta:   uint16(r.Int31()),
+		HtlcMinimumMsat: lnwire.MilliSatoshi(r.Int63()),
+		HtlcMaximumMsat: maxHtlc,
+		BaseFee:         uint32(r.Int31()),
+		FeeRate:         uint32(r.Int31()),
+		ExtraOpaqueData: createExtraData(t, r),
+		Signature:       testNodeSig,
+	}
+
+	_, err := r.Read(msg.ChainHash[:])
+	require.NoError(t, err, "unable to generate chain hash")
+
+	return msg
+}
+
+func newMsgAnnounceSignatures(t testing.TB,
+	r *rand.Rand) *lnwire.AnnounceSignatures {
+
+	t.Helper()
+
+	msg := &lnwire.AnnounceSignatures{
+		ShortChannelID: lnwire.NewShortChanIDFromInt(
+			uint64(r.Int63()),
+		),
+		ExtraOpaqueData:  createExtraData(t, r),
+		NodeSignature:    testNodeSig,
+		BitcoinSignature: testNodeSig,
+	}
+
+	_, err := r.Read(msg.ChannelID[:])
+	require.NoError(t, err, "unable to generate chan id")
+
+	return msg
+}
+
+func newMsgQueryShortChanIDs(t testing.TB,
+	r *rand.Rand) *lnwire.QueryShortChanIDs {
+
+	t.Helper()
+
+	msg := &lnwire.QueryShortChanIDs{
+		EncodingType: lnwire.EncodingSortedPlain,
+		ExtraData:    createExtraData(t, r),
+	}
+
+	_, err := rand.Read(msg.ChainHash[:])
+	require.NoError(t, err, "unable to read chain hash")
+
+	for i := 0; i < testNumChanIDs; i++ {
+		msg.ShortChanIDs = append(msg.ShortChanIDs,
+			lnwire.NewShortChanIDFromInt(uint64(r.Int63())))
+	}
+
+	return msg
+}
+
+func newMsgQueryShortChanIDsZlib(t testing.TB,
+	r *rand.Rand) *lnwire.QueryShortChanIDs {
+
+	t.Helper()
+
+	msg := &lnwire.QueryShortChanIDs{
+		EncodingType: lnwire.EncodingSortedZlib,
+		ExtraData:    createExtraData(t, r),
+	}
+
+	_, err := rand.Read(msg.ChainHash[:])
+	require.NoError(t, err, "unable to read chain hash")
+
+	for i := 0; i < testNumChanIDs; i++ {
+		msg.ShortChanIDs = append(msg.ShortChanIDs,
+			lnwire.NewShortChanIDFromInt(uint64(r.Int63())))
+	}
+
+	return msg
+}
+
+func newMsgReplyShortChanIDsEnd(t testing.TB,
+	r *rand.Rand) *lnwire.ReplyShortChanIDsEnd {
+
+	t.Helper()
+
+	msg := lnwire.NewReplyShortChanIDsEnd()
+
+	_, err := rand.Read(msg.ChainHash[:])
+	require.NoError(t, err, "unable to read chain hash")
+
+	msg.Complete = uint8(r.Int31n(2))
+	msg.ExtraData = createExtraData(t, r)
+
+	return msg
+}
+
+func newMsgQueryChannelRange(t testing.TB,
+	r *rand.Rand) *lnwire.QueryChannelRange {
+
+	t.Helper()
+
+	msg := lnwire.NewQueryChannelRange()
+
+	_, err := rand.Read(msg.ChainHash[:])
+	require.NoError(t, err, "unable to read chain hash")
+
+	msg.FirstBlockHeight = r.Uint32()
+	msg.NumBlocks = r.Uint32()
+	msg.ExtraData = createExtraData(t, r)
+
+	return msg
+}
+
+func newMsgReplyChannelRange(t testing.TB,
+	r *rand.Rand) *lnwire.ReplyChannelRange {
+
+	t.Helper()
+
+	msg := &lnwire.ReplyChannelRange{
+		EncodingType: lnwire.EncodingSortedPlain,
+		ExtraData:    createExtraData(t, r),
+	}
+
+	_, err := rand.Read(msg.ChainHash[:])
+	require.NoError(t, err, "unable to read chain hash")
+
+	msg.Complete = uint8(r.Int31n(2))
+
+	for i := 0; i < testNumChanIDs; i++ {
+		msg.ShortChanIDs = append(msg.ShortChanIDs,
+			lnwire.NewShortChanIDFromInt(uint64(r.Int63())))
+	}
+
+	return msg
+}
+
+func newMsgReplyChannelRangeZlib(t testing.TB,
+	r *rand.Rand) *lnwire.ReplyChannelRange {
+
+	t.Helper()
+
+	msg := &lnwire.ReplyChannelRange{
+		EncodingType: lnwire.EncodingSortedZlib,
+		ExtraData:    createExtraData(t, r),
+	}
+
+	_, err := rand.Read(msg.ChainHash[:])
+	require.NoError(t, err, "unable to read chain hash")
+
+	msg.Complete = uint8(r.Int31n(2))
+
+	for i := 0; i < testNumChanIDs; i++ {
+		msg.ShortChanIDs = append(msg.ShortChanIDs,
+			lnwire.NewShortChanIDFromInt(uint64(r.Int63())))
+	}
+
+	return msg
+}
+
+func newMsgGossipTimestampRange(t testing.TB,
+	r *rand.Rand) *lnwire.GossipTimestampRange {
+
+	t.Helper()
+
+	msg := lnwire.NewGossipTimestampRange()
+	msg.FirstTimestamp = r.Uint32()
+	msg.TimestampRange = r.Uint32()
+	msg.ExtraData = createExtraData(t, r)
+
+	_, err := r.Read(msg.ChainHash[:])
+	require.NoError(t, err, "unable to read chain hash")
+
+	return msg
+}
+
+func randRawKey(t testing.TB) [33]byte {
+	t.Helper()
+
+	var n [33]byte
+
+	priv, err := btcec.NewPrivateKey(btcec.S256())
+	require.NoError(t, err, "failed to create privKey")
+
+	copy(n[:], priv.PubKey().SerializeCompressed())
+
+	return n
+}
+
+func randPubKey(t testing.TB) *btcec.PublicKey {
+	t.Helper()
+
+	priv, err := btcec.NewPrivateKey(btcec.S256())
+	require.NoError(t, err, "failed to create pubkey")
+
+	return priv.PubKey()
+}
+
+func rawFeatureVector() *lnwire.RawFeatureVector {
+	// Get a slice of known feature bits.
+	featureBits := make([]lnwire.FeatureBit, 0, len(lnwire.Features))
+	for fb := range lnwire.Features {
+		featureBits = append(featureBits, fb)
+	}
+
+	featureVec := lnwire.NewRawFeatureVector(featureBits...)
+
+	return featureVec
+}
+
+func randDeliveryAddress(t testing.TB, r *rand.Rand) lnwire.DeliveryAddress {
+	t.Helper()
+
+	// Generate a max sized address.
+	size := r.Intn(deliveryAddressMaxSize) + 1
+	da := lnwire.DeliveryAddress(make([]byte, size))
+
+	_, err := r.Read(da)
+	require.NoError(t, err, "unable to read address")
+	return da
+}
+
+func randTCP4Addr(t testing.TB, r *rand.Rand) *net.TCPAddr {
+	t.Helper()
+
+	var ip [4]byte
+	_, err := r.Read(ip[:])
+	require.NoError(t, err, "unable to read ip")
+
+	var port [2]byte
+	_, err = r.Read(port[:])
+	require.NoError(t, err, "unable to read port")
+
+	addrIP := net.IP(ip[:])
+	addrPort := int(binary.BigEndian.Uint16(port[:]))
+
+	return &net.TCPAddr{IP: addrIP, Port: addrPort}
+}
+
+func randTCP6Addr(t testing.TB, r *rand.Rand) *net.TCPAddr {
+	t.Helper()
+
+	var ip [16]byte
+
+	_, err := r.Read(ip[:])
+	require.NoError(t, err, "unable to read ip")
+
+	var port [2]byte
+	_, err = r.Read(port[:])
+	require.NoError(t, err, "unable to read port")
+
+	addrIP := net.IP(ip[:])
+	addrPort := int(binary.BigEndian.Uint16(port[:]))
+
+	return &net.TCPAddr{IP: addrIP, Port: addrPort}
+}
+
+func randV2OnionAddr(t testing.TB, r *rand.Rand) *tor.OnionAddr {
+	t.Helper()
+
+	var serviceID [tor.V2DecodedLen]byte
+	_, err := r.Read(serviceID[:])
+	require.NoError(t, err, "unable to read serviceID")
+
+	var port [2]byte
+	_, err = r.Read(port[:])
+	require.NoError(t, err, "unable to read port")
+
+	onionService := tor.Base32Encoding.EncodeToString(serviceID[:])
+	onionService += tor.OnionSuffix
+	addrPort := int(binary.BigEndian.Uint16(port[:]))
+
+	return &tor.OnionAddr{OnionService: onionService, Port: addrPort}
+}
+
+func randV3OnionAddr(t testing.TB, r *rand.Rand) *tor.OnionAddr {
+	t.Helper()
+
+	var serviceID [tor.V3DecodedLen]byte
+	_, err := r.Read(serviceID[:])
+	require.NoError(t, err, "unable to read serviceID")
+
+	var port [2]byte
+	_, err = r.Read(port[:])
+	require.NoError(t, err, "unable to read port")
+
+	onionService := tor.Base32Encoding.EncodeToString(serviceID[:])
+	onionService += tor.OnionSuffix
+	addrPort := int(binary.BigEndian.Uint16(port[:]))
+
+	return &tor.OnionAddr{OnionService: onionService, Port: addrPort}
+}
+
+func randAddrs(t testing.TB, r *rand.Rand) []net.Addr {
+	tcp4Addr := randTCP4Addr(t, r)
+	tcp6Addr := randTCP6Addr(t, r)
+	v2OnionAddr := randV2OnionAddr(t, r)
+	v3OnionAddr := randV3OnionAddr(t, r)
+
+	return []net.Addr{tcp4Addr, tcp6Addr, v2OnionAddr, v3OnionAddr}
+}
+
+func randAlias(r *rand.Rand) lnwire.NodeAlias {
+	var a lnwire.NodeAlias
+	for i := range a {
+		a[i] = letterBytes[r.Intn(len(letterBytes))]
+	}
+
+	return a
+}
+
+func createExtraData(t testing.TB, r io.Reader) []byte {
+	t.Helper()
+
+	// Read random bytes.
+	extraData := make([]byte, testNumExtraBytes)
+	_, err := r.Read(extraData)
+	require.NoError(t, err, "unable to generate extra data")
+
+	// Encode the data length.
+	binary.BigEndian.PutUint16(extraData[:2], uint16(len(extraData[2:])))
+
+	return extraData
+}

From f212f1aa20f4f30c8504509459f61f4f3f8cc977 Mon Sep 17 00:00:00 2001
From: yyforyongyu <yy2452@columbia.edu>
Date: Thu, 17 Jun 2021 11:29:40 +0800
Subject: [PATCH 03/11] lnwire: refactor WriteMessage to use bytes.Buffer

This commit changes the method WriteMessage to use bytes.Buffer to save
heap allocations. A unit test is added to check the method is
implemented as expected.
---
 lnwire/message.go      |  84 +++++++++++++++---------
 lnwire/message_test.go | 144 +++++++++++++++++++++++++++++++++++++++++
 2 files changed, 199 insertions(+), 29 deletions(-)

diff --git a/lnwire/message.go b/lnwire/message.go
index c4f2fd36a..466dc50f2 100644
--- a/lnwire/message.go
+++ b/lnwire/message.go
@@ -52,6 +52,27 @@ const (
 	MsgGossipTimestampRange                = 265
 )
 
+// ErrorEncodeMessage is used when failed to encode the message payload.
+func ErrorEncodeMessage(err error) error {
+	return fmt.Errorf("failed to encode message to buffer, got %w", err)
+}
+
+// ErrorWriteMessageType is used when failed to write the message type.
+func ErrorWriteMessageType(err error) error {
+	return fmt.Errorf("failed to write message type, got %w", err)
+}
+
+// ErrorPayloadTooLarge is used when the payload size exceeds the
+// MaxMsgBody.
+func ErrorPayloadTooLarge(size int) error {
+	return fmt.Errorf(
+		"message payload is too large - encoded %d bytes, "+
+			"but maximum message payload is %d bytes",
+		size, MaxMsgBody,
+	)
+
+}
+
 // String return the string representation of message type.
 func (t MessageType) String() string {
 	switch t {
@@ -218,44 +239,49 @@ func makeEmptyMessage(msgType MessageType) (Message, error) {
 	return msg, nil
 }
 
-// WriteMessage writes a lightning Message to w including the necessary header
-// information and returns the number of bytes written.
-func WriteMessage(w io.Writer, msg Message, pver uint32) (int, error) {
-	totalBytes := 0
+// WriteMessage writes a lightning Message to a buffer including the necessary
+// header information and returns the number of bytes written. If any error is
+// encountered, the buffer passed will be reset to its original state since we
+// don't want any broken bytes left. In other words, no bytes will be written
+// if there's an error. Either all or none of the message bytes will be written
+// to the buffer.
+//
+// NOTE: this method is not concurrent safe.
+func WriteMessage(buf *bytes.Buffer, msg Message, pver uint32) (int, error) {
+	// Record the size of the bytes already written in buffer.
+	oldByteSize := buf.Len()
 
-	// Encode the message payload itself into a temporary buffer.
-	// TODO(roasbeef): create buffer pool
-	var bw bytes.Buffer
-	if err := msg.Encode(&bw, pver); err != nil {
-		return totalBytes, err
-	}
-	payload := bw.Bytes()
-	lenp := len(payload)
-
-	// Enforce maximum message payload, which means the body cannot be
-	// greater than MaxMsgBody.
-	if lenp > MaxMsgBody {
-		return totalBytes, fmt.Errorf("message payload is too large - "+
-			"encoded %d bytes, but maximum message body is %d bytes",
-			lenp, MaxMsgBody)
+	// cleanBrokenBytes is a helper closure that helps reset the buffer to
+	// its original state. It truncates all the bytes written in current
+	// scope.
+	var cleanBrokenBytes = func(b *bytes.Buffer) int {
+		b.Truncate(oldByteSize)
+		return 0
 	}
 
-	// With the initial sanity checks complete, we'll now write out the
-	// message type itself.
+	// Write the message type.
 	var mType [2]byte
 	binary.BigEndian.PutUint16(mType[:], uint16(msg.MsgType()))
-	n, err := w.Write(mType[:])
-	totalBytes += n
+	msgTypeBytes, err := buf.Write(mType[:])
 	if err != nil {
-		return totalBytes, err
+		return cleanBrokenBytes(buf), ErrorWriteMessageType(err)
 	}
 
-	// With the message type written, we'll now write out the raw payload
-	// itself.
-	n, err = w.Write(payload)
-	totalBytes += n
+	// Use the write buffer to encode our message.
+	if err := msg.Encode(buf, pver); err != nil {
+		return cleanBrokenBytes(buf), ErrorEncodeMessage(err)
+	}
 
-	return totalBytes, err
+	// Enforce maximum overall message payload. The write buffer now has
+	// the size of len(originalBytes) + len(payload) + len(type). We want
+	// to enforce the payload here, so we subtract it by the length of the
+	// type and old bytes.
+	lenp := buf.Len() - oldByteSize - msgTypeBytes
+	if lenp > MaxMsgBody {
+		return cleanBrokenBytes(buf), ErrorPayloadTooLarge(lenp)
+	}
+
+	return buf.Len() - oldByteSize, nil
 }
 
 // ReadMessage reads, validates, and parses the next Lightning message from r
diff --git a/lnwire/message_test.go b/lnwire/message_test.go
index fbade5082..4626eb3e7 100644
--- a/lnwire/message_test.go
+++ b/lnwire/message_test.go
@@ -3,6 +3,7 @@ package lnwire_test
 import (
 	"bytes"
 	"encoding/binary"
+	"errors"
 	"image/color"
 	"io"
 	"math"
@@ -16,6 +17,7 @@ import (
 	"github.com/btcsuite/btcutil"
 	"github.com/lightningnetwork/lnd/lnwire"
 	"github.com/lightningnetwork/lnd/tor"
+	"github.com/stretchr/testify/mock"
 	"github.com/stretchr/testify/require"
 )
 
@@ -41,6 +43,148 @@ var (
 	}
 )
 
+type mockMsg struct {
+	mock.Mock
+}
+
+func (m *mockMsg) Decode(r io.Reader, pver uint32) error {
+	args := m.Called(r, pver)
+	return args.Error(0)
+}
+
+func (m *mockMsg) Encode(w io.Writer, pver uint32) error {
+	args := m.Called(w, pver)
+	return args.Error(0)
+}
+
+func (m *mockMsg) MsgType() lnwire.MessageType {
+	args := m.Called()
+	return lnwire.MessageType(args.Int(0))
+}
+
+// A compile time check to ensure mockMsg implements the lnwire.Message
+// interface.
+var _ lnwire.Message = (*mockMsg)(nil)
+
+// TestWriteMessage tests the function lnwire.WriteMessage.
+func TestWriteMessage(t *testing.T) {
+	var (
+		buf = new(bytes.Buffer)
+
+		// encodeNormalSize specifies a message size that is normal.
+		encodeNormalSize = 1000
+
+		// encodeOversize specifies a message size that's too big.
+		encodeOversize = lnwire.MaxMsgBody + 1
+
+		// errDummy is returned by the msg.Encode when specified.
+		errDummy = errors.New("test error")
+
+		// oneByte is a dummy byte used to fill up the buffer.
+		oneByte = [1]byte{}
+	)
+
+	testCases := []struct {
+		name string
+
+		// encodeSize controls how many bytes are written to the buffer
+		// by the method msg.Encode(buf, pver).
+		encodeSize int
+
+		// encodeErr determines the return value of the method
+		// msg.Encode(buf, pver).
+		encodeErr error
+
+		errorExpected error
+	}{
+
+		{
+			name:          "successful write",
+			encodeSize:    encodeNormalSize,
+			encodeErr:     nil,
+			errorExpected: nil,
+		},
+		{
+			name:          "failed to encode payload",
+			encodeSize:    encodeNormalSize,
+			encodeErr:     errDummy,
+			errorExpected: lnwire.ErrorEncodeMessage(errDummy),
+		},
+		{
+			name:       "exceeds MaxMsgBody",
+			encodeSize: encodeOversize,
+			encodeErr:  nil,
+			errorExpected: lnwire.ErrorPayloadTooLarge(
+				encodeOversize,
+			),
+		},
+	}
+
+	for _, test := range testCases {
+		tc := test
+		t.Run(tc.name, func(t *testing.T) {
+			// Start the test by creating a mock message and patch
+			// the relevant methods.
+			msg := &mockMsg{}
+
+			// Use message type Ping here since all types are
+			// encoded using 2 bytes, it won't affect anything
+			// here.
+			msg.On("MsgType").Return(lnwire.MsgPing)
+
+			// Encode will return the specified error (could be
+			// nil) and has the side effect of filling up the
+			// buffer by repeating the oneByte encodeSize times.
+			msg.On("Encode", mock.Anything, mock.Anything).Return(
+				tc.encodeErr,
+			).Run(func(_ mock.Arguments) {
+				for i := 0; i < tc.encodeSize; i++ {
+					_, err := buf.Write(oneByte[:])
+					require.NoError(t, err)
+				}
+			})
+
+			// Record the initial state of the buffer and write the
+			// message.
+			oldBytesSize := buf.Len()
+			bytesWritten, err := lnwire.WriteMessage(
+				buf, msg, 1,
+			)
+
+			// Check that the returned error is expected.
+			require.Equal(
+				t, tc.errorExpected, err, "unexpected err",
+			)
+
+			// If there's an error, no bytes should be written to
+			// the buf.
+			if tc.errorExpected != nil {
+				require.Equal(
+					t, 0, bytesWritten,
+					"bytes written should be 0",
+				)
+
+				// We also check that the old buf was not
+				// affected.
+				require.Equal(
+					t, oldBytesSize, buf.Len(),
+					"original buffer should not change",
+				)
+			} else {
+				expected := buf.Len() - oldBytesSize
+				require.Equal(
+					t, expected, bytesWritten,
+					"bytes written not matched",
+				)
+			}
+
+			// Finally, check the mocked methods are called as
+			// expected.
+			msg.AssertExpectations(t)
+		})
+	}
+}
+
 // BenchmarkWriteMessage benchmarks the performance of lnwire.WriteMessage. It
 // generates a test message for each of the lnwire.Message, calls the
 // WriteMessage method and benchmark it.

From 879d3cc86c7d24f0ba9cf9fdaee9c606f9c53001 Mon Sep 17 00:00:00 2001
From: yyforyongyu <yy2452@columbia.edu>
Date: Thu, 17 Jun 2021 20:04:25 +0800
Subject: [PATCH 04/11] wtwire: remove reference from lnwire

---
 watchtower/wtwire/message.go | 11 ++++++++---
 1 file changed, 8 insertions(+), 3 deletions(-)

diff --git a/watchtower/wtwire/message.go b/watchtower/wtwire/message.go
index 2b3d22887..5bc8c75ce 100644
--- a/watchtower/wtwire/message.go
+++ b/watchtower/wtwire/message.go
@@ -5,8 +5,6 @@ import (
 	"encoding/binary"
 	"fmt"
 	"io"
-
-	"github.com/lightningnetwork/lnd/lnwire"
 )
 
 // MaxMessagePayload is the maximum bytes a message can be regardless of other
@@ -75,7 +73,14 @@ func (m MessageType) String() string {
 
 // Serializable is an interface which defines a lightning wire serializable
 // object.
-type Serializable = lnwire.Serializable
+type Serializable interface {
+	// Decode reads the bytes stream and converts it to the object.
+	Decode(io.Reader, uint32) error
+
+	// Encode converts object to the bytes stream and write it into the
+	// write buffer.
+	Encode(io.Writer, uint32) error
+}
 
 // Message is an interface that defines a lightning wire protocol message. The
 // interface is general in order to allow implementing types full control over

From aa1561c60d93db160399c2f22650c1e86761963f Mon Sep 17 00:00:00 2001
From: yyforyongyu <yy2452@columbia.edu>
Date: Thu, 17 Jun 2021 14:17:30 +0800
Subject: [PATCH 05/11] lnwire: use write buffer in Encode methods

This commit changes the WriteElement and WriteElements methods to take a
write buffer instead of io.Writer. The corresponding Encode methods are
changed to use the write buffer.
---
 chanbackup/single.go                 |  9 ++++++-
 channeldb/waitingproof.go            | 10 ++++++-
 lnwire/accept_channel.go             |  3 ++-
 lnwire/announcement_signatures.go    |  3 ++-
 lnwire/channel_announcement.go       |  2 +-
 lnwire/channel_reestablish.go        |  3 ++-
 lnwire/channel_update.go             |  2 +-
 lnwire/closing_signed.go             |  3 ++-
 lnwire/commit_sig.go                 |  3 ++-
 lnwire/error.go                      |  3 ++-
 lnwire/extra_bytes.go                |  2 +-
 lnwire/funding_created.go            |  3 ++-
 lnwire/funding_locked.go             |  3 ++-
 lnwire/funding_signed.go             |  7 +++--
 lnwire/gossip_timestamp_range.go     |  3 ++-
 lnwire/init_message.go               |  7 +++--
 lnwire/lnwire.go                     | 15 ++++-------
 lnwire/message.go                    |  4 +--
 lnwire/message_test.go               |  2 +-
 lnwire/node_announcement.go          |  2 +-
 lnwire/onion_error.go                | 40 +++++++++++++++++-----------
 lnwire/onion_error_test.go           |  4 ++-
 lnwire/open_channel.go               |  3 ++-
 lnwire/ping.go                       |  7 +++--
 lnwire/pong.go                       |  7 +++--
 lnwire/query_channel_range.go        |  3 ++-
 lnwire/query_short_chan_ids.go       | 35 +++++++++++++++---------
 lnwire/reply_channel_range.go        |  3 ++-
 lnwire/reply_short_chan_ids_end.go   |  3 ++-
 lnwire/revoke_and_ack.go             |  3 ++-
 lnwire/shutdown.go                   |  3 ++-
 lnwire/update_add_htlc.go            |  3 ++-
 lnwire/update_fail_htlc.go           |  3 ++-
 lnwire/update_fail_malformed_htlc.go |  3 ++-
 lnwire/update_fee.go                 |  3 ++-
 lnwire/update_fulfill_htlc.go        |  3 ++-
 36 files changed, 139 insertions(+), 76 deletions(-)

diff --git a/chanbackup/single.go b/chanbackup/single.go
index e51700c88..3007559e0 100644
--- a/chanbackup/single.go
+++ b/chanbackup/single.go
@@ -265,8 +265,15 @@ func (s *Single) Serialize(w io.Writer) error {
 		return err
 	}
 
+	// TODO(yy): remove the type assertion when we finished refactoring db
+	// into using write buffer.
+	buf, ok := w.(*bytes.Buffer)
+	if !ok {
+		return fmt.Errorf("expect io.Writer to be *bytes.Buffer")
+	}
+
 	return lnwire.WriteElements(
-		w,
+		buf,
 		byte(s.Version),
 		uint16(len(singleBytes.Bytes())),
 		singleBytes.Bytes(),
diff --git a/channeldb/waitingproof.go b/channeldb/waitingproof.go
index b35586bf0..e8a09b758 100644
--- a/channeldb/waitingproof.go
+++ b/channeldb/waitingproof.go
@@ -2,6 +2,7 @@ package channeldb
 
 import (
 	"encoding/binary"
+	"fmt"
 	"sync"
 
 	"io"
@@ -232,7 +233,14 @@ func (p *WaitingProof) Encode(w io.Writer) error {
 		return err
 	}
 
-	if err := p.AnnounceSignatures.Encode(w, 0); err != nil {
+	// TODO(yy): remove the type assertion when we finished refactoring db
+	// into using write buffer.
+	buf, ok := w.(*bytes.Buffer)
+	if !ok {
+		return fmt.Errorf("expect io.Writer to be *bytes.Buffer")
+	}
+
+	if err := p.AnnounceSignatures.Encode(buf, 0); err != nil {
 		return err
 	}
 
diff --git a/lnwire/accept_channel.go b/lnwire/accept_channel.go
index f88a1a08f..f356e8943 100644
--- a/lnwire/accept_channel.go
+++ b/lnwire/accept_channel.go
@@ -1,6 +1,7 @@
 package lnwire
 
 import (
+	"bytes"
 	"fmt"
 	"io"
 
@@ -115,7 +116,7 @@ var _ Message = (*AcceptChannel)(nil)
 // protocol version.
 //
 // This is part of the lnwire.Message interface.
-func (a *AcceptChannel) Encode(w io.Writer, pver uint32) error {
+func (a *AcceptChannel) Encode(w *bytes.Buffer, pver uint32) error {
 	// Since the upfront script is encoded as a TLV record, concatenate it
 	// with the ExtraData, and write them as one.
 	tlvRecords, err := packShutdownScript(
diff --git a/lnwire/announcement_signatures.go b/lnwire/announcement_signatures.go
index 41372db19..2550d1c1a 100644
--- a/lnwire/announcement_signatures.go
+++ b/lnwire/announcement_signatures.go
@@ -1,6 +1,7 @@
 package lnwire
 
 import (
+	"bytes"
 	"io"
 )
 
@@ -64,7 +65,7 @@ func (a *AnnounceSignatures) Decode(r io.Reader, pver uint32) error {
 // observing the protocol version specified.
 //
 // This is part of the lnwire.Message interface.
-func (a *AnnounceSignatures) Encode(w io.Writer, pver uint32) error {
+func (a *AnnounceSignatures) Encode(w *bytes.Buffer, pver uint32) error {
 	return WriteElements(w,
 		a.ChannelID,
 		a.ShortChannelID,
diff --git a/lnwire/channel_announcement.go b/lnwire/channel_announcement.go
index 865b7b52c..9e57d17bc 100644
--- a/lnwire/channel_announcement.go
+++ b/lnwire/channel_announcement.go
@@ -87,7 +87,7 @@ func (a *ChannelAnnouncement) Decode(r io.Reader, pver uint32) error {
 // observing the protocol version specified.
 //
 // This is part of the lnwire.Message interface.
-func (a *ChannelAnnouncement) Encode(w io.Writer, pver uint32) error {
+func (a *ChannelAnnouncement) Encode(w *bytes.Buffer, pver uint32) error {
 	return WriteElements(w,
 		a.NodeSig1,
 		a.NodeSig2,
diff --git a/lnwire/channel_reestablish.go b/lnwire/channel_reestablish.go
index ebf2fcc1f..44cc6f3f6 100644
--- a/lnwire/channel_reestablish.go
+++ b/lnwire/channel_reestablish.go
@@ -1,6 +1,7 @@
 package lnwire
 
 import (
+	"bytes"
 	"io"
 
 	"github.com/btcsuite/btcd/btcec"
@@ -75,7 +76,7 @@ var _ Message = (*ChannelReestablish)(nil)
 // observing the protocol version specified.
 //
 // This is part of the lnwire.Message interface.
-func (a *ChannelReestablish) Encode(w io.Writer, pver uint32) error {
+func (a *ChannelReestablish) Encode(w *bytes.Buffer, pver uint32) error {
 	err := WriteElements(w,
 		a.ChanID,
 		a.NextLocalCommitHeight,
diff --git a/lnwire/channel_update.go b/lnwire/channel_update.go
index 87a7d30a1..e1bac9f96 100644
--- a/lnwire/channel_update.go
+++ b/lnwire/channel_update.go
@@ -159,7 +159,7 @@ func (a *ChannelUpdate) Decode(r io.Reader, pver uint32) error {
 // observing the protocol version specified.
 //
 // This is part of the lnwire.Message interface.
-func (a *ChannelUpdate) Encode(w io.Writer, pver uint32) error {
+func (a *ChannelUpdate) Encode(w *bytes.Buffer, pver uint32) error {
 	err := WriteElements(w,
 		a.Signature,
 		a.ChainHash[:],
diff --git a/lnwire/closing_signed.go b/lnwire/closing_signed.go
index 3e4dcdf26..66ccb7bcb 100644
--- a/lnwire/closing_signed.go
+++ b/lnwire/closing_signed.go
@@ -1,6 +1,7 @@
 package lnwire
 
 import (
+	"bytes"
 	"io"
 
 	"github.com/btcsuite/btcutil"
@@ -63,7 +64,7 @@ func (c *ClosingSigned) Decode(r io.Reader, pver uint32) error {
 // observing the protocol version specified.
 //
 // This is part of the lnwire.Message interface.
-func (c *ClosingSigned) Encode(w io.Writer, pver uint32) error {
+func (c *ClosingSigned) Encode(w *bytes.Buffer, pver uint32) error {
 	return WriteElements(
 		w, c.ChannelID, c.FeeSatoshis, c.Signature, c.ExtraData,
 	)
diff --git a/lnwire/commit_sig.go b/lnwire/commit_sig.go
index 3251fe8dc..59a785a0d 100644
--- a/lnwire/commit_sig.go
+++ b/lnwire/commit_sig.go
@@ -1,6 +1,7 @@
 package lnwire
 
 import (
+	"bytes"
 	"io"
 )
 
@@ -69,7 +70,7 @@ func (c *CommitSig) Decode(r io.Reader, pver uint32) error {
 // observing the protocol version specified.
 //
 // This is part of the lnwire.Message interface.
-func (c *CommitSig) Encode(w io.Writer, pver uint32) error {
+func (c *CommitSig) Encode(w *bytes.Buffer, pver uint32) error {
 	return WriteElements(w,
 		c.ChanID,
 		c.CommitSig,
diff --git a/lnwire/error.go b/lnwire/error.go
index 0aa2d4c31..0b95607f2 100644
--- a/lnwire/error.go
+++ b/lnwire/error.go
@@ -1,6 +1,7 @@
 package lnwire
 
 import (
+	"bytes"
 	"fmt"
 	"io"
 )
@@ -103,7 +104,7 @@ func (c *Error) Decode(r io.Reader, pver uint32) error {
 // protocol version specified.
 //
 // This is part of the lnwire.Message interface.
-func (c *Error) Encode(w io.Writer, pver uint32) error {
+func (c *Error) Encode(w *bytes.Buffer, pver uint32) error {
 	return WriteElements(w,
 		c.ChanID,
 		c.Data,
diff --git a/lnwire/extra_bytes.go b/lnwire/extra_bytes.go
index 22fd20bd8..f2dbec45b 100644
--- a/lnwire/extra_bytes.go
+++ b/lnwire/extra_bytes.go
@@ -16,7 +16,7 @@ import (
 type ExtraOpaqueData []byte
 
 // Encode attempts to encode the raw extra bytes into the passed io.Writer.
-func (e *ExtraOpaqueData) Encode(w io.Writer) error {
+func (e *ExtraOpaqueData) Encode(w *bytes.Buffer) error {
 	eBytes := []byte((*e)[:])
 	if err := WriteElements(w, eBytes); err != nil {
 		return err
diff --git a/lnwire/funding_created.go b/lnwire/funding_created.go
index 0f10214ed..6ca8c1b01 100644
--- a/lnwire/funding_created.go
+++ b/lnwire/funding_created.go
@@ -1,6 +1,7 @@
 package lnwire
 
 import (
+	"bytes"
 	"io"
 
 	"github.com/btcsuite/btcd/wire"
@@ -40,7 +41,7 @@ var _ Message = (*FundingCreated)(nil)
 // protocol version.
 //
 // This is part of the lnwire.Message interface.
-func (f *FundingCreated) Encode(w io.Writer, pver uint32) error {
+func (f *FundingCreated) Encode(w *bytes.Buffer, pver uint32) error {
 	return WriteElements(
 		w, f.PendingChannelID[:], f.FundingPoint, f.CommitSig,
 		f.ExtraData,
diff --git a/lnwire/funding_locked.go b/lnwire/funding_locked.go
index 0f8e4ca1e..7c3b2eb3a 100644
--- a/lnwire/funding_locked.go
+++ b/lnwire/funding_locked.go
@@ -1,6 +1,7 @@
 package lnwire
 
 import (
+	"bytes"
 	"io"
 
 	"github.com/btcsuite/btcd/btcec"
@@ -58,7 +59,7 @@ func (c *FundingLocked) Decode(r io.Reader, pver uint32) error {
 // protocol version.
 //
 // This is part of the lnwire.Message interface.
-func (c *FundingLocked) Encode(w io.Writer, pver uint32) error {
+func (c *FundingLocked) Encode(w *bytes.Buffer, pver uint32) error {
 	return WriteElements(w,
 		c.ChanID,
 		c.NextPerCommitmentPoint,
diff --git a/lnwire/funding_signed.go b/lnwire/funding_signed.go
index 15c62cf05..d45fbd40e 100644
--- a/lnwire/funding_signed.go
+++ b/lnwire/funding_signed.go
@@ -1,6 +1,9 @@
 package lnwire
 
-import "io"
+import (
+	"bytes"
+	"io"
+)
 
 // FundingSigned is sent from Bob (the responder) to Alice (the initiator)
 // after receiving the funding outpoint and her signature for Bob's version of
@@ -29,7 +32,7 @@ var _ Message = (*FundingSigned)(nil)
 // protocol version.
 //
 // This is part of the lnwire.Message interface.
-func (f *FundingSigned) Encode(w io.Writer, pver uint32) error {
+func (f *FundingSigned) Encode(w *bytes.Buffer, pver uint32) error {
 	return WriteElements(w, f.ChanID, f.CommitSig, f.ExtraData)
 }
 
diff --git a/lnwire/gossip_timestamp_range.go b/lnwire/gossip_timestamp_range.go
index 6827b9818..e9665d157 100644
--- a/lnwire/gossip_timestamp_range.go
+++ b/lnwire/gossip_timestamp_range.go
@@ -1,6 +1,7 @@
 package lnwire
 
 import (
+	"bytes"
 	"io"
 
 	"github.com/btcsuite/btcd/chaincfg/chainhash"
@@ -57,7 +58,7 @@ func (g *GossipTimestampRange) Decode(r io.Reader, pver uint32) error {
 // observing the protocol version specified.
 //
 // This is part of the lnwire.Message interface.
-func (g *GossipTimestampRange) Encode(w io.Writer, pver uint32) error {
+func (g *GossipTimestampRange) Encode(w *bytes.Buffer, pver uint32) error {
 	return WriteElements(w,
 		g.ChainHash[:],
 		g.FirstTimestamp,
diff --git a/lnwire/init_message.go b/lnwire/init_message.go
index 18fa693b4..be25642e1 100644
--- a/lnwire/init_message.go
+++ b/lnwire/init_message.go
@@ -1,6 +1,9 @@
 package lnwire
 
-import "io"
+import (
+	"bytes"
+	"io"
+)
 
 // Init is the first message reveals the features supported or required by this
 // node. Nodes wait for receipt of the other's features to simplify error
@@ -56,7 +59,7 @@ func (msg *Init) Decode(r io.Reader, pver uint32) error {
 // the protocol version specified.
 //
 // This is part of the lnwire.Message interface.
-func (msg *Init) Encode(w io.Writer, pver uint32) error {
+func (msg *Init) Encode(w *bytes.Buffer, pver uint32) error {
 	return WriteElements(w,
 		msg.GlobalFeatures,
 		msg.Features,
diff --git a/lnwire/lnwire.go b/lnwire/lnwire.go
index c180cad38..8400c210e 100644
--- a/lnwire/lnwire.go
+++ b/lnwire/lnwire.go
@@ -75,13 +75,8 @@ func (a addressType) AddrLen() uint16 {
 }
 
 // WriteElement is a one-stop shop to write the big endian representation of
-// any element which is to be serialized for the wire protocol. The passed
-// io.Writer should be backed by an appropriately sized byte slice, or be able
-// to dynamically expand to accommodate additional data.
-//
-// TODO(roasbeef): this should eventually draw from a buffer pool for
-// serialization.
-func WriteElement(w io.Writer, element interface{}) error {
+// any element which is to be serialized for the wire protocol.
+func WriteElement(w *bytes.Buffer, element interface{}) error {
 	switch e := element.(type) {
 	case NodeAlias:
 		if _, err := w.Write(e[:]); err != nil {
@@ -437,10 +432,10 @@ func WriteElement(w io.Writer, element interface{}) error {
 }
 
 // WriteElements is writes each element in the elements slice to the passed
-// io.Writer using WriteElement.
-func WriteElements(w io.Writer, elements ...interface{}) error {
+// buffer using WriteElement.
+func WriteElements(buf *bytes.Buffer, elements ...interface{}) error {
 	for _, element := range elements {
-		err := WriteElement(w, element)
+		err := WriteElement(buf, element)
 		if err != nil {
 			return err
 		}
diff --git a/lnwire/message.go b/lnwire/message.go
index 466dc50f2..ccd24d10b 100644
--- a/lnwire/message.go
+++ b/lnwire/message.go
@@ -158,8 +158,8 @@ type Serializable interface {
 	Decode(io.Reader, uint32) error
 
 	// Encode converts object to the bytes stream and write it into the
-	// writer.
-	Encode(io.Writer, uint32) error
+	// write buffer.
+	Encode(*bytes.Buffer, uint32) error
 }
 
 // Message is an interface that defines a lightning wire protocol message. The
diff --git a/lnwire/message_test.go b/lnwire/message_test.go
index 4626eb3e7..55e0c0153 100644
--- a/lnwire/message_test.go
+++ b/lnwire/message_test.go
@@ -52,7 +52,7 @@ func (m *mockMsg) Decode(r io.Reader, pver uint32) error {
 	return args.Error(0)
 }
 
-func (m *mockMsg) Encode(w io.Writer, pver uint32) error {
+func (m *mockMsg) Encode(w *bytes.Buffer, pver uint32) error {
 	args := m.Called(w, pver)
 	return args.Error(0)
 }
diff --git a/lnwire/node_announcement.go b/lnwire/node_announcement.go
index 540137aae..664bda84d 100644
--- a/lnwire/node_announcement.go
+++ b/lnwire/node_announcement.go
@@ -124,7 +124,7 @@ func (a *NodeAnnouncement) Decode(r io.Reader, pver uint32) error {
 // Encode serializes the target NodeAnnouncement into the passed io.Writer
 // observing the protocol version specified.
 //
-func (a *NodeAnnouncement) Encode(w io.Writer, pver uint32) error {
+func (a *NodeAnnouncement) Encode(w *bytes.Buffer, pver uint32) error {
 	return WriteElements(w,
 		a.Signature,
 		a.Features,
diff --git a/lnwire/onion_error.go b/lnwire/onion_error.go
index 35555e266..f9c86db61 100644
--- a/lnwire/onion_error.go
+++ b/lnwire/onion_error.go
@@ -419,7 +419,7 @@ func (f *FailIncorrectDetails) Decode(r io.Reader, pver uint32) error {
 // Encode writes the failure in bytes stream.
 //
 // NOTE: Part of the Serializable interface.
-func (f *FailIncorrectDetails) Encode(w io.Writer, pver uint32) error {
+func (f *FailIncorrectDetails) Encode(w *bytes.Buffer, pver uint32) error {
 	return WriteElements(w, f.amount, f.height)
 }
 
@@ -485,7 +485,7 @@ func (f *FailInvalidOnionVersion) Decode(r io.Reader, pver uint32) error {
 // Encode writes the failure in bytes stream.
 //
 // NOTE: Part of the Serializable interface.
-func (f *FailInvalidOnionVersion) Encode(w io.Writer, pver uint32) error {
+func (f *FailInvalidOnionVersion) Encode(w *bytes.Buffer, pver uint32) error {
 	return WriteElement(w, f.OnionSHA256[:])
 }
 
@@ -519,7 +519,7 @@ func (f *FailInvalidOnionHmac) Decode(r io.Reader, pver uint32) error {
 // Encode writes the failure in bytes stream.
 //
 // NOTE: Part of the Serializable interface.
-func (f *FailInvalidOnionHmac) Encode(w io.Writer, pver uint32) error {
+func (f *FailInvalidOnionHmac) Encode(w *bytes.Buffer, pver uint32) error {
 	return WriteElement(w, f.OnionSHA256[:])
 }
 
@@ -561,7 +561,7 @@ func (f *FailInvalidOnionKey) Decode(r io.Reader, pver uint32) error {
 // Encode writes the failure in bytes stream.
 //
 // NOTE: Part of the Serializable interface.
-func (f *FailInvalidOnionKey) Encode(w io.Writer, pver uint32) error {
+func (f *FailInvalidOnionKey) Encode(w *bytes.Buffer, pver uint32) error {
 	return WriteElement(w, f.OnionSHA256[:])
 }
 
@@ -670,7 +670,9 @@ func (f *FailTemporaryChannelFailure) Decode(r io.Reader, pver uint32) error {
 // Encode writes the failure in bytes stream.
 //
 // NOTE: Part of the Serializable interface.
-func (f *FailTemporaryChannelFailure) Encode(w io.Writer, pver uint32) error {
+func (f *FailTemporaryChannelFailure) Encode(w *bytes.Buffer,
+	pver uint32) error {
+
 	var payload []byte
 	if f.Update != nil {
 		var bw bytes.Buffer
@@ -749,7 +751,7 @@ func (f *FailAmountBelowMinimum) Decode(r io.Reader, pver uint32) error {
 // Encode writes the failure in bytes stream.
 //
 // NOTE: Part of the Serializable interface.
-func (f *FailAmountBelowMinimum) Encode(w io.Writer, pver uint32) error {
+func (f *FailAmountBelowMinimum) Encode(w *bytes.Buffer, pver uint32) error {
 	if err := WriteElement(w, f.HtlcMsat); err != nil {
 		return err
 	}
@@ -817,7 +819,7 @@ func (f *FailFeeInsufficient) Decode(r io.Reader, pver uint32) error {
 // Encode writes the failure in bytes stream.
 //
 // NOTE: Part of the Serializable interface.
-func (f *FailFeeInsufficient) Encode(w io.Writer, pver uint32) error {
+func (f *FailFeeInsufficient) Encode(w *bytes.Buffer, pver uint32) error {
 	if err := WriteElement(w, f.HtlcMsat); err != nil {
 		return err
 	}
@@ -885,7 +887,7 @@ func (f *FailIncorrectCltvExpiry) Decode(r io.Reader, pver uint32) error {
 // Encode writes the failure in bytes stream.
 //
 // NOTE: Part of the Serializable interface.
-func (f *FailIncorrectCltvExpiry) Encode(w io.Writer, pver uint32) error {
+func (f *FailIncorrectCltvExpiry) Encode(w *bytes.Buffer, pver uint32) error {
 	if err := WriteElement(w, f.CltvExpiry); err != nil {
 		return err
 	}
@@ -942,7 +944,7 @@ func (f *FailExpiryTooSoon) Decode(r io.Reader, pver uint32) error {
 // Encode writes the failure in bytes stream.
 //
 // NOTE: Part of the Serializable interface.
-func (f *FailExpiryTooSoon) Encode(w io.Writer, pver uint32) error {
+func (f *FailExpiryTooSoon) Encode(w *bytes.Buffer, pver uint32) error {
 	return writeOnionErrorChanUpdate(w, &f.Update, pver)
 }
 
@@ -1006,7 +1008,7 @@ func (f *FailChannelDisabled) Decode(r io.Reader, pver uint32) error {
 // Encode writes the failure in bytes stream.
 //
 // NOTE: Part of the Serializable interface.
-func (f *FailChannelDisabled) Encode(w io.Writer, pver uint32) error {
+func (f *FailChannelDisabled) Encode(w *bytes.Buffer, pver uint32) error {
 	if err := WriteElement(w, f.Flags); err != nil {
 		return err
 	}
@@ -1056,7 +1058,9 @@ func (f *FailFinalIncorrectCltvExpiry) Decode(r io.Reader, pver uint32) error {
 // Encode writes the failure in bytes stream.
 //
 // NOTE: Part of the Serializable interface.
-func (f *FailFinalIncorrectCltvExpiry) Encode(w io.Writer, pver uint32) error {
+func (f *FailFinalIncorrectCltvExpiry) Encode(w *bytes.Buffer,
+	pver uint32) error {
+
 	return WriteElement(w, f.CltvExpiry)
 }
 
@@ -1102,7 +1106,9 @@ func (f *FailFinalIncorrectHtlcAmount) Decode(r io.Reader, pver uint32) error {
 // Encode writes the failure in bytes stream.
 //
 // NOTE: Part of the Serializable interface.
-func (f *FailFinalIncorrectHtlcAmount) Encode(w io.Writer, pver uint32) error {
+func (f *FailFinalIncorrectHtlcAmount) Encode(w *bytes.Buffer,
+	pver uint32) error {
+
 	return WriteElement(w, f.IncomingHTLCAmount)
 }
 
@@ -1177,7 +1183,7 @@ func (f *InvalidOnionPayload) Decode(r io.Reader, pver uint32) error {
 // Encode writes the failure in bytes stream.
 //
 // NOTE: Part of the Serializable interface.
-func (f *InvalidOnionPayload) Encode(w io.Writer, pver uint32) error {
+func (f *InvalidOnionPayload) Encode(w *bytes.Buffer, pver uint32) error {
 	var buf [8]byte
 	if err := tlv.WriteVarInt(w, f.Type, &buf); err != nil {
 		return err
@@ -1263,7 +1269,7 @@ func DecodeFailureMessage(r io.Reader, pver uint32) (FailureMessage, error) {
 
 // EncodeFailure encodes, including the necessary onion failure header
 // information.
-func EncodeFailure(w io.Writer, failure FailureMessage, pver uint32) error {
+func EncodeFailure(w *bytes.Buffer, failure FailureMessage, pver uint32) error {
 	var failureMessageBuffer bytes.Buffer
 
 	err := EncodeFailureMessage(&failureMessageBuffer, failure, pver)
@@ -1293,7 +1299,9 @@ func EncodeFailure(w io.Writer, failure FailureMessage, pver uint32) error {
 
 // EncodeFailureMessage encodes just the failure message without adding a length
 // and padding the message for the onion protocol.
-func EncodeFailureMessage(w io.Writer, failure FailureMessage, pver uint32) error {
+func EncodeFailureMessage(w *bytes.Buffer,
+	failure FailureMessage, pver uint32) error {
+
 	// First, we'll write out the error code itself into the failure
 	// buffer.
 	var codeBytes [2]byte
@@ -1401,7 +1409,7 @@ func makeEmptyOnionError(code FailCode) (FailureMessage, error) {
 // writeOnionErrorChanUpdate writes out a ChannelUpdate using the onion error
 // format. The format is that we first write out the true serialized length of
 // the channel update, followed by the serialized channel update itself.
-func writeOnionErrorChanUpdate(w io.Writer, chanUpdate *ChannelUpdate,
+func writeOnionErrorChanUpdate(w *bytes.Buffer, chanUpdate *ChannelUpdate,
 	pver uint32) error {
 
 	// First, we encode the channel update in a temporary buffer in order
diff --git a/lnwire/onion_error_test.go b/lnwire/onion_error_test.go
index 8c4c131c6..927b6ee9c 100644
--- a/lnwire/onion_error_test.go
+++ b/lnwire/onion_error_test.go
@@ -275,6 +275,8 @@ func (f *mockFailIncorrectDetailsNoHeight) Decode(r io.Reader, pver uint32) erro
 	return nil
 }
 
-func (f *mockFailIncorrectDetailsNoHeight) Encode(w io.Writer, pver uint32) error {
+func (f *mockFailIncorrectDetailsNoHeight) Encode(w *bytes.Buffer,
+	pver uint32) error {
+
 	return WriteElement(w, f.amount)
 }
diff --git a/lnwire/open_channel.go b/lnwire/open_channel.go
index 9407d0037..d975749d2 100644
--- a/lnwire/open_channel.go
+++ b/lnwire/open_channel.go
@@ -1,6 +1,7 @@
 package lnwire
 
 import (
+	"bytes"
 	"io"
 
 	"github.com/btcsuite/btcd/btcec"
@@ -150,7 +151,7 @@ var _ Message = (*OpenChannel)(nil)
 // protocol version.
 //
 // This is part of the lnwire.Message interface.
-func (o *OpenChannel) Encode(w io.Writer, pver uint32) error {
+func (o *OpenChannel) Encode(w *bytes.Buffer, pver uint32) error {
 	// Since the upfront script is encoded as a TLV record, concatenate it
 	// with the ExtraData, and write them as one.
 	tlvRecords, err := packShutdownScript(
diff --git a/lnwire/ping.go b/lnwire/ping.go
index e7e160d44..9218fe7f6 100644
--- a/lnwire/ping.go
+++ b/lnwire/ping.go
@@ -1,6 +1,9 @@
 package lnwire
 
-import "io"
+import (
+	"bytes"
+	"io"
+)
 
 // PingPayload is a set of opaque bytes used to pad out a ping message.
 type PingPayload []byte
@@ -44,7 +47,7 @@ func (p *Ping) Decode(r io.Reader, pver uint32) error {
 // protocol version specified.
 //
 // This is part of the lnwire.Message interface.
-func (p *Ping) Encode(w io.Writer, pver uint32) error {
+func (p *Ping) Encode(w *bytes.Buffer, pver uint32) error {
 	return WriteElements(w,
 		p.NumPongBytes,
 		p.PaddingBytes)
diff --git a/lnwire/pong.go b/lnwire/pong.go
index b0d54523b..6c70c120f 100644
--- a/lnwire/pong.go
+++ b/lnwire/pong.go
@@ -1,6 +1,9 @@
 package lnwire
 
-import "io"
+import (
+	"bytes"
+	"io"
+)
 
 // PongPayload is a set of opaque bytes sent in response to a ping message.
 type PongPayload []byte
@@ -40,7 +43,7 @@ func (p *Pong) Decode(r io.Reader, pver uint32) error {
 // protocol version specified.
 //
 // This is part of the lnwire.Message interface.
-func (p *Pong) Encode(w io.Writer, pver uint32) error {
+func (p *Pong) Encode(w *bytes.Buffer, pver uint32) error {
 	return WriteElements(w,
 		p.PongBytes,
 	)
diff --git a/lnwire/query_channel_range.go b/lnwire/query_channel_range.go
index 167a4c64e..a6a20e1fd 100644
--- a/lnwire/query_channel_range.go
+++ b/lnwire/query_channel_range.go
@@ -1,6 +1,7 @@
 package lnwire
 
 import (
+	"bytes"
 	"io"
 	"math"
 
@@ -58,7 +59,7 @@ func (q *QueryChannelRange) Decode(r io.Reader, pver uint32) error {
 // observing the protocol version specified.
 //
 // This is part of the lnwire.Message interface.
-func (q *QueryChannelRange) Encode(w io.Writer, pver uint32) error {
+func (q *QueryChannelRange) Encode(w *bytes.Buffer, pver uint32) error {
 	return WriteElements(w,
 		q.ChainHash[:],
 		q.FirstBlockHeight,
diff --git a/lnwire/query_short_chan_ids.go b/lnwire/query_short_chan_ids.go
index 92ffaf585..930becdbb 100644
--- a/lnwire/query_short_chan_ids.go
+++ b/lnwire/query_short_chan_ids.go
@@ -291,7 +291,7 @@ func decodeShortChanIDs(r io.Reader) (ShortChanIDEncoding, []ShortChannelID, err
 // observing the protocol version specified.
 //
 // This is part of the lnwire.Message interface.
-func (q *QueryShortChanIDs) Encode(w io.Writer, pver uint32) error {
+func (q *QueryShortChanIDs) Encode(w *bytes.Buffer, pver uint32) error {
 	// First, we'll write out the chain hash.
 	err := WriteElements(w, q.ChainHash[:])
 	if err != nil {
@@ -310,7 +310,7 @@ func (q *QueryShortChanIDs) Encode(w io.Writer, pver uint32) error {
 
 // encodeShortChanIDs encodes the passed short channel ID's into the passed
 // io.Writer, respecting the specified encoding type.
-func encodeShortChanIDs(w io.Writer, encodingType ShortChanIDEncoding,
+func encodeShortChanIDs(w *bytes.Buffer, encodingType ShortChanIDEncoding,
 	shortChanIDs []ShortChannelID, noSort bool) error {
 
 	// For both of the current encoding types, the channel ID's are to be
@@ -360,29 +360,40 @@ func encodeShortChanIDs(w io.Writer, encodingType ShortChanIDEncoding,
 	// TODO(roasbeef): assumes the caller knows the proper chunk size to
 	// pass to avoid bin-packing here
 	case EncodingSortedZlib:
-		// We'll make a new buffer, then wrap that with a zlib writer
-		// so we can write directly to the buffer and encode in a
-		// streaming manner.
-		var buf bytes.Buffer
-		zlibWriter := zlib.NewWriter(&buf)
-
 		// If we don't have anything at all to write, then we'll write
 		// an empty payload so we don't include things like the zlib
 		// header when the remote party is expecting no actual short
 		// channel IDs.
 		var compressedPayload []byte
 		if len(shortChanIDs) > 0 {
+			// We'll make a new write buffer to hold the bytes of
+			// shortChanIDs.
+			var wb bytes.Buffer
+
 			// Next, we'll write out all the channel ID's directly
 			// into the zlib writer, which will do compressing on
 			// the fly.
 			for _, chanID := range shortChanIDs {
-				err := WriteElements(zlibWriter, chanID)
+				err := WriteElements(&wb, chanID)
 				if err != nil {
-					return fmt.Errorf("unable to write short chan "+
-						"ID: %v", err)
+					return fmt.Errorf(
+						"unable to write short chan "+
+							"ID: %v", err,
+					)
 				}
 			}
 
+			// With shortChanIDs written into wb, we'll create a
+			// zlib writer and write all the compressed bytes.
+			var zlibBuffer bytes.Buffer
+			zlibWriter := zlib.NewWriter(&zlibBuffer)
+
+			if _, err := zlibWriter.Write(wb.Bytes()); err != nil {
+				return fmt.Errorf(
+					"unable to write compressed short chan"+
+						"ID: %w", err)
+			}
+
 			// Now that we've written all the elements, we'll
 			// ensure the compressed stream is written to the
 			// underlying buffer.
@@ -391,7 +402,7 @@ func encodeShortChanIDs(w io.Writer, encodingType ShortChanIDEncoding,
 					"compression: %v", err)
 			}
 
-			compressedPayload = buf.Bytes()
+			compressedPayload = zlibBuffer.Bytes()
 		}
 
 		// Now that we have all the items compressed, we can compute
diff --git a/lnwire/reply_channel_range.go b/lnwire/reply_channel_range.go
index 2e8085c65..00bb10eb2 100644
--- a/lnwire/reply_channel_range.go
+++ b/lnwire/reply_channel_range.go
@@ -1,6 +1,7 @@
 package lnwire
 
 import (
+	"bytes"
 	"io"
 	"math"
 
@@ -85,7 +86,7 @@ func (c *ReplyChannelRange) Decode(r io.Reader, pver uint32) error {
 // observing the protocol version specified.
 //
 // This is part of the lnwire.Message interface.
-func (c *ReplyChannelRange) Encode(w io.Writer, pver uint32) error {
+func (c *ReplyChannelRange) Encode(w *bytes.Buffer, pver uint32) error {
 	err := WriteElements(w,
 		c.ChainHash[:],
 		c.FirstBlockHeight,
diff --git a/lnwire/reply_short_chan_ids_end.go b/lnwire/reply_short_chan_ids_end.go
index 2280dee4b..6d909febf 100644
--- a/lnwire/reply_short_chan_ids_end.go
+++ b/lnwire/reply_short_chan_ids_end.go
@@ -1,6 +1,7 @@
 package lnwire
 
 import (
+	"bytes"
 	"io"
 
 	"github.com/btcsuite/btcd/chaincfg/chainhash"
@@ -54,7 +55,7 @@ func (c *ReplyShortChanIDsEnd) Decode(r io.Reader, pver uint32) error {
 // observing the protocol version specified.
 //
 // This is part of the lnwire.Message interface.
-func (c *ReplyShortChanIDsEnd) Encode(w io.Writer, pver uint32) error {
+func (c *ReplyShortChanIDsEnd) Encode(w *bytes.Buffer, pver uint32) error {
 	return WriteElements(w,
 		c.ChainHash[:],
 		c.Complete,
diff --git a/lnwire/revoke_and_ack.go b/lnwire/revoke_and_ack.go
index 6562547fe..acbfcae6e 100644
--- a/lnwire/revoke_and_ack.go
+++ b/lnwire/revoke_and_ack.go
@@ -1,6 +1,7 @@
 package lnwire
 
 import (
+	"bytes"
 	"io"
 
 	"github.com/btcsuite/btcd/btcec"
@@ -65,7 +66,7 @@ func (c *RevokeAndAck) Decode(r io.Reader, pver uint32) error {
 // observing the protocol version specified.
 //
 // This is part of the lnwire.Message interface.
-func (c *RevokeAndAck) Encode(w io.Writer, pver uint32) error {
+func (c *RevokeAndAck) Encode(w *bytes.Buffer, pver uint32) error {
 	return WriteElements(w,
 		c.ChanID,
 		c.Revocation[:],
diff --git a/lnwire/shutdown.go b/lnwire/shutdown.go
index 4eec13523..e7b26283b 100644
--- a/lnwire/shutdown.go
+++ b/lnwire/shutdown.go
@@ -1,6 +1,7 @@
 package lnwire
 
 import (
+	"bytes"
 	"io"
 )
 
@@ -46,7 +47,7 @@ func (s *Shutdown) Decode(r io.Reader, pver uint32) error {
 // the protocol version specified.
 //
 // This is part of the lnwire.Message interface.
-func (s *Shutdown) Encode(w io.Writer, pver uint32) error {
+func (s *Shutdown) Encode(w *bytes.Buffer, pver uint32) error {
 	return WriteElements(w, s.ChannelID, s.Address, s.ExtraData)
 }
 
diff --git a/lnwire/update_add_htlc.go b/lnwire/update_add_htlc.go
index 35fc5fc36..a7756994c 100644
--- a/lnwire/update_add_htlc.go
+++ b/lnwire/update_add_htlc.go
@@ -1,6 +1,7 @@
 package lnwire
 
 import (
+	"bytes"
 	"io"
 )
 
@@ -88,7 +89,7 @@ func (c *UpdateAddHTLC) Decode(r io.Reader, pver uint32) error {
 // the protocol version specified.
 //
 // This is part of the lnwire.Message interface.
-func (c *UpdateAddHTLC) Encode(w io.Writer, pver uint32) error {
+func (c *UpdateAddHTLC) Encode(w *bytes.Buffer, pver uint32) error {
 	return WriteElements(w,
 		c.ChanID,
 		c.ID,
diff --git a/lnwire/update_fail_htlc.go b/lnwire/update_fail_htlc.go
index 8706c1dc8..2344b5683 100644
--- a/lnwire/update_fail_htlc.go
+++ b/lnwire/update_fail_htlc.go
@@ -1,6 +1,7 @@
 package lnwire
 
 import (
+	"bytes"
 	"io"
 )
 
@@ -54,7 +55,7 @@ func (c *UpdateFailHTLC) Decode(r io.Reader, pver uint32) error {
 // the protocol version specified.
 //
 // This is part of the lnwire.Message interface.
-func (c *UpdateFailHTLC) Encode(w io.Writer, pver uint32) error {
+func (c *UpdateFailHTLC) Encode(w *bytes.Buffer, pver uint32) error {
 	return WriteElements(w,
 		c.ChanID,
 		c.ID,
diff --git a/lnwire/update_fail_malformed_htlc.go b/lnwire/update_fail_malformed_htlc.go
index e7994fccc..120f89541 100644
--- a/lnwire/update_fail_malformed_htlc.go
+++ b/lnwire/update_fail_malformed_htlc.go
@@ -1,6 +1,7 @@
 package lnwire
 
 import (
+	"bytes"
 	"crypto/sha256"
 	"io"
 )
@@ -53,7 +54,7 @@ func (c *UpdateFailMalformedHTLC) Decode(r io.Reader, pver uint32) error {
 // io.Writer observing the protocol version specified.
 //
 // This is part of the lnwire.Message interface.
-func (c *UpdateFailMalformedHTLC) Encode(w io.Writer, pver uint32) error {
+func (c *UpdateFailMalformedHTLC) Encode(w *bytes.Buffer, pver uint32) error {
 	return WriteElements(w,
 		c.ChanID,
 		c.ID,
diff --git a/lnwire/update_fee.go b/lnwire/update_fee.go
index 375b5c671..a30634d79 100644
--- a/lnwire/update_fee.go
+++ b/lnwire/update_fee.go
@@ -1,6 +1,7 @@
 package lnwire
 
 import (
+	"bytes"
 	"io"
 )
 
@@ -51,7 +52,7 @@ func (c *UpdateFee) Decode(r io.Reader, pver uint32) error {
 // observing the protocol version specified.
 //
 // This is part of the lnwire.Message interface.
-func (c *UpdateFee) Encode(w io.Writer, pver uint32) error {
+func (c *UpdateFee) Encode(w *bytes.Buffer, pver uint32) error {
 	return WriteElements(w,
 		c.ChanID,
 		c.FeePerKw,
diff --git a/lnwire/update_fulfill_htlc.go b/lnwire/update_fulfill_htlc.go
index 2b7f2e494..4cd3a9720 100644
--- a/lnwire/update_fulfill_htlc.go
+++ b/lnwire/update_fulfill_htlc.go
@@ -1,6 +1,7 @@
 package lnwire
 
 import (
+	"bytes"
 	"io"
 )
 
@@ -60,7 +61,7 @@ func (c *UpdateFulfillHTLC) Decode(r io.Reader, pver uint32) error {
 // observing the protocol version specified.
 //
 // This is part of the lnwire.Message interface.
-func (c *UpdateFulfillHTLC) Encode(w io.Writer, pver uint32) error {
+func (c *UpdateFulfillHTLC) Encode(w *bytes.Buffer, pver uint32) error {
 	return WriteElements(w,
 		c.ChanID,
 		c.ID,

From e381b33b215d6f7e3a1d3e5c0d673d576dd597ac Mon Sep 17 00:00:00 2001
From: yyforyongyu <yy2452@columbia.edu>
Date: Thu, 17 Jun 2021 20:21:48 +0800
Subject: [PATCH 06/11] lnwire: fix typo in docs

---
 lnwire/features.go | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/lnwire/features.go b/lnwire/features.go
index 9dfa07ee1..675ba0f50 100644
--- a/lnwire/features.go
+++ b/lnwire/features.go
@@ -91,7 +91,7 @@ const (
 	// attacks on the receiver of a payment.
 	PaymentAddrOptional FeatureBit = 15
 
-	// MPPOptional is a required feature bit that signals that the receiver
+	// MPPRequired is a required feature bit that signals that the receiver
 	// of a payment requires settlement of an invoice with more than one
 	// HTLC.
 	MPPRequired FeatureBit = 16
@@ -124,7 +124,7 @@ const (
 	// transactions, which also imply anchor commitments.
 	AnchorsZeroFeeHtlcTxRequired FeatureBit = 22
 
-	// AnchorsZeroFeeHtlcTxRequired is an optional feature bit that signals
+	// AnchorsZeroFeeHtlcTxOptional is an optional feature bit that signals
 	// that the node supports channels having zero-fee second-level HTLC
 	// transactions, which also imply anchor commitments.
 	AnchorsZeroFeeHtlcTxOptional FeatureBit = 23

From f04410c546c10b748ac7b8125e62b56642fb1b3f Mon Sep 17 00:00:00 2001
From: yyforyongyu <yy2452@columbia.edu>
Date: Fri, 18 Jun 2021 14:40:59 +0800
Subject: [PATCH 07/11] lnwire: add independent writers

This commit breaks the method WriteElement and adds specific writers for
each of the data types.
---
 lnwire/lnwire.go      |  11 +-
 lnwire/writer.go      | 431 +++++++++++++++++++++++++++++
 lnwire/writer_test.go | 618 ++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 1059 insertions(+), 1 deletion(-)
 create mode 100644 lnwire/writer.go
 create mode 100644 lnwire/writer_test.go

diff --git a/lnwire/lnwire.go b/lnwire/lnwire.go
index 8400c210e..20edcbab4 100644
--- a/lnwire/lnwire.go
+++ b/lnwire/lnwire.go
@@ -76,6 +76,9 @@ func (a addressType) AddrLen() uint16 {
 
 // WriteElement is a one-stop shop to write the big endian representation of
 // any element which is to be serialized for the wire protocol.
+//
+// TODO(yy): rm this method once we finish dereferencing it from other
+// packages.
 func WriteElement(w *bytes.Buffer, element interface{}) error {
 	switch e := element.(type) {
 	case NodeAlias:
@@ -433,6 +436,9 @@ func WriteElement(w *bytes.Buffer, element interface{}) error {
 
 // WriteElements is writes each element in the elements slice to the passed
 // buffer using WriteElement.
+//
+// TODO(yy): rm this method once we finish dereferencing it from other
+// packages.
 func WriteElements(buf *bytes.Buffer, elements ...interface{}) error {
 	for _, element := range elements {
 		err := WriteElement(buf, element)
@@ -823,8 +829,11 @@ func ReadElement(r io.Reader, element interface{}) error {
 		length := binary.BigEndian.Uint16(addrLen[:])
 
 		var addrBytes [deliveryAddressMaxSize]byte
+
 		if length > deliveryAddressMaxSize {
-			return fmt.Errorf("cannot read %d bytes into addrBytes", length)
+			return fmt.Errorf(
+				"cannot read %d bytes into addrBytes", length,
+			)
 		}
 		if _, err = io.ReadFull(r, addrBytes[:length]); err != nil {
 			return err
diff --git a/lnwire/writer.go b/lnwire/writer.go
new file mode 100644
index 000000000..6a20c5715
--- /dev/null
+++ b/lnwire/writer.go
@@ -0,0 +1,431 @@
+package lnwire
+
+import (
+	"bytes"
+	"encoding/binary"
+	"errors"
+	"fmt"
+	"image/color"
+	"math"
+	"net"
+
+	"github.com/btcsuite/btcd/btcec"
+	"github.com/btcsuite/btcd/wire"
+	"github.com/btcsuite/btcutil"
+	"github.com/lightningnetwork/lnd/tor"
+)
+
+var (
+	// ErrNilFeatureVector is returned when the supplied feature is nil.
+	ErrNilFeatureVector = errors.New("cannot write nil feature vector")
+
+	// ErrPkScriptTooLong is returned when the length of the provided
+	// script exceeds 34.
+	ErrPkScriptTooLong = errors.New("'PkScript' too long")
+
+	// ErrNilTCPAddress is returned when the supplied address is nil.
+	ErrNilTCPAddress = errors.New("cannot write nil TCPAddr")
+
+	// ErrNilOnionAddress is returned when the supplied address is nil.
+	ErrNilOnionAddress = errors.New("cannot write nil onion address")
+
+	// ErrNilNetAddress is returned when a nil value is used in []net.Addr.
+	ErrNilNetAddress = errors.New("cannot write nil address")
+
+	// ErrNilPublicKey is returned when a nil pubkey is used.
+	ErrNilPublicKey = errors.New("cannot write nil pubkey")
+
+	// ErrUnknownServiceLength is returned when the onion service length is
+	// unknown.
+	ErrUnknownServiceLength = errors.New("unknown onion service length")
+)
+
+// ErrOutpointIndexTooBig is used when the outpoint index exceeds the max value
+// of uint16.
+func ErrOutpointIndexTooBig(index uint32) error {
+	return fmt.Errorf(
+		"index for outpoint (%v) is greater than "+
+			"max index of %v", index, math.MaxUint16,
+	)
+}
+
+// WriteBytes appends the given bytes to the provided buffer.
+//
+// Note: We intentionally skip the interfacer linter check here because we want
+// to have concrete type (bytes.Buffer) rather than interface type (io.Write)
+// due to performance concern.
+func WriteBytes(buf *bytes.Buffer, b []byte) error { // nolint: interfacer
+	_, err := buf.Write(b)
+	return err
+}
+
+// WriteUint8 appends the uint8 to the provided buffer.
+//
+// Note: We intentionally skip the interfacer linter check here because we want
+// to have concrete type (bytes.Buffer) rather than interface type (io.Write)
+// due to performance concern.
+func WriteUint8(buf *bytes.Buffer, n uint8) error { // nolint: interfacer
+	_, err := buf.Write([]byte{n})
+	return err
+}
+
+// WriteUint16 appends the uint16 to the provided buffer. It encodes the
+// integer using big endian byte order.
+//
+// Note: We intentionally skip the interfacer linter check here because we want
+// to have concrete type (bytes.Buffer) rather than interface type (io.Write)
+// due to performance concern.
+func WriteUint16(buf *bytes.Buffer, n uint16) error { // nolint: interfacer
+	var b [2]byte
+	binary.BigEndian.PutUint16(b[:], n)
+	_, err := buf.Write(b[:])
+	return err
+}
+
+// WriteUint32 appends the uint32 to the provided buffer. It encodes the
+// integer using big endian byte order.
+func WriteUint32(buf *bytes.Buffer, n uint32) error {
+	var b [4]byte
+	binary.BigEndian.PutUint32(b[:], n)
+	_, err := buf.Write(b[:])
+	return err
+}
+
+// WriteUint64 appends the uint64 to the provided buffer. It encodes the
+// integer using big endian byte order.
+//
+// Note: We intentionally skip the interfacer linter check here because we want
+// to have concrete type (bytes.Buffer) rather than interface type (io.Write)
+// due to performance concern.
+func WriteUint64(buf *bytes.Buffer, n uint64) error { // nolint: interfacer
+	var b [8]byte
+	binary.BigEndian.PutUint64(b[:], n)
+	_, err := buf.Write(b[:])
+	return err
+}
+
+// WriteSatoshi appends the Satoshi value to the provided buffer.
+func WriteSatoshi(buf *bytes.Buffer, amount btcutil.Amount) error {
+	return WriteUint64(buf, uint64(amount))
+}
+
+// WriteMilliSatoshi appends the MilliSatoshi value to the provided buffer.
+func WriteMilliSatoshi(buf *bytes.Buffer, amount MilliSatoshi) error {
+	return WriteUint64(buf, uint64(amount))
+}
+
+// WritePublicKey appends the compressed public key to the provided buffer.
+func WritePublicKey(buf *bytes.Buffer, pub *btcec.PublicKey) error {
+	if pub == nil {
+		return ErrNilPublicKey
+	}
+
+	serializedPubkey := pub.SerializeCompressed()
+	return WriteBytes(buf, serializedPubkey)
+
+}
+
+// WriteChannelID appends the ChannelID to the provided buffer.
+func WriteChannelID(buf *bytes.Buffer, channelID ChannelID) error {
+	return WriteBytes(buf, channelID[:])
+}
+
+// WriteNodeAlias appends the alias to the provided buffer.
+func WriteNodeAlias(buf *bytes.Buffer, alias NodeAlias) error {
+	return WriteBytes(buf, alias[:])
+}
+
+// WriteShortChannelID appends the ShortChannelID to the provided buffer. It
+// encodes the BlockHeight and TxIndex each using 3 bytes with big endian byte
+// order, and encodes txPosition using 2 bytes with big endian byte order.
+func WriteShortChannelID(buf *bytes.Buffer, shortChanID ShortChannelID) error {
+	// Check that field fit in 3 bytes and write the blockHeight
+	if shortChanID.BlockHeight > ((1 << 24) - 1) {
+		return errors.New("block height should fit in 3 bytes")
+	}
+
+	var blockHeight [4]byte
+	binary.BigEndian.PutUint32(blockHeight[:], shortChanID.BlockHeight)
+
+	if _, err := buf.Write(blockHeight[1:]); err != nil {
+		return err
+	}
+
+	// Check that field fit in 3 bytes and write the txIndex
+	if shortChanID.TxIndex > ((1 << 24) - 1) {
+		return errors.New("tx index should fit in 3 bytes")
+	}
+
+	var txIndex [4]byte
+	binary.BigEndian.PutUint32(txIndex[:], shortChanID.TxIndex)
+	if _, err := buf.Write(txIndex[1:]); err != nil {
+		return err
+	}
+
+	// Write the TxPosition
+	return WriteUint16(buf, shortChanID.TxPosition)
+}
+
+// WriteSig appends the signature to the provided buffer.
+func WriteSig(buf *bytes.Buffer, sig Sig) error {
+	return WriteBytes(buf, sig[:])
+}
+
+// WriteSigs appends the slice of signatures to the provided buffer with its
+// length.
+func WriteSigs(buf *bytes.Buffer, sigs []Sig) error {
+	// Write the length of the sigs.
+	if err := WriteUint16(buf, uint16(len(sigs))); err != nil {
+		return err
+	}
+
+	for _, sig := range sigs {
+		if err := WriteSig(buf, sig); err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
+// WriteFailCode appends the FailCode to the provided buffer.
+func WriteFailCode(buf *bytes.Buffer, e FailCode) error {
+	return WriteUint16(buf, uint16(e))
+}
+
+// WriteRawFeatureVector encodes the feature using the feature's Encode method
+// and appends the data to the provided buffer. An error will return if the
+// passed feature is nil.
+//
+// Note: We intentionally skip the interfacer linter check here because we want
+// to have concrete type (bytes.Buffer) rather than interface type (io.Write)
+// due to performance concern.
+func WriteRawFeatureVector(buf *bytes.Buffer, // nolint: interfacer
+	feature *RawFeatureVector) error {
+
+	if feature == nil {
+		return ErrNilFeatureVector
+	}
+
+	return feature.Encode(buf)
+}
+
+// WriteColorRGBA appends the RGBA color using three bytes.
+func WriteColorRGBA(buf *bytes.Buffer, e color.RGBA) error {
+	// Write R
+	if err := WriteUint8(buf, e.R); err != nil {
+		return err
+	}
+
+	// Write G
+	if err := WriteUint8(buf, e.G); err != nil {
+		return err
+	}
+
+	// Write B
+	return WriteUint8(buf, e.B)
+}
+
+// WriteShortChanIDEncoding appends the ShortChanIDEncoding to the provided
+// buffer.
+func WriteShortChanIDEncoding(buf *bytes.Buffer, e ShortChanIDEncoding) error {
+	return WriteUint8(buf, uint8(e))
+}
+
+// WriteFundingFlag appends the FundingFlag to the provided buffer.
+func WriteFundingFlag(buf *bytes.Buffer, flag FundingFlag) error {
+	return WriteUint8(buf, uint8(flag))
+}
+
+// WriteChanUpdateMsgFlags appends the update flag to the provided buffer.
+func WriteChanUpdateMsgFlags(buf *bytes.Buffer, f ChanUpdateMsgFlags) error {
+	return WriteUint8(buf, uint8(f))
+}
+
+// WriteChanUpdateChanFlags appends the update flag to the provided buffer.
+func WriteChanUpdateChanFlags(buf *bytes.Buffer, f ChanUpdateChanFlags) error {
+	return WriteUint8(buf, uint8(f))
+}
+
+// WriteDeliveryAddress appends the address to the provided buffer.
+func WriteDeliveryAddress(buf *bytes.Buffer, addr DeliveryAddress) error {
+	return writeDataWithLength(buf, addr)
+}
+
+// WritePingPayload appends the payload to the provided buffer.
+func WritePingPayload(buf *bytes.Buffer, payload PingPayload) error {
+	return writeDataWithLength(buf, payload)
+}
+
+// WritePongPayload appends the payload to the provided buffer.
+func WritePongPayload(buf *bytes.Buffer, payload PongPayload) error {
+	return writeDataWithLength(buf, payload)
+}
+
+// WriteErrorData appends the data to the provided buffer.
+func WriteErrorData(buf *bytes.Buffer, data ErrorData) error {
+	return writeDataWithLength(buf, data)
+}
+
+// WriteOpaqueReason appends the reason to the provided buffer.
+func WriteOpaqueReason(buf *bytes.Buffer, reason OpaqueReason) error {
+	return writeDataWithLength(buf, reason)
+}
+
+// WriteBool appends the boolean to the provided buffer.
+func WriteBool(buf *bytes.Buffer, b bool) error {
+	if b {
+		return WriteBytes(buf, []byte{1})
+	}
+	return WriteBytes(buf, []byte{0})
+}
+
+// WritePkScript appends the script to the provided buffer. Returns an error if
+// the provided script exceeds 34 bytes.
+//
+// Note: We intentionally skip the interfacer linter check here because we want
+// to have concrete type (bytes.Buffer) rather than interface type (io.Write)
+// due to performance concern.
+func WritePkScript(buf *bytes.Buffer, s PkScript) error { // nolint: interfacer
+	// The largest script we'll accept is a p2wsh which is exactly
+	// 34 bytes long.
+	scriptLength := len(s)
+	if scriptLength > 34 {
+		return ErrPkScriptTooLong
+	}
+
+	return wire.WriteVarBytes(buf, 0, s)
+}
+
+// WriteOutPoint appends the outpoint to the provided buffer.
+func WriteOutPoint(buf *bytes.Buffer, p wire.OutPoint) error {
+	// Before we write anything to the buffer, check the Index is sane.
+	if p.Index > math.MaxUint16 {
+		return ErrOutpointIndexTooBig(p.Index)
+	}
+
+	var h [32]byte
+	copy(h[:], p.Hash[:])
+	if _, err := buf.Write(h[:]); err != nil {
+		return err
+	}
+
+	// Write the index using two bytes.
+	return WriteUint16(buf, uint16(p.Index))
+}
+
+// WriteTCPAddr appends the TCP address to the provided buffer, either a IPv4
+// or a IPv6.
+func WriteTCPAddr(buf *bytes.Buffer, addr *net.TCPAddr) error {
+	if addr == nil {
+		return ErrNilTCPAddress
+	}
+
+	// Make a slice of bytes to hold the data of descriptor and ip. At
+	// most, we need 17 bytes - 1 byte for the descriptor, 16 bytes for
+	// IPv6.
+	data := make([]byte, 0, 17)
+
+	if addr.IP.To4() != nil {
+		data = append(data, uint8(tcp4Addr))
+		data = append(data, addr.IP.To4()...)
+	} else {
+		data = append(data, uint8(tcp6Addr))
+		data = append(data, addr.IP.To16()...)
+	}
+
+	if _, err := buf.Write(data); err != nil {
+		return err
+	}
+
+	return WriteUint16(buf, uint16(addr.Port))
+}
+
+// WriteOnionAddr appends the onion address to the provided buffer.
+func WriteOnionAddr(buf *bytes.Buffer, addr *tor.OnionAddr) error {
+	if addr == nil {
+		return ErrNilOnionAddress
+	}
+
+	var (
+		suffixIndex int
+		descriptor  []byte
+	)
+
+	// Decide the suffixIndex and descriptor.
+	switch len(addr.OnionService) {
+	case tor.V2Len:
+		descriptor = []byte{byte(v2OnionAddr)}
+		suffixIndex = tor.V2Len - tor.OnionSuffixLen
+
+	case tor.V3Len:
+		descriptor = []byte{byte(v3OnionAddr)}
+		suffixIndex = tor.V3Len - tor.OnionSuffixLen
+
+	default:
+		return ErrUnknownServiceLength
+	}
+
+	// Decode the address.
+	host, err := tor.Base32Encoding.DecodeString(
+		addr.OnionService[:suffixIndex],
+	)
+	if err != nil {
+		return err
+	}
+
+	// Perform the actual write when the above checks passed.
+	if _, err := buf.Write(descriptor); err != nil {
+		return err
+	}
+	if _, err := buf.Write(host); err != nil {
+		return err
+	}
+
+	return WriteUint16(buf, uint16(addr.Port))
+}
+
+// WriteNetAddrs appends a slice of addresses to the provided buffer with the
+// length info.
+func WriteNetAddrs(buf *bytes.Buffer, addresses []net.Addr) error {
+	// First, we'll encode all the addresses into an intermediate
+	// buffer. We need to do this in order to compute the total
+	// length of the addresses.
+	buffer := make([]byte, 0, MaxMsgBody)
+	addrBuf := bytes.NewBuffer(buffer)
+
+	for _, address := range addresses {
+		switch a := address.(type) {
+		case *net.TCPAddr:
+			if err := WriteTCPAddr(addrBuf, a); err != nil {
+				return err
+			}
+		case *tor.OnionAddr:
+			if err := WriteOnionAddr(addrBuf, a); err != nil {
+				return err
+			}
+		default:
+			return ErrNilNetAddress
+		}
+	}
+
+	// With the addresses fully encoded, we can now write out data.
+	return writeDataWithLength(buf, addrBuf.Bytes())
+}
+
+// writeDataWithLength writes the data and its length to the buffer.
+//
+// Note: We intentionally skip the interfacer linter check here because we want
+// to have concrete type (bytes.Buffer) rather than interface type (io.Write)
+// due to performance concern.
+func writeDataWithLength(buf *bytes.Buffer, // nolint: interfacer
+	data []byte) error {
+
+	var l [2]byte
+	binary.BigEndian.PutUint16(l[:], uint16(len(data)))
+	if _, err := buf.Write(l[:]); err != nil {
+		return err
+	}
+
+	_, err := buf.Write(data)
+	return err
+}
diff --git a/lnwire/writer_test.go b/lnwire/writer_test.go
new file mode 100644
index 000000000..96ef66103
--- /dev/null
+++ b/lnwire/writer_test.go
@@ -0,0 +1,618 @@
+package lnwire
+
+import (
+	"bytes"
+	"encoding/base32"
+	"image/color"
+	"math"
+	"net"
+	"testing"
+
+	"github.com/btcsuite/btcd/chaincfg/chainhash"
+	"github.com/btcsuite/btcd/wire"
+	"github.com/btcsuite/btcutil"
+	"github.com/lightningnetwork/lnd/tor"
+	"github.com/stretchr/testify/require"
+)
+
+func TestWriteBytes(t *testing.T) {
+	buf := new(bytes.Buffer)
+	data := []byte{1, 2, 3}
+
+	err := WriteBytes(buf, data)
+
+	require.NoError(t, err)
+	require.Equal(t, data, buf.Bytes())
+}
+
+func TestWriteUint8(t *testing.T) {
+	buf := new(bytes.Buffer)
+	data := uint8(1)
+	expectedBytes := []byte{1}
+
+	err := WriteUint8(buf, data)
+
+	require.NoError(t, err)
+	require.Equal(t, expectedBytes, buf.Bytes())
+}
+
+func TestWriteUint16(t *testing.T) {
+	buf := new(bytes.Buffer)
+	data := uint16(1)
+	expectedBytes := []byte{0, 1}
+
+	err := WriteUint16(buf, data)
+
+	require.NoError(t, err)
+	require.Equal(t, expectedBytes, buf.Bytes())
+}
+
+func TestWriteUint32(t *testing.T) {
+	buf := new(bytes.Buffer)
+	data := uint32(1)
+	expectedBytes := []byte{0, 0, 0, 1}
+
+	err := WriteUint32(buf, data)
+
+	require.NoError(t, err)
+	require.Equal(t, expectedBytes, buf.Bytes())
+}
+
+func TestWriteUint64(t *testing.T) {
+	buf := new(bytes.Buffer)
+	data := uint64(1)
+	expectedBytes := []byte{0, 0, 0, 0, 0, 0, 0, 1}
+
+	err := WriteUint64(buf, data)
+
+	require.NoError(t, err)
+	require.Equal(t, expectedBytes, buf.Bytes())
+}
+
+func TestWriteSatoshi(t *testing.T) {
+	buf := new(bytes.Buffer)
+	data := btcutil.Amount(1)
+	expectedBytes := []byte{0, 0, 0, 0, 0, 0, 0, 1}
+
+	err := WriteSatoshi(buf, data)
+
+	require.NoError(t, err)
+	require.Equal(t, expectedBytes, buf.Bytes())
+}
+
+func TestWriteMilliSatoshi(t *testing.T) {
+	buf := new(bytes.Buffer)
+	data := MilliSatoshi(1)
+	expectedBytes := []byte{0, 0, 0, 0, 0, 0, 0, 1}
+
+	err := WriteMilliSatoshi(buf, data)
+
+	require.NoError(t, err)
+	require.Equal(t, expectedBytes, buf.Bytes())
+}
+
+func TestWritePublicKey(t *testing.T) {
+	buf := new(bytes.Buffer)
+
+	// Check that when nil pubkey is used, an error will return.
+	err := WritePublicKey(buf, nil)
+	require.Equal(t, ErrNilPublicKey, err)
+
+	pub, err := randPubKey()
+	require.NoError(t, err)
+	expectedBytes := pub.SerializeCompressed()
+
+	err = WritePublicKey(buf, pub)
+
+	require.NoError(t, err)
+	require.Equal(t, expectedBytes, buf.Bytes())
+}
+
+func TestWriteChannelID(t *testing.T) {
+	buf := new(bytes.Buffer)
+	data := ChannelID{1}
+	expectedBytes := [32]byte{1}
+
+	err := WriteChannelID(buf, data)
+
+	require.NoError(t, err)
+	require.Equal(t, expectedBytes[:], buf.Bytes())
+}
+
+func TestWriteNodeAlias(t *testing.T) {
+	buf := new(bytes.Buffer)
+	data := NodeAlias{1}
+	expectedBytes := [32]byte{1}
+
+	err := WriteNodeAlias(buf, data)
+
+	require.NoError(t, err)
+	require.Equal(t, expectedBytes[:], buf.Bytes())
+}
+
+func TestWriteShortChannelID(t *testing.T) {
+	buf := new(bytes.Buffer)
+	data := ShortChannelID{BlockHeight: 1, TxIndex: 2, TxPosition: 3}
+	expectedBytes := []byte{
+		0, 0, 1, // First three bytes encodes BlockHeight.
+		0, 0, 2, // Second three bytes encodes TxIndex.
+		0, 3, // Final two bytes encodes TxPosition.
+	}
+
+	err := WriteShortChannelID(buf, data)
+
+	require.NoError(t, err)
+	require.Equal(t, expectedBytes, buf.Bytes())
+}
+
+func TestWriteSig(t *testing.T) {
+	buf := new(bytes.Buffer)
+	data := Sig{1, 2, 3}
+	expectedBytes := [64]byte{1, 2, 3}
+
+	err := WriteSig(buf, data)
+
+	require.NoError(t, err)
+	require.Equal(t, expectedBytes[:], buf.Bytes())
+}
+
+func TestWriteSigs(t *testing.T) {
+	buf := new(bytes.Buffer)
+	sig1, sig2, sig3 := Sig{1}, Sig{2}, Sig{3}
+	data := []Sig{sig1, sig2, sig3}
+
+	// First two bytes encode the length of the slice.
+	expectedBytes := []byte{0, 3}
+	expectedBytes = append(expectedBytes, sig1[:]...)
+	expectedBytes = append(expectedBytes, sig2[:]...)
+	expectedBytes = append(expectedBytes, sig3[:]...)
+
+	err := WriteSigs(buf, data)
+
+	require.NoError(t, err)
+	require.Equal(t, expectedBytes, buf.Bytes())
+}
+
+func TestWriteFailCode(t *testing.T) {
+	buf := new(bytes.Buffer)
+	data := FailCode(1)
+	expectedBytes := []byte{0, 1}
+
+	err := WriteFailCode(buf, data)
+
+	require.NoError(t, err)
+	require.Equal(t, expectedBytes, buf.Bytes())
+}
+
+// TODO(yy): expand the test to cover more encoding scenarios.
+func TestWriteRawFeatureVector(t *testing.T) {
+	buf := new(bytes.Buffer)
+
+	// Check that when nil feature is used, an error will return.
+	err := WriteRawFeatureVector(buf, nil)
+	require.Equal(t, ErrNilFeatureVector, err)
+
+	// Create a raw feature vector.
+	feature := &RawFeatureVector{features: map[FeatureBit]bool{
+		InitialRoutingSync: true, // FeatureBit 3.
+	}}
+	expectedBytes := []byte{
+		0, 1, // First two bytes encode the length.
+		8, // Last byte encodes the feature bit (1 << 3).
+	}
+
+	err = WriteRawFeatureVector(buf, feature)
+	require.NoError(t, err)
+	require.Equal(t, expectedBytes, buf.Bytes())
+}
+
+func TestWriteColorRGBA(t *testing.T) {
+	buf := new(bytes.Buffer)
+	data := color.RGBA{R: 1, G: 2, B: 3}
+	expectedBytes := []byte{1, 2, 3}
+
+	err := WriteColorRGBA(buf, data)
+
+	require.NoError(t, err)
+	require.Equal(t, expectedBytes, buf.Bytes())
+}
+
+func TestWriteShortChanIDEncoding(t *testing.T) {
+	buf := new(bytes.Buffer)
+	data := ShortChanIDEncoding(1)
+	expectedBytes := []byte{1}
+
+	err := WriteShortChanIDEncoding(buf, data)
+
+	require.NoError(t, err)
+	require.Equal(t, expectedBytes, buf.Bytes())
+}
+
+func TestWriteFundingFlag(t *testing.T) {
+	buf := new(bytes.Buffer)
+	data := FundingFlag(1)
+	expectedBytes := []byte{1}
+
+	err := WriteFundingFlag(buf, data)
+
+	require.NoError(t, err)
+	require.Equal(t, expectedBytes, buf.Bytes())
+}
+
+func TestWriteChanUpdateMsgFlags(t *testing.T) {
+	buf := new(bytes.Buffer)
+	data := ChanUpdateMsgFlags(1)
+	expectedBytes := []byte{1}
+
+	err := WriteChanUpdateMsgFlags(buf, data)
+
+	require.NoError(t, err)
+	require.Equal(t, expectedBytes, buf.Bytes())
+}
+
+func TestWriteChanUpdateChanFlags(t *testing.T) {
+	buf := new(bytes.Buffer)
+	data := ChanUpdateChanFlags(1)
+	expectedBytes := []byte{1}
+
+	err := WriteChanUpdateChanFlags(buf, data)
+
+	require.NoError(t, err)
+	require.Equal(t, expectedBytes, buf.Bytes())
+}
+
+func TestWriteDeliveryAddress(t *testing.T) {
+	buf := new(bytes.Buffer)
+	data := DeliveryAddress{1, 1, 1}
+	expectedBytes := []byte{
+		0, 3, // First two bytes encode the length.
+		1, 1, 1, // The actual data.
+	}
+
+	err := WriteDeliveryAddress(buf, data)
+
+	require.NoError(t, err)
+	require.Equal(t, expectedBytes, buf.Bytes())
+}
+
+func TestWritePingPayload(t *testing.T) {
+	buf := new(bytes.Buffer)
+	data := PingPayload{1, 1, 1}
+	expectedBytes := []byte{
+		0, 3, // First two bytes encode the length.
+		1, 1, 1, // The actual data.
+	}
+
+	err := WritePingPayload(buf, data)
+
+	require.NoError(t, err)
+	require.Equal(t, expectedBytes, buf.Bytes())
+}
+
+func TestWritePongPayload(t *testing.T) {
+	buf := new(bytes.Buffer)
+	data := PongPayload{1, 1, 1}
+	expectedBytes := []byte{
+		0, 3, // First two bytes encode the length.
+		1, 1, 1, // The actual data.
+	}
+
+	err := WritePongPayload(buf, data)
+
+	require.NoError(t, err)
+	require.Equal(t, expectedBytes, buf.Bytes())
+}
+
+func TestWriteErrorData(t *testing.T) {
+	buf := new(bytes.Buffer)
+	data := ErrorData{1, 1, 1}
+	expectedBytes := []byte{
+		0, 3, // First two bytes encode the length.
+		1, 1, 1, // The actual data.
+	}
+
+	err := WriteErrorData(buf, data)
+
+	require.NoError(t, err)
+	require.Equal(t, expectedBytes, buf.Bytes())
+}
+
+func TestWriteOpaqueReason(t *testing.T) {
+	buf := new(bytes.Buffer)
+	data := OpaqueReason{1, 1, 1}
+	expectedBytes := []byte{
+		0, 3, // First two bytes encode the length.
+		1, 1, 1, // The actual data.
+	}
+
+	err := WriteOpaqueReason(buf, data)
+
+	require.NoError(t, err)
+	require.Equal(t, expectedBytes, buf.Bytes())
+}
+
+func TestWriteBool(t *testing.T) {
+	buf := new(bytes.Buffer)
+
+	// Test write true.
+	data := true
+	expectedBytes := []byte{1}
+
+	err := WriteBool(buf, data)
+
+	require.NoError(t, err)
+	require.Equal(t, expectedBytes, buf.Bytes())
+
+	// Test write false.
+	data = false
+	expectedBytes = append(expectedBytes, 0)
+
+	err = WriteBool(buf, data)
+
+	require.NoError(t, err)
+	require.Equal(t, expectedBytes, buf.Bytes())
+}
+
+func TestWritePkScript(t *testing.T) {
+	buf := new(bytes.Buffer)
+
+	// Write a very long script to check the error is returned as expected.
+	script := PkScript{}
+	zeros := [35]byte{}
+	script = append(script, zeros[:]...)
+	err := WritePkScript(buf, script)
+	require.Equal(t, ErrPkScriptTooLong, err)
+
+	data := PkScript{1, 1, 1}
+	expectedBytes := []byte{
+		3,       // First byte encodes the length.
+		1, 1, 1, // The actual data.
+	}
+
+	err = WritePkScript(buf, data)
+
+	require.NoError(t, err)
+	require.Equal(t, expectedBytes, buf.Bytes())
+}
+
+func TestWriteOutPoint(t *testing.T) {
+	buf := new(bytes.Buffer)
+
+	// Create an outpoint with very large index to check the error is
+	// returned as expected.
+	outpointWrong := wire.OutPoint{Index: math.MaxUint16 + 1}
+	err := WriteOutPoint(buf, outpointWrong)
+	require.Equal(t, ErrOutpointIndexTooBig(outpointWrong.Index), err)
+
+	// Now check the normal write succeeds.
+	hash := chainhash.Hash{1}
+	data := wire.OutPoint{Index: 2, Hash: hash}
+	expectedBytes := []byte{}
+	expectedBytes = append(expectedBytes, hash[:]...)
+	expectedBytes = append(expectedBytes, []byte{0, 2}...)
+
+	err = WriteOutPoint(buf, data)
+
+	require.NoError(t, err)
+	require.Equal(t, expectedBytes, buf.Bytes())
+}
+
+func TestWriteTCPAddr(t *testing.T) {
+	buf := new(bytes.Buffer)
+
+	testCases := []struct {
+		name string
+		addr *net.TCPAddr
+
+		expectedErr   error
+		expectedBytes []byte
+	}{
+		{
+			// Check that the error is returned when nil address is
+			// used.
+			name:          "nil address err",
+			addr:          nil,
+			expectedErr:   ErrNilTCPAddress,
+			expectedBytes: nil,
+		},
+		{
+			// Check write IPv4.
+			name: "write ipv4",
+			addr: &net.TCPAddr{
+				IP:   net.IP{127, 0, 0, 1},
+				Port: 8080,
+			},
+			expectedErr: nil,
+			expectedBytes: []byte{
+				0x1,                 // The addressType.
+				0x7f, 0x0, 0x0, 0x1, // The IP.
+				0x1f, 0x90, // The port (31 * 256 + 144).
+			},
+		},
+		{
+			// Check write IPv6.
+			name: "write ipv6",
+			addr: &net.TCPAddr{
+				IP: net.IP{
+					1, 1, 1, 1, 1, 1, 1, 1,
+					1, 1, 1, 1, 1, 1, 1, 1,
+				},
+				Port: 8080,
+			},
+			expectedErr: nil,
+			expectedBytes: []byte{
+				0x2, // The addressType.
+				// The IP.
+				0x1, 0x1, 0x1, 0x1, 0x1, 0x1, 0x1, 0x1,
+				0x1, 0x1, 0x1, 0x1, 0x1, 0x1, 0x1, 0x1,
+				0x1f, 0x90, // The port (31 * 256 + 144).
+			},
+		},
+	}
+
+	for _, tc := range testCases {
+		tc := tc
+		t.Run(tc.name, func(t *testing.T) {
+			oldLen := buf.Len()
+
+			err := WriteTCPAddr(buf, tc.addr)
+			require.Equal(t, tc.expectedErr, err)
+
+			bytesWritten := buf.Bytes()[oldLen:buf.Len()]
+			require.Equal(t, tc.expectedBytes, bytesWritten)
+		})
+	}
+}
+
+func TestWriteOnionAddr(t *testing.T) {
+	buf := new(bytes.Buffer)
+
+	testCases := []struct {
+		name string
+		addr *tor.OnionAddr
+
+		expectedErr   error
+		expectedBytes []byte
+	}{
+		{
+			// Check that the error is returned when nil address is
+			// used.
+			name:          "nil address err",
+			addr:          nil,
+			expectedErr:   ErrNilOnionAddress,
+			expectedBytes: nil,
+		},
+		{
+			// Check the error is returned when an invalid onion
+			// address is used.
+			name:          "wrong onion service length",
+			addr:          &tor.OnionAddr{OnionService: "wrong"},
+			expectedErr:   ErrUnknownServiceLength,
+			expectedBytes: nil,
+		},
+		{
+			// Check when the address has invalid base32 encoding,
+			// the error is returned.
+			name: "invalid base32 encoding",
+			addr: &tor.OnionAddr{
+				OnionService: "1234567890123456.onion",
+			},
+			expectedErr:   base32.CorruptInputError(0),
+			expectedBytes: nil,
+		},
+		{
+			// Check write onion v2.
+			name: "onion address v2",
+			addr: &tor.OnionAddr{
+				OnionService: "abcdefghijklmnop.onion",
+				Port:         9065,
+			},
+			expectedErr: nil,
+			expectedBytes: []byte{
+				0x3,                         // The descriptor.
+				0x0, 0x44, 0x32, 0x14, 0xc7, // The host.
+				0x42, 0x54, 0xb6, 0x35, 0xcf,
+				0x23, 0x69, // The port.
+			},
+		},
+		{
+			// Check write onion v3.
+			name: "onion address v3",
+			addr: &tor.OnionAddr{
+				OnionService: "abcdefghij" +
+					"abcdefghijabcdefghij" +
+					"abcdefghijabcdefghij" +
+					"234567.onion",
+				Port: 9065,
+			},
+			expectedErr: nil,
+			expectedBytes: []byte{
+				0x4, // The descriptor.
+				0x0, 0x44, 0x32, 0x14, 0xc7, 0x42, 0x40,
+				0x11, 0xc, 0x85, 0x31, 0xd0, 0x90, 0x4,
+				0x43, 0x21, 0x4c, 0x74, 0x24, 0x1, 0x10,
+				0xc8, 0x53, 0x1d, 0x9, 0x0, 0x44, 0x32,
+				0x14, 0xc7, 0x42, 0x75, 0xbe, 0x77, 0xdf,
+				0x23, 0x69, // The port.
+			},
+		},
+	}
+
+	for _, tc := range testCases {
+		tc := tc
+		t.Run(tc.name, func(t *testing.T) {
+			oldLen := buf.Len()
+
+			err := WriteOnionAddr(buf, tc.addr)
+			require.Equal(t, tc.expectedErr, err)
+
+			bytesWritten := buf.Bytes()[oldLen:buf.Len()]
+			require.Equal(t, tc.expectedBytes, bytesWritten)
+		})
+	}
+}
+
+func TestWriteNetAddrs(t *testing.T) {
+	buf := new(bytes.Buffer)
+	tcpAddr := &net.TCPAddr{
+		IP:   net.IP{127, 0, 0, 1},
+		Port: 8080,
+	}
+	onionAddr := &tor.OnionAddr{
+		OnionService: "abcdefghijklmnop.onion",
+		Port:         9065,
+	}
+
+	testCases := []struct {
+		name string
+		addr []net.Addr
+
+		expectedErr   error
+		expectedBytes []byte
+	}{
+		{
+			// Check that the error is returned when nil address is
+			// used.
+			name:          "nil address err",
+			addr:          []net.Addr{nil, tcpAddr, onionAddr},
+			expectedErr:   ErrNilNetAddress,
+			expectedBytes: nil,
+		},
+		{
+			// Check empty address slice.
+			name:        "empty address slice",
+			addr:        []net.Addr{},
+			expectedErr: nil,
+			// Use two bytes to encode the address size.
+			expectedBytes: []byte{0, 0},
+		},
+		{
+			// Check a successful writes of a slice of addresses.
+			name:        "two addresses",
+			addr:        []net.Addr{tcpAddr, onionAddr},
+			expectedErr: nil,
+			expectedBytes: []byte{
+				// 7 bytes for TCP and 13 bytes for onion.
+				0x0, 0x14,
+				// TCP address.
+				0x1, 0x7f, 0x0, 0x0, 0x1, 0x1f, 0x90,
+				// Onion address.
+				0x3, 0x0, 0x44, 0x32, 0x14, 0xc7, 0x42,
+				0x54, 0xb6, 0x35, 0xcf, 0x23, 0x69,
+			},
+		},
+	}
+
+	for _, tc := range testCases {
+		tc := tc
+		t.Run(tc.name, func(t *testing.T) {
+			oldLen := buf.Len()
+
+			err := WriteNetAddrs(buf, tc.addr)
+			require.Equal(t, tc.expectedErr, err)
+
+			bytesWritten := buf.Bytes()[oldLen:buf.Len()]
+			require.Equal(t, tc.expectedBytes, bytesWritten)
+		})
+	}
+}

From 563ff7266a57026af3bec82b17e6b4c4870da5cf Mon Sep 17 00:00:00 2001
From: yyforyongyu <yy2452@columbia.edu>
Date: Fri, 18 Jun 2021 15:03:43 +0800
Subject: [PATCH 08/11] lnwire: refactor Encode to use specific writers - I

This commit takes 10 types of messages and refactors their Encode method
to use specific writers. The following commits will refactor the rest.
---
 lnwire/accept_channel.go          | 74 ++++++++++++++++++++++++-------
 lnwire/announcement_signatures.go | 24 +++++++---
 lnwire/closing_signed.go          | 16 +++++--
 lnwire/commit_sig.go              | 19 +++++---
 lnwire/error.go                   |  9 ++--
 lnwire/funding_created.go         | 17 +++++--
 lnwire/funding_locked.go          | 14 +++---
 lnwire/funding_signed.go          | 10 ++++-
 lnwire/gossip_timestamp_range.go  | 19 +++++---
 lnwire/init_message.go            | 14 +++---
 10 files changed, 158 insertions(+), 58 deletions(-)

diff --git a/lnwire/accept_channel.go b/lnwire/accept_channel.go
index f356e8943..f5786b918 100644
--- a/lnwire/accept_channel.go
+++ b/lnwire/accept_channel.go
@@ -126,23 +126,63 @@ func (a *AcceptChannel) Encode(w *bytes.Buffer, pver uint32) error {
 		return err
 	}
 
-	return WriteElements(w,
-		a.PendingChannelID[:],
-		a.DustLimit,
-		a.MaxValueInFlight,
-		a.ChannelReserve,
-		a.HtlcMinimum,
-		a.MinAcceptDepth,
-		a.CsvDelay,
-		a.MaxAcceptedHTLCs,
-		a.FundingKey,
-		a.RevocationPoint,
-		a.PaymentPoint,
-		a.DelayedPaymentPoint,
-		a.HtlcPoint,
-		a.FirstCommitmentPoint,
-		tlvRecords,
-	)
+	if err := WriteBytes(w, a.PendingChannelID[:]); err != nil {
+		return err
+	}
+
+	if err := WriteSatoshi(w, a.DustLimit); err != nil {
+		return err
+	}
+
+	if err := WriteMilliSatoshi(w, a.MaxValueInFlight); err != nil {
+		return err
+	}
+
+	if err := WriteSatoshi(w, a.ChannelReserve); err != nil {
+		return err
+	}
+
+	if err := WriteMilliSatoshi(w, a.HtlcMinimum); err != nil {
+		return err
+	}
+
+	if err := WriteUint32(w, a.MinAcceptDepth); err != nil {
+		return err
+	}
+
+	if err := WriteUint16(w, a.CsvDelay); err != nil {
+		return err
+	}
+
+	if err := WriteUint16(w, a.MaxAcceptedHTLCs); err != nil {
+		return err
+	}
+
+	if err := WritePublicKey(w, a.FundingKey); err != nil {
+		return err
+	}
+
+	if err := WritePublicKey(w, a.RevocationPoint); err != nil {
+		return err
+	}
+
+	if err := WritePublicKey(w, a.PaymentPoint); err != nil {
+		return err
+	}
+
+	if err := WritePublicKey(w, a.DelayedPaymentPoint); err != nil {
+		return err
+	}
+
+	if err := WritePublicKey(w, a.HtlcPoint); err != nil {
+		return err
+	}
+
+	if err := WritePublicKey(w, a.FirstCommitmentPoint); err != nil {
+		return err
+	}
+
+	return WriteBytes(w, tlvRecords)
 }
 
 // Decode deserializes the serialized AcceptChannel stored in the passed
diff --git a/lnwire/announcement_signatures.go b/lnwire/announcement_signatures.go
index 2550d1c1a..49b610611 100644
--- a/lnwire/announcement_signatures.go
+++ b/lnwire/announcement_signatures.go
@@ -66,13 +66,23 @@ func (a *AnnounceSignatures) Decode(r io.Reader, pver uint32) error {
 //
 // This is part of the lnwire.Message interface.
 func (a *AnnounceSignatures) Encode(w *bytes.Buffer, pver uint32) error {
-	return WriteElements(w,
-		a.ChannelID,
-		a.ShortChannelID,
-		a.NodeSignature,
-		a.BitcoinSignature,
-		a.ExtraOpaqueData,
-	)
+	if err := WriteChannelID(w, a.ChannelID); err != nil {
+		return err
+	}
+
+	if err := WriteShortChannelID(w, a.ShortChannelID); err != nil {
+		return err
+	}
+
+	if err := WriteSig(w, a.NodeSignature); err != nil {
+		return err
+	}
+
+	if err := WriteSig(w, a.BitcoinSignature); err != nil {
+		return err
+	}
+
+	return WriteBytes(w, a.ExtraOpaqueData)
 }
 
 // MsgType returns the integer uniquely identifying this message type on the
diff --git a/lnwire/closing_signed.go b/lnwire/closing_signed.go
index 66ccb7bcb..9ba170ba5 100644
--- a/lnwire/closing_signed.go
+++ b/lnwire/closing_signed.go
@@ -65,9 +65,19 @@ func (c *ClosingSigned) Decode(r io.Reader, pver uint32) error {
 //
 // This is part of the lnwire.Message interface.
 func (c *ClosingSigned) Encode(w *bytes.Buffer, pver uint32) error {
-	return WriteElements(
-		w, c.ChannelID, c.FeeSatoshis, c.Signature, c.ExtraData,
-	)
+	if err := WriteChannelID(w, c.ChannelID); err != nil {
+		return err
+	}
+
+	if err := WriteSatoshi(w, c.FeeSatoshis); err != nil {
+		return err
+	}
+
+	if err := WriteSig(w, c.Signature); err != nil {
+		return err
+	}
+
+	return WriteBytes(w, c.ExtraData)
 }
 
 // MsgType returns the integer uniquely identifying this message type on the
diff --git a/lnwire/commit_sig.go b/lnwire/commit_sig.go
index 59a785a0d..ca105f71a 100644
--- a/lnwire/commit_sig.go
+++ b/lnwire/commit_sig.go
@@ -71,12 +71,19 @@ func (c *CommitSig) Decode(r io.Reader, pver uint32) error {
 //
 // This is part of the lnwire.Message interface.
 func (c *CommitSig) Encode(w *bytes.Buffer, pver uint32) error {
-	return WriteElements(w,
-		c.ChanID,
-		c.CommitSig,
-		c.HtlcSigs,
-		c.ExtraData,
-	)
+	if err := WriteChannelID(w, c.ChanID); err != nil {
+		return err
+	}
+
+	if err := WriteSig(w, c.CommitSig); err != nil {
+		return err
+	}
+
+	if err := WriteSigs(w, c.HtlcSigs); err != nil {
+		return err
+	}
+
+	return WriteBytes(w, c.ExtraData)
 }
 
 // MsgType returns the integer uniquely identifying this message type on the
diff --git a/lnwire/error.go b/lnwire/error.go
index 0b95607f2..74814b59e 100644
--- a/lnwire/error.go
+++ b/lnwire/error.go
@@ -105,10 +105,11 @@ func (c *Error) Decode(r io.Reader, pver uint32) error {
 //
 // This is part of the lnwire.Message interface.
 func (c *Error) Encode(w *bytes.Buffer, pver uint32) error {
-	return WriteElements(w,
-		c.ChanID,
-		c.Data,
-	)
+	if err := WriteBytes(w, c.ChanID[:]); err != nil {
+		return err
+	}
+
+	return WriteErrorData(w, c.Data)
 }
 
 // MsgType returns the integer uniquely identifying an Error message on the
diff --git a/lnwire/funding_created.go b/lnwire/funding_created.go
index 6ca8c1b01..02b513471 100644
--- a/lnwire/funding_created.go
+++ b/lnwire/funding_created.go
@@ -42,10 +42,19 @@ var _ Message = (*FundingCreated)(nil)
 //
 // This is part of the lnwire.Message interface.
 func (f *FundingCreated) Encode(w *bytes.Buffer, pver uint32) error {
-	return WriteElements(
-		w, f.PendingChannelID[:], f.FundingPoint, f.CommitSig,
-		f.ExtraData,
-	)
+	if err := WriteBytes(w, f.PendingChannelID[:]); err != nil {
+		return err
+	}
+
+	if err := WriteOutPoint(w, f.FundingPoint); err != nil {
+		return err
+	}
+
+	if err := WriteSig(w, f.CommitSig); err != nil {
+		return err
+	}
+
+	return WriteBytes(w, f.ExtraData)
 }
 
 // Decode deserializes the serialized FundingCreated stored in the passed
diff --git a/lnwire/funding_locked.go b/lnwire/funding_locked.go
index 7c3b2eb3a..bdb0a9b06 100644
--- a/lnwire/funding_locked.go
+++ b/lnwire/funding_locked.go
@@ -60,11 +60,15 @@ func (c *FundingLocked) Decode(r io.Reader, pver uint32) error {
 //
 // This is part of the lnwire.Message interface.
 func (c *FundingLocked) Encode(w *bytes.Buffer, pver uint32) error {
-	return WriteElements(w,
-		c.ChanID,
-		c.NextPerCommitmentPoint,
-		c.ExtraData,
-	)
+	if err := WriteChannelID(w, c.ChanID); err != nil {
+		return err
+	}
+
+	if err := WritePublicKey(w, c.NextPerCommitmentPoint); err != nil {
+		return err
+	}
+
+	return WriteBytes(w, c.ExtraData)
 }
 
 // MsgType returns the uint32 code which uniquely identifies this message as a
diff --git a/lnwire/funding_signed.go b/lnwire/funding_signed.go
index d45fbd40e..d7386f2ed 100644
--- a/lnwire/funding_signed.go
+++ b/lnwire/funding_signed.go
@@ -33,7 +33,15 @@ var _ Message = (*FundingSigned)(nil)
 //
 // This is part of the lnwire.Message interface.
 func (f *FundingSigned) Encode(w *bytes.Buffer, pver uint32) error {
-	return WriteElements(w, f.ChanID, f.CommitSig, f.ExtraData)
+	if err := WriteChannelID(w, f.ChanID); err != nil {
+		return err
+	}
+
+	if err := WriteSig(w, f.CommitSig); err != nil {
+		return err
+	}
+
+	return WriteBytes(w, f.ExtraData)
 }
 
 // Decode deserializes the serialized FundingSigned stored in the passed
diff --git a/lnwire/gossip_timestamp_range.go b/lnwire/gossip_timestamp_range.go
index e9665d157..d2dbddecb 100644
--- a/lnwire/gossip_timestamp_range.go
+++ b/lnwire/gossip_timestamp_range.go
@@ -59,12 +59,19 @@ func (g *GossipTimestampRange) Decode(r io.Reader, pver uint32) error {
 //
 // This is part of the lnwire.Message interface.
 func (g *GossipTimestampRange) Encode(w *bytes.Buffer, pver uint32) error {
-	return WriteElements(w,
-		g.ChainHash[:],
-		g.FirstTimestamp,
-		g.TimestampRange,
-		g.ExtraData,
-	)
+	if err := WriteBytes(w, g.ChainHash[:]); err != nil {
+		return err
+	}
+
+	if err := WriteUint32(w, g.FirstTimestamp); err != nil {
+		return err
+	}
+
+	if err := WriteUint32(w, g.TimestampRange); err != nil {
+		return err
+	}
+
+	return WriteBytes(w, g.ExtraData)
 }
 
 // MsgType returns the integer uniquely identifying this message type on the
diff --git a/lnwire/init_message.go b/lnwire/init_message.go
index be25642e1..dbbddea2b 100644
--- a/lnwire/init_message.go
+++ b/lnwire/init_message.go
@@ -60,11 +60,15 @@ func (msg *Init) Decode(r io.Reader, pver uint32) error {
 //
 // This is part of the lnwire.Message interface.
 func (msg *Init) Encode(w *bytes.Buffer, pver uint32) error {
-	return WriteElements(w,
-		msg.GlobalFeatures,
-		msg.Features,
-		msg.ExtraData,
-	)
+	if err := WriteRawFeatureVector(w, msg.GlobalFeatures); err != nil {
+		return err
+	}
+
+	if err := WriteRawFeatureVector(w, msg.Features); err != nil {
+		return err
+	}
+
+	return WriteBytes(w, msg.ExtraData)
 }
 
 // MsgType returns the integer uniquely identifying this message type on the

From c1ad9cc60fda589ad440c9bdfdad06c846431c81 Mon Sep 17 00:00:00 2001
From: yyforyongyu <yy2452@columbia.edu>
Date: Fri, 18 Jun 2021 15:11:43 +0800
Subject: [PATCH 09/11] lnwire: refactor Encode to use specific writers - II

This commit takes another 10 message types and refactors their Encode
method to use specific writers. The following commit will refactor the
rest.
---
 lnwire/node_announcement.go        | 81 ++++++++++++++++++-------
 lnwire/onion_error.go              | 48 +++++++++------
 lnwire/onion_error_test.go         |  2 +-
 lnwire/open_channel.go             | 95 +++++++++++++++++++++++-------
 lnwire/ping.go                     |  8 ++-
 lnwire/pong.go                     |  4 +-
 lnwire/query_channel_range.go      | 19 ++++--
 lnwire/reply_short_chan_ids_end.go | 14 +++--
 lnwire/revoke_and_ack.go           | 19 ++++--
 lnwire/shutdown.go                 | 10 +++-
 10 files changed, 213 insertions(+), 87 deletions(-)

diff --git a/lnwire/node_announcement.go b/lnwire/node_announcement.go
index 664bda84d..4f1620281 100644
--- a/lnwire/node_announcement.go
+++ b/lnwire/node_announcement.go
@@ -124,17 +124,37 @@ func (a *NodeAnnouncement) Decode(r io.Reader, pver uint32) error {
 // Encode serializes the target NodeAnnouncement into the passed io.Writer
 // observing the protocol version specified.
 //
+// This is part of the lnwire.Message interface.
 func (a *NodeAnnouncement) Encode(w *bytes.Buffer, pver uint32) error {
-	return WriteElements(w,
-		a.Signature,
-		a.Features,
-		a.Timestamp,
-		a.NodeID,
-		a.RGBColor,
-		a.Alias,
-		a.Addresses,
-		a.ExtraOpaqueData,
-	)
+	if err := WriteSig(w, a.Signature); err != nil {
+		return err
+	}
+
+	if err := WriteRawFeatureVector(w, a.Features); err != nil {
+		return err
+	}
+
+	if err := WriteUint32(w, a.Timestamp); err != nil {
+		return err
+	}
+
+	if err := WriteBytes(w, a.NodeID[:]); err != nil {
+		return err
+	}
+
+	if err := WriteColorRGBA(w, a.RGBColor); err != nil {
+		return err
+	}
+
+	if err := WriteNodeAlias(w, a.Alias); err != nil {
+		return err
+	}
+
+	if err := WriteNetAddrs(w, a.Addresses); err != nil {
+		return err
+	}
+
+	return WriteBytes(w, a.ExtraOpaqueData)
 }
 
 // MsgType returns the integer uniquely identifying this message type on the
@@ -149,19 +169,36 @@ func (a *NodeAnnouncement) MsgType() MessageType {
 func (a *NodeAnnouncement) DataToSign() ([]byte, error) {
 
 	// We should not include the signatures itself.
-	var w bytes.Buffer
-	err := WriteElements(&w,
-		a.Features,
-		a.Timestamp,
-		a.NodeID,
-		a.RGBColor,
-		a.Alias[:],
-		a.Addresses,
-		a.ExtraOpaqueData,
-	)
-	if err != nil {
+	buffer := make([]byte, 0, MaxMsgBody)
+	buf := bytes.NewBuffer(buffer)
+
+	if err := WriteRawFeatureVector(buf, a.Features); err != nil {
 		return nil, err
 	}
 
-	return w.Bytes(), nil
+	if err := WriteUint32(buf, a.Timestamp); err != nil {
+		return nil, err
+	}
+
+	if err := WriteBytes(buf, a.NodeID[:]); err != nil {
+		return nil, err
+	}
+
+	if err := WriteColorRGBA(buf, a.RGBColor); err != nil {
+		return nil, err
+	}
+
+	if err := WriteNodeAlias(buf, a.Alias); err != nil {
+		return nil, err
+	}
+
+	if err := WriteNetAddrs(buf, a.Addresses); err != nil {
+		return nil, err
+	}
+
+	if err := WriteBytes(buf, a.ExtraOpaqueData); err != nil {
+		return nil, err
+	}
+
+	return buf.Bytes(), nil
 }
diff --git a/lnwire/onion_error.go b/lnwire/onion_error.go
index f9c86db61..60d61184d 100644
--- a/lnwire/onion_error.go
+++ b/lnwire/onion_error.go
@@ -420,7 +420,11 @@ func (f *FailIncorrectDetails) Decode(r io.Reader, pver uint32) error {
 //
 // NOTE: Part of the Serializable interface.
 func (f *FailIncorrectDetails) Encode(w *bytes.Buffer, pver uint32) error {
-	return WriteElements(w, f.amount, f.height)
+	if err := WriteMilliSatoshi(w, f.amount); err != nil {
+		return err
+	}
+
+	return WriteUint32(w, f.height)
 }
 
 // FailFinalExpiryTooSoon is returned if the cltv_expiry is too low, the final
@@ -486,7 +490,7 @@ func (f *FailInvalidOnionVersion) Decode(r io.Reader, pver uint32) error {
 //
 // NOTE: Part of the Serializable interface.
 func (f *FailInvalidOnionVersion) Encode(w *bytes.Buffer, pver uint32) error {
-	return WriteElement(w, f.OnionSHA256[:])
+	return WriteBytes(w, f.OnionSHA256[:])
 }
 
 // FailInvalidOnionHmac is return if the onion HMAC is incorrect.
@@ -520,7 +524,7 @@ func (f *FailInvalidOnionHmac) Decode(r io.Reader, pver uint32) error {
 //
 // NOTE: Part of the Serializable interface.
 func (f *FailInvalidOnionHmac) Encode(w *bytes.Buffer, pver uint32) error {
-	return WriteElement(w, f.OnionSHA256[:])
+	return WriteBytes(w, f.OnionSHA256[:])
 }
 
 // Returns a human readable string describing the target FailureMessage.
@@ -562,7 +566,7 @@ func (f *FailInvalidOnionKey) Decode(r io.Reader, pver uint32) error {
 //
 // NOTE: Part of the Serializable interface.
 func (f *FailInvalidOnionKey) Encode(w *bytes.Buffer, pver uint32) error {
-	return WriteElement(w, f.OnionSHA256[:])
+	return WriteBytes(w, f.OnionSHA256[:])
 }
 
 // Returns a human readable string describing the target FailureMessage.
@@ -682,7 +686,7 @@ func (f *FailTemporaryChannelFailure) Encode(w *bytes.Buffer,
 		payload = bw.Bytes()
 	}
 
-	if err := WriteElement(w, uint16(len(payload))); err != nil {
+	if err := WriteUint16(w, uint16(len(payload))); err != nil {
 		return err
 	}
 
@@ -752,7 +756,7 @@ func (f *FailAmountBelowMinimum) Decode(r io.Reader, pver uint32) error {
 //
 // NOTE: Part of the Serializable interface.
 func (f *FailAmountBelowMinimum) Encode(w *bytes.Buffer, pver uint32) error {
-	if err := WriteElement(w, f.HtlcMsat); err != nil {
+	if err := WriteMilliSatoshi(w, f.HtlcMsat); err != nil {
 		return err
 	}
 
@@ -820,7 +824,7 @@ func (f *FailFeeInsufficient) Decode(r io.Reader, pver uint32) error {
 //
 // NOTE: Part of the Serializable interface.
 func (f *FailFeeInsufficient) Encode(w *bytes.Buffer, pver uint32) error {
-	if err := WriteElement(w, f.HtlcMsat); err != nil {
+	if err := WriteMilliSatoshi(w, f.HtlcMsat); err != nil {
 		return err
 	}
 
@@ -888,7 +892,7 @@ func (f *FailIncorrectCltvExpiry) Decode(r io.Reader, pver uint32) error {
 //
 // NOTE: Part of the Serializable interface.
 func (f *FailIncorrectCltvExpiry) Encode(w *bytes.Buffer, pver uint32) error {
-	if err := WriteElement(w, f.CltvExpiry); err != nil {
+	if err := WriteUint32(w, f.CltvExpiry); err != nil {
 		return err
 	}
 
@@ -1009,7 +1013,7 @@ func (f *FailChannelDisabled) Decode(r io.Reader, pver uint32) error {
 //
 // NOTE: Part of the Serializable interface.
 func (f *FailChannelDisabled) Encode(w *bytes.Buffer, pver uint32) error {
-	if err := WriteElement(w, f.Flags); err != nil {
+	if err := WriteUint16(w, f.Flags); err != nil {
 		return err
 	}
 
@@ -1061,7 +1065,7 @@ func (f *FailFinalIncorrectCltvExpiry) Decode(r io.Reader, pver uint32) error {
 func (f *FailFinalIncorrectCltvExpiry) Encode(w *bytes.Buffer,
 	pver uint32) error {
 
-	return WriteElement(w, f.CltvExpiry)
+	return WriteUint32(w, f.CltvExpiry)
 }
 
 // FailFinalIncorrectHtlcAmount is returned if the amt_to_forward is higher
@@ -1109,7 +1113,7 @@ func (f *FailFinalIncorrectHtlcAmount) Decode(r io.Reader, pver uint32) error {
 func (f *FailFinalIncorrectHtlcAmount) Encode(w *bytes.Buffer,
 	pver uint32) error {
 
-	return WriteElement(w, f.IncomingHTLCAmount)
+	return WriteMilliSatoshi(w, f.IncomingHTLCAmount)
 }
 
 // FailExpiryTooFar is returned if the CLTV expiry in the HTLC is too far in the
@@ -1189,7 +1193,7 @@ func (f *InvalidOnionPayload) Encode(w *bytes.Buffer, pver uint32) error {
 		return err
 	}
 
-	return WriteElements(w, f.Offset)
+	return WriteUint16(w, f.Offset)
 }
 
 // FailMPPTimeout is returned if the complete amount for a multi part payment
@@ -1289,12 +1293,18 @@ func EncodeFailure(w *bytes.Buffer, failure FailureMessage, pver uint32) error {
 	// messages are fixed size.
 	pad := make([]byte, FailureMessageLength-len(failureMessage))
 
-	return WriteElements(w,
-		uint16(len(failureMessage)),
-		failureMessage,
-		uint16(len(pad)),
-		pad,
-	)
+	if err := WriteUint16(w, uint16(len(failureMessage))); err != nil {
+		return err
+	}
+
+	if err := WriteBytes(w, failureMessage); err != nil {
+		return err
+	}
+	if err := WriteUint16(w, uint16(len(pad))); err != nil {
+		return err
+	}
+
+	return WriteBytes(w, pad)
 }
 
 // EncodeFailureMessage encodes just the failure message without adding a length
@@ -1422,7 +1432,7 @@ func writeOnionErrorChanUpdate(w *bytes.Buffer, chanUpdate *ChannelUpdate,
 	// Now that we know the size, we can write the length out in the main
 	// writer.
 	updateLen := b.Len()
-	if err := WriteElement(w, uint16(updateLen)); err != nil {
+	if err := WriteUint16(w, uint16(updateLen)); err != nil {
 		return err
 	}
 
diff --git a/lnwire/onion_error_test.go b/lnwire/onion_error_test.go
index 927b6ee9c..16f9d4c36 100644
--- a/lnwire/onion_error_test.go
+++ b/lnwire/onion_error_test.go
@@ -278,5 +278,5 @@ func (f *mockFailIncorrectDetailsNoHeight) Decode(r io.Reader, pver uint32) erro
 func (f *mockFailIncorrectDetailsNoHeight) Encode(w *bytes.Buffer,
 	pver uint32) error {
 
-	return WriteElement(w, f.amount)
+	return WriteUint64(w, f.amount)
 }
diff --git a/lnwire/open_channel.go b/lnwire/open_channel.go
index d975749d2..035297ba2 100644
--- a/lnwire/open_channel.go
+++ b/lnwire/open_channel.go
@@ -161,27 +161,80 @@ func (o *OpenChannel) Encode(w *bytes.Buffer, pver uint32) error {
 		return err
 	}
 
-	return WriteElements(w,
-		o.ChainHash[:],
-		o.PendingChannelID[:],
-		o.FundingAmount,
-		o.PushAmount,
-		o.DustLimit,
-		o.MaxValueInFlight,
-		o.ChannelReserve,
-		o.HtlcMinimum,
-		o.FeePerKiloWeight,
-		o.CsvDelay,
-		o.MaxAcceptedHTLCs,
-		o.FundingKey,
-		o.RevocationPoint,
-		o.PaymentPoint,
-		o.DelayedPaymentPoint,
-		o.HtlcPoint,
-		o.FirstCommitmentPoint,
-		o.ChannelFlags,
-		tlvRecords,
-	)
+	if err := WriteBytes(w, o.ChainHash[:]); err != nil {
+		return err
+	}
+
+	if err := WriteBytes(w, o.PendingChannelID[:]); err != nil {
+		return err
+	}
+
+	if err := WriteSatoshi(w, o.FundingAmount); err != nil {
+		return err
+	}
+
+	if err := WriteMilliSatoshi(w, o.PushAmount); err != nil {
+		return err
+	}
+
+	if err := WriteSatoshi(w, o.DustLimit); err != nil {
+		return err
+	}
+
+	if err := WriteMilliSatoshi(w, o.MaxValueInFlight); err != nil {
+		return err
+	}
+
+	if err := WriteSatoshi(w, o.ChannelReserve); err != nil {
+		return err
+	}
+
+	if err := WriteMilliSatoshi(w, o.HtlcMinimum); err != nil {
+		return err
+	}
+
+	if err := WriteUint32(w, o.FeePerKiloWeight); err != nil {
+		return err
+	}
+
+	if err := WriteUint16(w, o.CsvDelay); err != nil {
+		return err
+	}
+
+	if err := WriteUint16(w, o.MaxAcceptedHTLCs); err != nil {
+		return err
+	}
+
+	if err := WritePublicKey(w, o.FundingKey); err != nil {
+		return err
+	}
+
+	if err := WritePublicKey(w, o.RevocationPoint); err != nil {
+		return err
+	}
+
+	if err := WritePublicKey(w, o.PaymentPoint); err != nil {
+		return err
+	}
+
+	if err := WritePublicKey(w, o.DelayedPaymentPoint); err != nil {
+		return err
+	}
+
+	if err := WritePublicKey(w, o.HtlcPoint); err != nil {
+		return err
+
+	}
+
+	if err := WritePublicKey(w, o.FirstCommitmentPoint); err != nil {
+		return err
+	}
+
+	if err := WriteFundingFlag(w, o.ChannelFlags); err != nil {
+		return err
+	}
+
+	return WriteBytes(w, tlvRecords)
 }
 
 // Decode deserializes the serialized OpenChannel stored in the passed
diff --git a/lnwire/ping.go b/lnwire/ping.go
index 9218fe7f6..1e2877d0c 100644
--- a/lnwire/ping.go
+++ b/lnwire/ping.go
@@ -48,9 +48,11 @@ func (p *Ping) Decode(r io.Reader, pver uint32) error {
 //
 // This is part of the lnwire.Message interface.
 func (p *Ping) Encode(w *bytes.Buffer, pver uint32) error {
-	return WriteElements(w,
-		p.NumPongBytes,
-		p.PaddingBytes)
+	if err := WriteUint16(w, p.NumPongBytes); err != nil {
+		return err
+	}
+
+	return WritePingPayload(w, p.PaddingBytes)
 }
 
 // MsgType returns the integer uniquely identifying this message type on the
diff --git a/lnwire/pong.go b/lnwire/pong.go
index 6c70c120f..953eef5f5 100644
--- a/lnwire/pong.go
+++ b/lnwire/pong.go
@@ -44,9 +44,7 @@ func (p *Pong) Decode(r io.Reader, pver uint32) error {
 //
 // This is part of the lnwire.Message interface.
 func (p *Pong) Encode(w *bytes.Buffer, pver uint32) error {
-	return WriteElements(w,
-		p.PongBytes,
-	)
+	return WritePongPayload(w, p.PongBytes)
 }
 
 // MsgType returns the integer uniquely identifying this message type on the
diff --git a/lnwire/query_channel_range.go b/lnwire/query_channel_range.go
index a6a20e1fd..cfe88e9bf 100644
--- a/lnwire/query_channel_range.go
+++ b/lnwire/query_channel_range.go
@@ -60,12 +60,19 @@ func (q *QueryChannelRange) Decode(r io.Reader, pver uint32) error {
 //
 // This is part of the lnwire.Message interface.
 func (q *QueryChannelRange) Encode(w *bytes.Buffer, pver uint32) error {
-	return WriteElements(w,
-		q.ChainHash[:],
-		q.FirstBlockHeight,
-		q.NumBlocks,
-		q.ExtraData,
-	)
+	if err := WriteBytes(w, q.ChainHash[:]); err != nil {
+		return err
+	}
+
+	if err := WriteUint32(w, q.FirstBlockHeight); err != nil {
+		return err
+	}
+
+	if err := WriteUint32(w, q.NumBlocks); err != nil {
+		return err
+	}
+
+	return WriteBytes(w, q.ExtraData)
 }
 
 // MsgType returns the integer uniquely identifying this message type on the
diff --git a/lnwire/reply_short_chan_ids_end.go b/lnwire/reply_short_chan_ids_end.go
index 6d909febf..53676f719 100644
--- a/lnwire/reply_short_chan_ids_end.go
+++ b/lnwire/reply_short_chan_ids_end.go
@@ -56,11 +56,15 @@ func (c *ReplyShortChanIDsEnd) Decode(r io.Reader, pver uint32) error {
 //
 // This is part of the lnwire.Message interface.
 func (c *ReplyShortChanIDsEnd) Encode(w *bytes.Buffer, pver uint32) error {
-	return WriteElements(w,
-		c.ChainHash[:],
-		c.Complete,
-		c.ExtraData,
-	)
+	if err := WriteBytes(w, c.ChainHash[:]); err != nil {
+		return err
+	}
+
+	if err := WriteUint8(w, c.Complete); err != nil {
+		return err
+	}
+
+	return WriteBytes(w, c.ExtraData)
 }
 
 // MsgType returns the integer uniquely identifying this message type on the
diff --git a/lnwire/revoke_and_ack.go b/lnwire/revoke_and_ack.go
index acbfcae6e..0acf3ef7b 100644
--- a/lnwire/revoke_and_ack.go
+++ b/lnwire/revoke_and_ack.go
@@ -67,12 +67,19 @@ func (c *RevokeAndAck) Decode(r io.Reader, pver uint32) error {
 //
 // This is part of the lnwire.Message interface.
 func (c *RevokeAndAck) Encode(w *bytes.Buffer, pver uint32) error {
-	return WriteElements(w,
-		c.ChanID,
-		c.Revocation[:],
-		c.NextRevocationKey,
-		c.ExtraData,
-	)
+	if err := WriteChannelID(w, c.ChanID); err != nil {
+		return err
+	}
+
+	if err := WriteBytes(w, c.Revocation[:]); err != nil {
+		return err
+	}
+
+	if err := WritePublicKey(w, c.NextRevocationKey); err != nil {
+		return err
+	}
+
+	return WriteBytes(w, c.ExtraData)
 }
 
 // MsgType returns the integer uniquely identifying this message type on the
diff --git a/lnwire/shutdown.go b/lnwire/shutdown.go
index e7b26283b..2adb6a082 100644
--- a/lnwire/shutdown.go
+++ b/lnwire/shutdown.go
@@ -48,7 +48,15 @@ func (s *Shutdown) Decode(r io.Reader, pver uint32) error {
 //
 // This is part of the lnwire.Message interface.
 func (s *Shutdown) Encode(w *bytes.Buffer, pver uint32) error {
-	return WriteElements(w, s.ChannelID, s.Address, s.ExtraData)
+	if err := WriteChannelID(w, s.ChannelID); err != nil {
+		return err
+	}
+
+	if err := WriteDeliveryAddress(w, s.Address); err != nil {
+		return err
+	}
+
+	return WriteBytes(w, s.ExtraData)
 }
 
 // MsgType returns the integer uniquely identifying this message type on the

From 2cf6969dbc76fb770607b0f6e604d8a5c09d5b7b Mon Sep 17 00:00:00 2001
From: yyforyongyu <yy2452@columbia.edu>
Date: Fri, 18 Jun 2021 15:15:44 +0800
Subject: [PATCH 10/11] lnwire: refactor Encode to use specific writers - III

This commit refactors the remaining usage of WriteElements. By
replacing the interface types with concrete types for the params used in
the methods, most of the encoding of the messages now takes zero heap
allocations.
---
 lnwire/channel_announcement.go       | 105 ++++++++++++++++++-------
 lnwire/channel_reestablish.go        |  30 +++++---
 lnwire/channel_update.go             | 111 +++++++++++++++++++--------
 lnwire/extra_bytes.go                |   2 +-
 lnwire/lnwire_test.go                |   2 +-
 lnwire/query_short_chan_ids.go       |  24 +++---
 lnwire/reply_channel_range.go        |  30 +++++---
 lnwire/update_add_htlc.go            |  38 ++++++---
 lnwire/update_fail_htlc.go           |  19 +++--
 lnwire/update_fail_malformed_htlc.go |  28 +++++--
 lnwire/update_fee.go                 |  14 ++--
 lnwire/update_fulfill_htlc.go        |  19 +++--
 12 files changed, 289 insertions(+), 133 deletions(-)

diff --git a/lnwire/channel_announcement.go b/lnwire/channel_announcement.go
index 9e57d17bc..2b34c0f99 100644
--- a/lnwire/channel_announcement.go
+++ b/lnwire/channel_announcement.go
@@ -88,20 +88,51 @@ func (a *ChannelAnnouncement) Decode(r io.Reader, pver uint32) error {
 //
 // This is part of the lnwire.Message interface.
 func (a *ChannelAnnouncement) Encode(w *bytes.Buffer, pver uint32) error {
-	return WriteElements(w,
-		a.NodeSig1,
-		a.NodeSig2,
-		a.BitcoinSig1,
-		a.BitcoinSig2,
-		a.Features,
-		a.ChainHash[:],
-		a.ShortChannelID,
-		a.NodeID1,
-		a.NodeID2,
-		a.BitcoinKey1,
-		a.BitcoinKey2,
-		a.ExtraOpaqueData,
-	)
+	if err := WriteSig(w, a.NodeSig1); err != nil {
+		return err
+	}
+
+	if err := WriteSig(w, a.NodeSig2); err != nil {
+		return err
+	}
+
+	if err := WriteSig(w, a.BitcoinSig1); err != nil {
+		return err
+	}
+
+	if err := WriteSig(w, a.BitcoinSig2); err != nil {
+		return err
+	}
+
+	if err := WriteRawFeatureVector(w, a.Features); err != nil {
+		return err
+	}
+
+	if err := WriteBytes(w, a.ChainHash[:]); err != nil {
+		return err
+	}
+
+	if err := WriteShortChannelID(w, a.ShortChannelID); err != nil {
+		return err
+	}
+
+	if err := WriteBytes(w, a.NodeID1[:]); err != nil {
+		return err
+	}
+
+	if err := WriteBytes(w, a.NodeID2[:]); err != nil {
+		return err
+	}
+
+	if err := WriteBytes(w, a.BitcoinKey1[:]); err != nil {
+		return err
+	}
+
+	if err := WriteBytes(w, a.BitcoinKey2[:]); err != nil {
+		return err
+	}
+
+	return WriteBytes(w, a.ExtraOpaqueData)
 }
 
 // MsgType returns the integer uniquely identifying this message type on the
@@ -116,20 +147,40 @@ func (a *ChannelAnnouncement) MsgType() MessageType {
 // be signed.
 func (a *ChannelAnnouncement) DataToSign() ([]byte, error) {
 	// We should not include the signatures itself.
-	var w bytes.Buffer
-	err := WriteElements(&w,
-		a.Features,
-		a.ChainHash[:],
-		a.ShortChannelID,
-		a.NodeID1,
-		a.NodeID2,
-		a.BitcoinKey1,
-		a.BitcoinKey2,
-		a.ExtraOpaqueData,
-	)
-	if err != nil {
+	b := make([]byte, 0, MaxMsgBody)
+	buf := bytes.NewBuffer(b)
+
+	if err := WriteRawFeatureVector(buf, a.Features); err != nil {
 		return nil, err
 	}
 
-	return w.Bytes(), nil
+	if err := WriteBytes(buf, a.ChainHash[:]); err != nil {
+		return nil, err
+	}
+
+	if err := WriteShortChannelID(buf, a.ShortChannelID); err != nil {
+		return nil, err
+	}
+
+	if err := WriteBytes(buf, a.NodeID1[:]); err != nil {
+		return nil, err
+	}
+
+	if err := WriteBytes(buf, a.NodeID2[:]); err != nil {
+		return nil, err
+	}
+
+	if err := WriteBytes(buf, a.BitcoinKey1[:]); err != nil {
+		return nil, err
+	}
+
+	if err := WriteBytes(buf, a.BitcoinKey2[:]); err != nil {
+		return nil, err
+	}
+
+	if err := WriteBytes(buf, a.ExtraOpaqueData); err != nil {
+		return nil, err
+	}
+
+	return buf.Bytes(), nil
 }
diff --git a/lnwire/channel_reestablish.go b/lnwire/channel_reestablish.go
index 44cc6f3f6..0de16a365 100644
--- a/lnwire/channel_reestablish.go
+++ b/lnwire/channel_reestablish.go
@@ -77,12 +77,15 @@ var _ Message = (*ChannelReestablish)(nil)
 //
 // This is part of the lnwire.Message interface.
 func (a *ChannelReestablish) Encode(w *bytes.Buffer, pver uint32) error {
-	err := WriteElements(w,
-		a.ChanID,
-		a.NextLocalCommitHeight,
-		a.RemoteCommitTailHeight,
-	)
-	if err != nil {
+	if err := WriteChannelID(w, a.ChanID); err != nil {
+		return err
+	}
+
+	if err := WriteUint64(w, a.NextLocalCommitHeight); err != nil {
+		return err
+	}
+
+	if err := WriteUint64(w, a.RemoteCommitTailHeight); err != nil {
 		return err
 	}
 
@@ -94,15 +97,18 @@ func (a *ChannelReestablish) Encode(w *bytes.Buffer, pver uint32) error {
 		//
 		// NOTE: This is here primarily for the quickcheck tests, in
 		// practice, we'll always populate this field.
-		return WriteElements(w, a.ExtraData)
+		return WriteBytes(w, a.ExtraData)
 	}
 
 	// Otherwise, we'll write out the remaining elements.
-	return WriteElements(w,
-		a.LastRemoteCommitSecret[:],
-		a.LocalUnrevokedCommitPoint,
-		a.ExtraData,
-	)
+	if err := WriteBytes(w, a.LastRemoteCommitSecret[:]); err != nil {
+		return err
+	}
+
+	if err := WritePublicKey(w, a.LocalUnrevokedCommitPoint); err != nil {
+		return err
+	}
+	return WriteBytes(w, a.ExtraData)
 }
 
 // Decode deserializes a serialized ChannelReestablish stored in the passed
diff --git a/lnwire/channel_update.go b/lnwire/channel_update.go
index e1bac9f96..7881f972f 100644
--- a/lnwire/channel_update.go
+++ b/lnwire/channel_update.go
@@ -160,32 +160,57 @@ func (a *ChannelUpdate) Decode(r io.Reader, pver uint32) error {
 //
 // This is part of the lnwire.Message interface.
 func (a *ChannelUpdate) Encode(w *bytes.Buffer, pver uint32) error {
-	err := WriteElements(w,
-		a.Signature,
-		a.ChainHash[:],
-		a.ShortChannelID,
-		a.Timestamp,
-		a.MessageFlags,
-		a.ChannelFlags,
-		a.TimeLockDelta,
-		a.HtlcMinimumMsat,
-		a.BaseFee,
-		a.FeeRate,
-	)
-	if err != nil {
+	if err := WriteSig(w, a.Signature); err != nil {
+		return err
+	}
+
+	if err := WriteBytes(w, a.ChainHash[:]); err != nil {
+		return err
+	}
+
+	if err := WriteShortChannelID(w, a.ShortChannelID); err != nil {
+		return err
+	}
+
+	if err := WriteUint32(w, a.Timestamp); err != nil {
+		return err
+	}
+
+	if err := WriteChanUpdateMsgFlags(w, a.MessageFlags); err != nil {
+		return err
+	}
+
+	if err := WriteChanUpdateChanFlags(w, a.ChannelFlags); err != nil {
+		return err
+	}
+
+	if err := WriteUint16(w, a.TimeLockDelta); err != nil {
+		return err
+	}
+
+	if err := WriteMilliSatoshi(w, a.HtlcMinimumMsat); err != nil {
+		return err
+	}
+
+	if err := WriteUint32(w, a.BaseFee); err != nil {
+		return err
+	}
+
+	if err := WriteUint32(w, a.FeeRate); err != nil {
 		return err
 	}
 
 	// Now append optional fields if they are set. Currently, the only
 	// optional field is max HTLC.
 	if a.MessageFlags.HasMaxHtlc() {
-		if err := WriteElements(w, a.HtlcMaximumMsat); err != nil {
+		err := WriteMilliSatoshi(w, a.HtlcMaximumMsat)
+		if err != nil {
 			return err
 		}
 	}
 
 	// Finally, append any extra opaque data.
-	return a.ExtraOpaqueData.Encode(w)
+	return WriteBytes(w, a.ExtraOpaqueData)
 }
 
 // MsgType returns the integer uniquely identifying this message type on the
@@ -199,36 +224,58 @@ func (a *ChannelUpdate) MsgType() MessageType {
 // DataToSign is used to retrieve part of the announcement message which should
 // be signed.
 func (a *ChannelUpdate) DataToSign() ([]byte, error) {
-
 	// We should not include the signatures itself.
-	var w bytes.Buffer
-	err := WriteElements(&w,
-		a.ChainHash[:],
-		a.ShortChannelID,
-		a.Timestamp,
-		a.MessageFlags,
-		a.ChannelFlags,
-		a.TimeLockDelta,
-		a.HtlcMinimumMsat,
-		a.BaseFee,
-		a.FeeRate,
-	)
-	if err != nil {
+	b := make([]byte, 0, MaxMsgBody)
+	buf := bytes.NewBuffer(b)
+	if err := WriteBytes(buf, a.ChainHash[:]); err != nil {
+		return nil, err
+	}
+
+	if err := WriteShortChannelID(buf, a.ShortChannelID); err != nil {
+		return nil, err
+	}
+
+	if err := WriteUint32(buf, a.Timestamp); err != nil {
+		return nil, err
+	}
+
+	if err := WriteChanUpdateMsgFlags(buf, a.MessageFlags); err != nil {
+		return nil, err
+	}
+
+	if err := WriteChanUpdateChanFlags(buf, a.ChannelFlags); err != nil {
+		return nil, err
+	}
+
+	if err := WriteUint16(buf, a.TimeLockDelta); err != nil {
+		return nil, err
+	}
+
+	if err := WriteMilliSatoshi(buf, a.HtlcMinimumMsat); err != nil {
+		return nil, err
+	}
+
+	if err := WriteUint32(buf, a.BaseFee); err != nil {
+		return nil, err
+	}
+
+	if err := WriteUint32(buf, a.FeeRate); err != nil {
 		return nil, err
 	}
 
 	// Now append optional fields if they are set. Currently, the only
 	// optional field is max HTLC.
 	if a.MessageFlags.HasMaxHtlc() {
-		if err := WriteElements(&w, a.HtlcMaximumMsat); err != nil {
+		err := WriteMilliSatoshi(buf, a.HtlcMaximumMsat)
+		if err != nil {
 			return nil, err
 		}
 	}
 
 	// Finally, append any extra opaque data.
-	if err := a.ExtraOpaqueData.Encode(&w); err != nil {
+	if err := WriteBytes(buf, a.ExtraOpaqueData); err != nil {
 		return nil, err
 	}
 
-	return w.Bytes(), nil
+	return buf.Bytes(), nil
 }
diff --git a/lnwire/extra_bytes.go b/lnwire/extra_bytes.go
index f2dbec45b..70554f4f5 100644
--- a/lnwire/extra_bytes.go
+++ b/lnwire/extra_bytes.go
@@ -18,7 +18,7 @@ type ExtraOpaqueData []byte
 // Encode attempts to encode the raw extra bytes into the passed io.Writer.
 func (e *ExtraOpaqueData) Encode(w *bytes.Buffer) error {
 	eBytes := []byte((*e)[:])
-	if err := WriteElements(w, eBytes); err != nil {
+	if err := WriteBytes(w, eBytes); err != nil {
 		return err
 	}
 
diff --git a/lnwire/lnwire_test.go b/lnwire/lnwire_test.go
index 14082bfe9..4475b3382 100644
--- a/lnwire/lnwire_test.go
+++ b/lnwire/lnwire_test.go
@@ -232,7 +232,7 @@ func TestMaxOutPointIndex(t *testing.T) {
 	}
 
 	var b bytes.Buffer
-	if err := WriteElement(&b, op); err == nil {
+	if err := WriteOutPoint(&b, op); err == nil {
 		t.Fatalf("write of outPoint should fail, index exceeds 16-bits")
 	}
 }
diff --git a/lnwire/query_short_chan_ids.go b/lnwire/query_short_chan_ids.go
index 930becdbb..323a936db 100644
--- a/lnwire/query_short_chan_ids.go
+++ b/lnwire/query_short_chan_ids.go
@@ -293,19 +293,18 @@ func decodeShortChanIDs(r io.Reader) (ShortChanIDEncoding, []ShortChannelID, err
 // This is part of the lnwire.Message interface.
 func (q *QueryShortChanIDs) Encode(w *bytes.Buffer, pver uint32) error {
 	// First, we'll write out the chain hash.
-	err := WriteElements(w, q.ChainHash[:])
-	if err != nil {
+	if err := WriteBytes(w, q.ChainHash[:]); err != nil {
 		return err
 	}
 
 	// Base on our encoding type, we'll write out the set of short channel
 	// ID's.
-	err = encodeShortChanIDs(w, q.EncodingType, q.ShortChanIDs, q.noSort)
+	err := encodeShortChanIDs(w, q.EncodingType, q.ShortChanIDs, q.noSort)
 	if err != nil {
 		return err
 	}
 
-	return q.ExtraData.Encode(w)
+	return WriteBytes(w, q.ExtraData)
 }
 
 // encodeShortChanIDs encodes the passed short channel ID's into the passed
@@ -332,20 +331,21 @@ func encodeShortChanIDs(w *bytes.Buffer, encodingType ShortChanIDEncoding,
 		// body. We add 1 as the response will have the encoding type
 		// prepended to it.
 		numBytesBody := uint16(len(shortChanIDs)*8) + 1
-		if err := WriteElements(w, numBytesBody); err != nil {
+		if err := WriteUint16(w, numBytesBody); err != nil {
 			return err
 		}
 
 		// We'll then write out the encoding that that follows the
 		// actual encoded short channel ID's.
-		if err := WriteElements(w, encodingType); err != nil {
+		err := WriteShortChanIDEncoding(w, encodingType)
+		if err != nil {
 			return err
 		}
 
 		// Now that we know they're sorted, we can write out each short
 		// channel ID to the buffer.
 		for _, chanID := range shortChanIDs {
-			if err := WriteElements(w, chanID); err != nil {
+			if err := WriteShortChannelID(w, chanID); err != nil {
 				return fmt.Errorf("unable to write short chan "+
 					"ID: %v", err)
 			}
@@ -374,7 +374,7 @@ func encodeShortChanIDs(w *bytes.Buffer, encodingType ShortChanIDEncoding,
 			// into the zlib writer, which will do compressing on
 			// the fly.
 			for _, chanID := range shortChanIDs {
-				err := WriteElements(&wb, chanID)
+				err := WriteShortChannelID(&wb, chanID)
 				if err != nil {
 					return fmt.Errorf(
 						"unable to write short chan "+
@@ -418,15 +418,15 @@ func encodeShortChanIDs(w *bytes.Buffer, encodingType ShortChanIDEncoding,
 
 		// Finally, we can write out the number of bytes, the
 		// compression type, and finally the buffer itself.
-		if err := WriteElements(w, uint16(numBytesBody)); err != nil {
+		if err := WriteUint16(w, uint16(numBytesBody)); err != nil {
 			return err
 		}
-		if err := WriteElements(w, encodingType); err != nil {
+		err := WriteShortChanIDEncoding(w, encodingType)
+		if err != nil {
 			return err
 		}
 
-		_, err := w.Write(compressedPayload)
-		return err
+		return WriteBytes(w, compressedPayload)
 
 	default:
 		// If we're trying to encode with an encoding type that we
diff --git a/lnwire/reply_channel_range.go b/lnwire/reply_channel_range.go
index 00bb10eb2..9dc0fca9c 100644
--- a/lnwire/reply_channel_range.go
+++ b/lnwire/reply_channel_range.go
@@ -87,22 +87,28 @@ func (c *ReplyChannelRange) Decode(r io.Reader, pver uint32) error {
 //
 // This is part of the lnwire.Message interface.
 func (c *ReplyChannelRange) Encode(w *bytes.Buffer, pver uint32) error {
-	err := WriteElements(w,
-		c.ChainHash[:],
-		c.FirstBlockHeight,
-		c.NumBlocks,
-		c.Complete,
-	)
+	if err := WriteBytes(w, c.ChainHash[:]); err != nil {
+		return err
+	}
+
+	if err := WriteUint32(w, c.FirstBlockHeight); err != nil {
+		return err
+	}
+
+	if err := WriteUint32(w, c.NumBlocks); err != nil {
+		return err
+	}
+
+	if err := WriteUint8(w, c.Complete); err != nil {
+		return err
+	}
+
+	err := encodeShortChanIDs(w, c.EncodingType, c.ShortChanIDs, c.noSort)
 	if err != nil {
 		return err
 	}
 
-	err = encodeShortChanIDs(w, c.EncodingType, c.ShortChanIDs, c.noSort)
-	if err != nil {
-		return err
-	}
-
-	return c.ExtraData.Encode(w)
+	return WriteBytes(w, c.ExtraData)
 }
 
 // MsgType returns the integer uniquely identifying this message type on the
diff --git a/lnwire/update_add_htlc.go b/lnwire/update_add_htlc.go
index a7756994c..666a54942 100644
--- a/lnwire/update_add_htlc.go
+++ b/lnwire/update_add_htlc.go
@@ -85,20 +85,36 @@ func (c *UpdateAddHTLC) Decode(r io.Reader, pver uint32) error {
 	)
 }
 
-// Encode serializes the target UpdateAddHTLC into the passed io.Writer observing
-// the protocol version specified.
+// Encode serializes the target UpdateAddHTLC into the passed io.Writer
+// observing the protocol version specified.
 //
 // This is part of the lnwire.Message interface.
 func (c *UpdateAddHTLC) Encode(w *bytes.Buffer, pver uint32) error {
-	return WriteElements(w,
-		c.ChanID,
-		c.ID,
-		c.Amount,
-		c.PaymentHash[:],
-		c.Expiry,
-		c.OnionBlob[:],
-		c.ExtraData,
-	)
+	if err := WriteChannelID(w, c.ChanID); err != nil {
+		return err
+	}
+
+	if err := WriteUint64(w, c.ID); err != nil {
+		return err
+	}
+
+	if err := WriteMilliSatoshi(w, c.Amount); err != nil {
+		return err
+	}
+
+	if err := WriteBytes(w, c.PaymentHash[:]); err != nil {
+		return err
+	}
+
+	if err := WriteUint32(w, c.Expiry); err != nil {
+		return err
+	}
+
+	if err := WriteBytes(w, c.OnionBlob[:]); err != nil {
+		return err
+	}
+
+	return WriteBytes(w, c.ExtraData)
 }
 
 // MsgType returns the integer uniquely identifying this message type on the
diff --git a/lnwire/update_fail_htlc.go b/lnwire/update_fail_htlc.go
index 2344b5683..61f02bac9 100644
--- a/lnwire/update_fail_htlc.go
+++ b/lnwire/update_fail_htlc.go
@@ -56,12 +56,19 @@ func (c *UpdateFailHTLC) Decode(r io.Reader, pver uint32) error {
 //
 // This is part of the lnwire.Message interface.
 func (c *UpdateFailHTLC) Encode(w *bytes.Buffer, pver uint32) error {
-	return WriteElements(w,
-		c.ChanID,
-		c.ID,
-		c.Reason,
-		c.ExtraData,
-	)
+	if err := WriteChannelID(w, c.ChanID); err != nil {
+		return err
+	}
+
+	if err := WriteUint64(w, c.ID); err != nil {
+		return err
+	}
+
+	if err := WriteOpaqueReason(w, c.Reason); err != nil {
+		return err
+	}
+
+	return WriteBytes(w, c.ExtraData)
 }
 
 // MsgType returns the integer uniquely identifying this message type on the
diff --git a/lnwire/update_fail_malformed_htlc.go b/lnwire/update_fail_malformed_htlc.go
index 120f89541..f28107a9a 100644
--- a/lnwire/update_fail_malformed_htlc.go
+++ b/lnwire/update_fail_malformed_htlc.go
@@ -54,14 +54,26 @@ func (c *UpdateFailMalformedHTLC) Decode(r io.Reader, pver uint32) error {
 // io.Writer observing the protocol version specified.
 //
 // This is part of the lnwire.Message interface.
-func (c *UpdateFailMalformedHTLC) Encode(w *bytes.Buffer, pver uint32) error {
-	return WriteElements(w,
-		c.ChanID,
-		c.ID,
-		c.ShaOnionBlob[:],
-		c.FailureCode,
-		c.ExtraData,
-	)
+func (c *UpdateFailMalformedHTLC) Encode(w *bytes.Buffer,
+	pver uint32) error {
+
+	if err := WriteChannelID(w, c.ChanID); err != nil {
+		return err
+	}
+
+	if err := WriteUint64(w, c.ID); err != nil {
+		return err
+	}
+
+	if err := WriteBytes(w, c.ShaOnionBlob[:]); err != nil {
+		return err
+	}
+
+	if err := WriteFailCode(w, c.FailureCode); err != nil {
+		return err
+	}
+
+	return WriteBytes(w, c.ExtraData)
 }
 
 // MsgType returns the integer uniquely identifying this message type on the
diff --git a/lnwire/update_fee.go b/lnwire/update_fee.go
index a30634d79..a7026044c 100644
--- a/lnwire/update_fee.go
+++ b/lnwire/update_fee.go
@@ -53,11 +53,15 @@ func (c *UpdateFee) Decode(r io.Reader, pver uint32) error {
 //
 // This is part of the lnwire.Message interface.
 func (c *UpdateFee) Encode(w *bytes.Buffer, pver uint32) error {
-	return WriteElements(w,
-		c.ChanID,
-		c.FeePerKw,
-		c.ExtraData,
-	)
+	if err := WriteChannelID(w, c.ChanID); err != nil {
+		return err
+	}
+
+	if err := WriteUint32(w, c.FeePerKw); err != nil {
+		return err
+	}
+
+	return WriteBytes(w, c.ExtraData)
 }
 
 // MsgType returns the integer uniquely identifying this message type on the
diff --git a/lnwire/update_fulfill_htlc.go b/lnwire/update_fulfill_htlc.go
index 4cd3a9720..275a37c87 100644
--- a/lnwire/update_fulfill_htlc.go
+++ b/lnwire/update_fulfill_htlc.go
@@ -62,12 +62,19 @@ func (c *UpdateFulfillHTLC) Decode(r io.Reader, pver uint32) error {
 //
 // This is part of the lnwire.Message interface.
 func (c *UpdateFulfillHTLC) Encode(w *bytes.Buffer, pver uint32) error {
-	return WriteElements(w,
-		c.ChanID,
-		c.ID,
-		c.PaymentPreimage[:],
-		c.ExtraData,
-	)
+	if err := WriteChannelID(w, c.ChanID); err != nil {
+		return err
+	}
+
+	if err := WriteUint64(w, c.ID); err != nil {
+		return err
+	}
+
+	if err := WriteBytes(w, c.PaymentPreimage[:]); err != nil {
+		return err
+	}
+
+	return WriteBytes(w, c.ExtraData)
 }
 
 // MsgType returns the integer uniquely identifying this message type on the

From a0e958f96f0d4e60909db70fcecd5c2e6fc1f23a Mon Sep 17 00:00:00 2001
From: yyforyongyu <yy2452@columbia.edu>
Date: Tue, 10 Aug 2021 05:54:45 +0800
Subject: [PATCH 11/11] docs: add release notes for lnwire gc

---
 docs/release-notes/release-notes-0.14.0.md | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/docs/release-notes/release-notes-0.14.0.md b/docs/release-notes/release-notes-0.14.0.md
index 99fbc481c..aac58c21f 100644
--- a/docs/release-notes/release-notes-0.14.0.md
+++ b/docs/release-notes/release-notes-0.14.0.md
@@ -119,6 +119,10 @@ you.
 * [Update MC store in blocks](https://github.com/lightningnetwork/lnd/pull/5515)
   to make payment throughput better when using etcd.
 
+* [The `lnwire` package now uses a write buffer pool](https://github.com/lightningnetwork/lnd/pull/4884)
+  when encoding/decoding messages. Such that most of the heap escapes are fixed,
+  resulting in less memory being used when running `lnd`.
+
 ## Bug Fixes
 
 A bug has been fixed that would cause `lnd` to [try to bootstrap using the