From 86cbc2a0ad52de11e62bc282a3a8ea1c5536a174 Mon Sep 17 00:00:00 2001
From: Olaoluwa Osuntokun <laolu32@gmail.com>
Date: Sun, 23 Jun 2024 23:30:52 -0700
Subject: [PATCH] contractcourt: update GenSweepScript to return internal key

For the upcoming aux sweeper integration, the internal key is needed for
the call backs.
---
 contractcourt/breach_arbitrator.go      |  6 +++---
 contractcourt/breach_arbitrator_test.go | 22 +++++++++++++---------
 server.go                               | 13 +++----------
 3 files changed, 19 insertions(+), 22 deletions(-)

diff --git a/contractcourt/breach_arbitrator.go b/contractcourt/breach_arbitrator.go
index f86eab397..792cc63f9 100644
--- a/contractcourt/breach_arbitrator.go
+++ b/contractcourt/breach_arbitrator.go
@@ -149,7 +149,7 @@ type BreachConfig struct {
 	Estimator chainfee.Estimator
 
 	// GenSweepScript generates the receiving scripts for swept outputs.
-	GenSweepScript func() ([]byte, error)
+	GenSweepScript func() fn.Result[lnwallet.AddrWithKey]
 
 	// Notifier provides a publish/subscribe interface for event driven
 	// notifications regarding the confirmation of txids.
@@ -1523,7 +1523,7 @@ func (b *BreachArbitrator) sweepSpendableOutputsTxn(txWeight lntypes.WeightUnit,
 	// sweep the funds to.
 	// TODO(roasbeef): possibly create many outputs to minimize change in
 	// the future?
-	pkScript, err := b.cfg.GenSweepScript()
+	pkScript, err := b.cfg.GenSweepScript().Unpack()
 	if err != nil {
 		return nil, err
 	}
@@ -1551,7 +1551,7 @@ func (b *BreachArbitrator) sweepSpendableOutputsTxn(txWeight lntypes.WeightUnit,
 
 	// We begin by adding the output to which our funds will be deposited.
 	txn.AddTxOut(&wire.TxOut{
-		PkScript: pkScript,
+		PkScript: pkScript.DeliveryAddress,
 		Value:    sweepAmt,
 	})
 
diff --git a/contractcourt/breach_arbitrator_test.go b/contractcourt/breach_arbitrator_test.go
index 896342d42..66b934c9c 100644
--- a/contractcourt/breach_arbitrator_test.go
+++ b/contractcourt/breach_arbitrator_test.go
@@ -2131,15 +2131,19 @@ func createTestArbiter(t *testing.T, contractBreaches chan *ContractBreachEvent,
 	// Assemble our test arbiter.
 	notifier := mock.MakeMockSpendNotifier()
 	ba := NewBreachArbitrator(&BreachConfig{
-		CloseLink:          func(_ *wire.OutPoint, _ ChannelCloseType) {},
-		DB:                 db.ChannelStateDB(),
-		Estimator:          chainfee.NewStaticEstimator(12500, 0),
-		GenSweepScript:     func() ([]byte, error) { return nil, nil },
-		ContractBreaches:   contractBreaches,
-		Signer:             signer,
-		Notifier:           notifier,
-		PublishTransaction: func(_ *wire.MsgTx, _ string) error { return nil },
-		Store:              store,
+		CloseLink: func(_ *wire.OutPoint, _ ChannelCloseType) {},
+		DB:        db.ChannelStateDB(),
+		Estimator: chainfee.NewStaticEstimator(12500, 0),
+		GenSweepScript: func() fn.Result[lnwallet.AddrWithKey] {
+			return fn.Ok(lnwallet.AddrWithKey{})
+		},
+		ContractBreaches: contractBreaches,
+		Signer:           signer,
+		Notifier:         notifier,
+		PublishTransaction: func(_ *wire.MsgTx, _ string) error {
+			return nil
+		},
+		Store: store,
 	})
 
 	if err := ba.Start(); err != nil {
diff --git a/server.go b/server.go
index 2dbc85d63..6022d4f7a 100644
--- a/server.go
+++ b/server.go
@@ -1152,16 +1152,9 @@ func newServer(cfg *Config, listenAddrs []net.Addr,
 			CloseLink: closeLink,
 			DB:        s.chanStateDB,
 			Estimator: s.cc.FeeEstimator,
-			GenSweepScript: func() ([]byte, error) {
-				addr, err := newSweepPkScriptGen(
-					cc.Wallet, netParams,
-				)().Unpack()
-				if err != nil {
-					return nil, err
-				}
-
-				return addr.DeliveryAddress, nil
-			},
+			GenSweepScript: newSweepPkScriptGen(
+				cc.Wallet, s.cfg.ActiveNetParams.Params,
+			),
 			Notifier:           cc.ChainNotifier,
 			PublishTransaction: cc.Wallet.PublishTransaction,
 			ContractBreaches:   contractBreaches,