Merge pull request #2606 from cfromknecht/wtwire-init-connection-features

watchtower/multi: send connection features + chain hash in Init

watchtower/wtwire/init.go

@@ -1,23 +1,59 @@
 package wtwire
 
-import "github.com/lightningnetwork/lnd/lnwire"
+import (
+	"io"
+
+	"github.com/btcsuite/btcd/chaincfg/chainhash"
+	"github.com/lightningnetwork/lnd/lnwire"
+)
 
 // Init is the first message sent over the watchtower wire protocol, and
-// specifies features and level of requiredness maintained by the sending node.
-// The watchtower Init message is identical to the LN Init message, except it
-// uses a different message type to ensure the two are not conflated.
+// specifies connection features bits and level of requiredness maintained by
+// the sending node. The Init message also sends the chain hash identifying the
+// network that the sender is on.
 type Init struct {
-	*lnwire.Init
+	// ConnFeatures are the feature bits being advertised for the duration
+	// of a single connection with a peer.
+	ConnFeatures *lnwire.RawFeatureVector
+
+	// ChainHash is the genesis hash of the chain that the advertiser claims
+	// to be on.
+	ChainHash chainhash.Hash
 }
 
-// NewInitMessage generates a new Init message from raw global and local feature
-// vectors.
-func NewInitMessage(gf, lf *lnwire.RawFeatureVector) *Init {
+// NewInitMessage generates a new Init message from a raw connection feature
+// vector and chain hash.
+func NewInitMessage(connFeatures *lnwire.RawFeatureVector,
+	chainHash chainhash.Hash) *Init {
+
 	return &Init{
-		Init: lnwire.NewInitMessage(gf, lf),
+		ConnFeatures: connFeatures,
+		ChainHash:    chainHash,
 	}
 }
 
+// Encode serializes the target Init into the passed io.Writer observing the
+// protocol version specified.
+//
+// This is part of the wtwire.Message interface.
+func (msg *Init) Encode(w io.Writer, pver uint32) error {
+	return WriteElements(w,
+		msg.ConnFeatures,
+		msg.ChainHash,
+	)
+}
+
+// Decode deserializes a serialized Init message stored in the passed io.Reader
+// observing the specified protocol version.
+//
+// This is part of the wtwire.Message interface.
+func (msg *Init) Decode(r io.Reader, pver uint32) error {
+	return ReadElements(r,
+		&msg.ConnFeatures,
+		&msg.ChainHash,
+	)
+}
+
 // MsgType returns the integer uniquely identifying this message type on the
 // wire.
 //
@@ -26,5 +62,13 @@ func (msg *Init) MsgType() MessageType {
 	return MsgInit
 }
 
+// MaxPayloadLength returns the maximum allowed payload size for an Init
+// complete message observing the specified protocol version.
+//
+// This is part of the wtwire.Message interface.
+func (msg *Init) MaxPayloadLength(uint32) uint32 {
+	return MaxMessagePayload
+}
+
 // A compile-time constraint to ensure Init implements the Message interface.
 var _ Message = (*Init)(nil)

watchtower/wtwire/message.go

@@ -88,7 +88,7 @@ func makeEmptyMessage(msgType MessageType) (Message, error) {
 
 	switch msgType {
 	case MsgInit:
-		msg = &Init{&lnwire.Init{}}
+		msg = &Init{}
 	case MsgCreateSession:
 		msg = &CreateSession{}
 	case MsgCreateSessionReply:

watchtower/wtwire/wtwire.go

@@ -6,8 +6,10 @@ import (
 	"io"
 
 	"github.com/btcsuite/btcd/btcec"
+	"github.com/btcsuite/btcd/chaincfg/chainhash"
 	"github.com/btcsuite/btcd/wire"
 	"github.com/lightningnetwork/lnd/lnwallet"
+	"github.com/lightningnetwork/lnd/lnwire"
 	"github.com/lightningnetwork/lnd/watchtower/blob"
 )
 
@@ -86,6 +88,20 @@ func WriteElement(w io.Writer, element interface{}) error {
 			return err
 		}
 
+	case chainhash.Hash:
+		if _, err := w.Write(e[:]); err != nil {
+			return err
+		}
+
+	case *lnwire.RawFeatureVector:
+		if e == nil {
+			return fmt.Errorf("cannot write nil feature vector")
+		}
+
+		if err := e.Encode(w); err != nil {
+			return err
+		}
+
 	case *btcec.PublicKey:
 		if e == nil {
 			return fmt.Errorf("cannot write nil pubkey")
@@ -192,6 +208,20 @@ func ReadElement(r io.Reader, element interface{}) error {
 		}
 		*e = ErrorCode(binary.BigEndian.Uint16(b[:]))
 
+	case *chainhash.Hash:
+		if _, err := io.ReadFull(r, e[:]); err != nil {
+			return err
+		}
+
+	case **lnwire.RawFeatureVector:
+		f := lnwire.NewRawFeatureVector()
+		err := f.Decode(r)
+		if err != nil {
+			return err
+		}
+
+		*e = f
+
 	case **btcec.PublicKey:
 		var b [btcec.PubKeyBytesLenCompressed]byte
 		if _, err := io.ReadFull(r, b[:]); err != nil {

watchtower/wtwire/wtwire_test.go

@@ -8,6 +8,7 @@ import (
 	"testing/quick"
 	"time"
 
+	"github.com/btcsuite/btcd/chaincfg/chainhash"
 	"github.com/davecgh/go-spew/spew"
 	"github.com/lightningnetwork/lnd/lnwire"
 	"github.com/lightningnetwork/lnd/watchtower/wtwire"
@@ -23,6 +24,12 @@ func randRawFeatureVector(r *rand.Rand) *lnwire.RawFeatureVector {
 	return featureVec
 }
 
+func randChainHash(r *rand.Rand) chainhash.Hash {
+	var hash chainhash.Hash
+	r.Read(hash[:])
+	return hash
+}
+
 // TestWatchtowerWireProtocol uses the testing/quick package to create a series
 // of fuzz tests to attempt to break a primary scenario which is implemented as
 // property based testing scenario.
@@ -73,7 +80,7 @@ func TestWatchtowerWireProtocol(t *testing.T) {
 		wtwire.MsgInit: func(v []reflect.Value, r *rand.Rand) {
 			req := wtwire.NewInitMessage(
 				randRawFeatureVector(r),
-				randRawFeatureVector(r),
+				randChainHash(r),
 			)
 
 			v[0] = reflect.ValueOf(*req)