From fe9f703b5b51b9c2298d0a2f7cf527eb6a0dc665 Mon Sep 17 00:00:00 2001
From: Conner Fromknecht <conner@lightning.engineering>
Date: Wed, 6 Feb 2019 18:19:58 -0800
Subject: [PATCH] watchtower/wtclient/backup_task: correct txweight + update
 test vectors

---
 watchtower/wtclient/backup_task.go               | 14 ++++++++++----
 watchtower/wtclient/backup_task_internal_test.go |  8 ++++----
 2 files changed, 14 insertions(+), 8 deletions(-)

diff --git a/watchtower/wtclient/backup_task.go b/watchtower/wtclient/backup_task.go
index dd1d4de8..1e5dd4b2 100644
--- a/watchtower/wtclient/backup_task.go
+++ b/watchtower/wtclient/backup_task.go
@@ -130,9 +130,6 @@ func (t *backupTask) bindSession(session *wtdb.SessionInfo,
 	// the watchtower is taking a reward.
 	var weightEstimate input.TxWeightEstimator
 
-	// All justice transactions have a p2wkh output paying to the victim.
-	weightEstimate.AddP2WKHOutput()
-
 	// Next, add the contribution from the inputs that are present on this
 	// breach transaction.
 	if t.toLocalInput != nil {
@@ -142,6 +139,15 @@ func (t *backupTask) bindSession(session *wtdb.SessionInfo,
 		weightEstimate.AddWitnessInput(input.P2WKHWitnessSize)
 	}
 
+	// All justice transactions have a p2wkh output paying to the victim.
+	weightEstimate.AddP2WKHOutput()
+
+	// If the justice transaction has a reward output, add the output's
+	// contribution to the weight estimate.
+	if session.Policy.BlobType.Has(blob.FlagReward) {
+		weightEstimate.AddP2WKHOutput()
+	}
+
 	// Now, compute the output values depending on whether FlagReward is set
 	// in the current session's policy.
 	outputs, err := session.Policy.ComputeJusticeTxOuts(
@@ -152,8 +158,8 @@ func (t *backupTask) bindSession(session *wtdb.SessionInfo,
 		return err
 	}
 
-	t.outputs = outputs
 	t.blobType = session.Policy.BlobType
+	t.outputs = outputs
 
 	return nil
 }
diff --git a/watchtower/wtclient/backup_task_internal_test.go b/watchtower/wtclient/backup_task_internal_test.go
index 2d39c600..26172620 100644
--- a/watchtower/wtclient/backup_task_internal_test.go
+++ b/watchtower/wtclient/backup_task_internal_test.go
@@ -312,7 +312,7 @@ var backupTaskTests = []backupTaskTest{
 		blobTypeCommitReward, // blobType
 		1000,                 // sweepFeeRate
 		addrScript,           // rewardScript
-		296241,               // expSweepAmt
+		296117,               // expSweepAmt
 		3000,                 // expRewardAmt
 		nil,                  // bindErr
 	),
@@ -324,7 +324,7 @@ var backupTaskTests = []backupTaskTest{
 		blobTypeCommitReward, // blobType
 		1000,                 // sweepFeeRate
 		addrScript,           // rewardScript
-		197514,               // expSweepAmt
+		197390,               // expSweepAmt
 		2000,                 // expRewardAmt
 		nil,                  // bindErr
 	),
@@ -336,7 +336,7 @@ var backupTaskTests = []backupTaskTest{
 		blobTypeCommitReward, // blobType
 		1000,                 // sweepFeeRate
 		addrScript,           // rewardScript
-		98561,                // expSweepAmt
+		98437,                // expSweepAmt
 		1000,                 // expRewardAmt
 		nil,                  // bindErr
 	),
@@ -346,7 +346,7 @@ var backupTaskTests = []backupTaskTest{
 		0,                       // toLocalAmt
 		100000,                  // toRemoteAmt
 		blobTypeCommitReward,    // blobType
-		225000,                  // sweepFeeRate
+		175000,                  // sweepFeeRate
 		addrScript,              // rewardScript
 		0,                       // expSweepAmt
 		0,                       // expRewardAmt