From bbca53507fdc14fe4e6dfa34e7f0abec84044794 Mon Sep 17 00:00:00 2001
From: Olaoluwa Osuntokun <laolu32@gmail.com>
Date: Fri, 19 Jan 2018 17:17:14 -0800
Subject: [PATCH] contractcourt: extend the ChainArbitratorConfig with
 IsOurAddress closure

In this commit, we add the IsOurAddress field into the config of the
chain arb. With this new function closure, the chain arb is able to
detect co-op on chain closes automatically.
---
 contractcourt/chain_arbitrator.go | 14 +++++++++++++-
 server.go                         |  4 ++++
 2 files changed, 17 insertions(+), 1 deletion(-)

diff --git a/contractcourt/chain_arbitrator.go b/contractcourt/chain_arbitrator.go
index df3a7cc07..3b984a8d2 100644
--- a/contractcourt/chain_arbitrator.go
+++ b/contractcourt/chain_arbitrator.go
@@ -11,6 +11,7 @@ import (
 	"github.com/lightningnetwork/lnd/lnwire"
 	"github.com/roasbeef/btcd/chaincfg/chainhash"
 	"github.com/roasbeef/btcd/wire"
+	"github.com/roasbeef/btcutil"
 )
 
 // ResolutionMsg is a message sent by resolvers to outside sub-systems once an
@@ -77,9 +78,14 @@ type ChainArbitratorConfig struct {
 	// ChannelArbitrator decides that it needs to go to chain in order to
 	// resolve contracts.
 	//
-	// TODO(Roasbeef): rename, routing based
+	// TODO(roasbeef): rename, routing based
 	MarkLinkInactive func(wire.OutPoint) error
 
+	// IsOurAddress is a function that returns true if the passed address
+	// is known to the underlying wallet. Otherwise, false should be
+	// returned.
+	IsOurAddress func(btcutil.Address) bool
+
 	// IncubateOutput sends either a incoming HTLC, an outgoing HTLC, or
 	// both to the utxo nursery. Once this function returns, the nursery
 	// should have safely persisted the outputs to disk, and should start
@@ -323,6 +329,9 @@ func (c *ChainArbitrator) Start() error {
 		// to ensure that we detect any relevant on chain events.
 		chainWatcher, err := newChainWatcher(
 			channel, c.cfg.Notifier, c.cfg.PreimageDB, c.cfg.Signer,
+			c.cfg.IsOurAddress, func() error {
+				return c.resolveContract(channel.FundingOutpoint, nil)
+			},
 		)
 		if err != nil {
 			return err
@@ -572,6 +581,9 @@ func (c *ChainArbitrator) WatchNewChannel(newChan *channeldb.OpenChannel) error
 	// that we detect any relevant on chain events.
 	chainWatcher, err := newChainWatcher(
 		newChan, c.cfg.Notifier, c.cfg.PreimageDB, c.cfg.Signer,
+		c.cfg.IsOurAddress, func() error {
+			return c.resolveContract(chanPoint, nil)
+		},
 	)
 	if err != nil {
 		return err
diff --git a/server.go b/server.go
index 820dc3156..247987c62 100644
--- a/server.go
+++ b/server.go
@@ -400,6 +400,10 @@ func newServer(listenAddrs []string, chanDB *channeldb.DB, cc *chainControl,
 			chanID := lnwire.NewChanIDFromOutPoint(&chanPoint)
 			return s.htlcSwitch.RemoveLink(chanID)
 		},
+		IsOurAddress: func(addr btcutil.Address) bool {
+			_, err := cc.wallet.GetPrivKey(addr)
+			return err == nil
+		},
 	}, chanDB)
 
 	s.breachArbiter = newBreachArbiter(&BreachConfig{