From 956d20ebdcea83528065a2fa7b77d572176faf9a Mon Sep 17 00:00:00 2001
From: Olaoluwa Osuntokun <laolu32@gmail.com>
Date: Fri, 30 Mar 2018 18:42:27 -0700
Subject: [PATCH] funding+lnd: ensure we reconnect to new channel peers

In this commit, we fix a minor bug in the prior versions of lnd. Before
this commit, if we received a new inbound connection for channel
creation, the channel was created, and then the peer disconnected, we
wouldn't automatically reconnect.

In this commit we fix this issue by overloading the WatchNewChannel
method to also accept the peer's ID so we can add it to the set of
persistent connections.
---
 fundingmanager.go | 14 +++++++-------
 lnd.go            | 15 ++++++++++++++-
 2 files changed, 21 insertions(+), 8 deletions(-)

diff --git a/fundingmanager.go b/fundingmanager.go
index 2fc2e1081..0171ebb12 100644
--- a/fundingmanager.go
+++ b/fundingmanager.go
@@ -299,8 +299,9 @@ type fundingConfig struct {
 	// WatchNewChannel is to be called once a new channel enters the final
 	// funding stage: waiting for on-chain confirmation. This method sends
 	// the channel to the ChainArbitrator so it can watch for any on-chain
-	// events related to the channel.
-	WatchNewChannel func(*channeldb.OpenChannel) error
+	// events related to the channel. We also provide the address of the
+	// node we're establishing a channel with for reconnection purposes.
+	WatchNewChannel func(*channeldb.OpenChannel, *lnwire.NetAddress) error
 
 	// ReportShortChanID allows the funding manager to report the newly
 	// discovered short channel ID of a formerly pending channel to outside
@@ -952,9 +953,6 @@ func (f *fundingManager) handleFundingOpen(fmsg *fundingOpenMsg) {
 	// reservation attempt may be rejected. Note that since we're on the
 	// responding side of a single funder workflow, we don't commit any
 	// funds to the channel ourselves.
-	//
-	// TODO(roasbeef): assuming this was an inbound connection, replace
-	// port with default advertised port
 	chainHash := chainhash.Hash(msg.ChainHash)
 	reservation, err := f.cfg.Wallet.InitChannelReservation(
 		amt, 0, msg.PushAmount,
@@ -1355,7 +1353,8 @@ func (f *fundingManager) handleFundingCreated(fmsg *fundingCreatedMsg) {
 	// Now that we've sent over our final signature for this channel, we'll
 	// send it to the ChainArbitrator so it can watch for any on-chain
 	// actions during this final confirmation stage.
-	if err := f.cfg.WatchNewChannel(completeChan); err != nil {
+	peerAddr := resCtx.peerAddress
+	if err := f.cfg.WatchNewChannel(completeChan, peerAddr); err != nil {
 		fndgLog.Errorf("Unable to send new ChannelPoint(%v) for "+
 			"arbitration: %v", fundingOut, err)
 	}
@@ -1503,7 +1502,8 @@ func (f *fundingManager) handleFundingSigned(fmsg *fundingSignedMsg) {
 	// we'll send the to be active channel to the ChainArbitrator so it can
 	// watch for any on-chin actions before the channel has fully
 	// confirmed.
-	if err := f.cfg.WatchNewChannel(completeChan); err != nil {
+	peerAddr := resCtx.peerAddress
+	if err := f.cfg.WatchNewChannel(completeChan, peerAddr); err != nil {
 		fndgLog.Errorf("Unable to send new ChannelPoint(%v) for "+
 			"arbitration: %v", fundingPoint, err)
 	}
diff --git a/lnd.go b/lnd.go
index 11c77bdd5..c911315b7 100644
--- a/lnd.go
+++ b/lnd.go
@@ -430,7 +430,20 @@ func lndMain() error {
 			}
 			return delay
 		},
-		WatchNewChannel: server.chainArb.WatchNewChannel,
+		WatchNewChannel: func(channel *channeldb.OpenChannel,
+			addr *lnwire.NetAddress) error {
+
+			// First, we'll mark this new peer as a persistent
+			// re-connection purposes.
+			server.mu.Lock()
+			pubStr := string(addr.IdentityKey.SerializeCompressed())
+			server.persistentPeers[pubStr] = struct{}{}
+			server.mu.Unlock()
+
+			// With that taken care of, we'll send this channel to
+			// the chain arb so it can react to on-chain events.
+			return server.chainArb.WatchNewChannel(channel)
+		},
 		ReportShortChanID: func(chanPoint wire.OutPoint,
 			sid lnwire.ShortChannelID) error {