From 5147a6d63d4c55a8b70584d4e7f0c222fdaa1cb5 Mon Sep 17 00:00:00 2001
From: Conner Fromknecht <conner@lightning.engineering>
Date: Tue, 17 Mar 2020 16:25:22 -0700
Subject: [PATCH] netann/channel_update: set timestamp via modifier

---
 netann/chan_status_manager.go |  2 +-
 netann/channel_update.go      | 30 +++++++++++++++++-------------
 netann/channel_update_test.go |  4 +++-
 3 files changed, 21 insertions(+), 15 deletions(-)

diff --git a/netann/chan_status_manager.go b/netann/chan_status_manager.go
index 146fef47..51797c87 100644
--- a/netann/chan_status_manager.go
+++ b/netann/chan_status_manager.go
@@ -527,7 +527,7 @@ func (m *ChanStatusManager) signAndSendNextUpdate(outpoint wire.OutPoint,
 
 	err = SignChannelUpdate(
 		m.cfg.MessageSigner, m.cfg.OurPubKey, chanUpdate,
-		ChannelUpdateSetDisable(disabled),
+		ChanUpdSetDisable(disabled), ChanUpdSetTimestamp,
 	)
 	if err != nil {
 		return err
diff --git a/netann/channel_update.go b/netann/channel_update.go
index b3c53136..a40abc25 100644
--- a/netann/channel_update.go
+++ b/netann/channel_update.go
@@ -15,9 +15,9 @@ import (
 // lnwire.ChannelUpdate.
 type ChannelUpdateModifier func(*lnwire.ChannelUpdate)
 
-// ChannelUpdateSetDisable sets the disabled channel flag if disabled is true,
-// and clears the bit otherwise.
-func ChannelUpdateSetDisable(disabled bool) ChannelUpdateModifier {
+// ChanUpdSetDisable is a functional option that sets the disabled channel flag
+// if disabled is true, and clears the bit otherwise.
+func ChanUpdSetDisable(disabled bool) ChannelUpdateModifier {
 	return func(update *lnwire.ChannelUpdate) {
 		if disabled {
 			// Set the bit responsible for marking a channel as
@@ -31,6 +31,20 @@ func ChannelUpdateSetDisable(disabled bool) ChannelUpdateModifier {
 	}
 }
 
+// ChanUpdSetTimestamp is a functional option that sets the timestamp of the
+// update to the current time, or increments it if the timestamp is already in
+// the future.
+func ChanUpdSetTimestamp(update *lnwire.ChannelUpdate) {
+	newTimestamp := uint32(time.Now().Unix())
+	if newTimestamp <= update.Timestamp {
+		// Increment the prior value to ensure the timestamp
+		// monotonically increases, otherwise the update won't
+		// propagate.
+		newTimestamp = update.Timestamp + 1
+	}
+	update.Timestamp = newTimestamp
+}
+
 // SignChannelUpdate applies the given modifiers to the passed
 // lnwire.ChannelUpdate, then signs the resulting update. The provided update
 // should be the most recent, valid update, otherwise the timestamp may not
@@ -45,16 +59,6 @@ func SignChannelUpdate(signer lnwallet.MessageSigner, pubKey *btcec.PublicKey,
 		modifier(update)
 	}
 
-	// Update the message's timestamp to the current time. If the update's
-	// current time is already in the future, we increment the prior value
-	// to ensure the timestamps monotonically increase, otherwise the
-	// update won't propagate.
-	newTimestamp := uint32(time.Now().Unix())
-	if newTimestamp <= update.Timestamp {
-		newTimestamp = update.Timestamp + 1
-	}
-	update.Timestamp = newTimestamp
-
 	// Create the DER-encoded ECDSA signature over the message digest.
 	sig, err := SignAnnouncement(signer, pubKey, update)
 	if err != nil {
diff --git a/netann/channel_update_test.go b/netann/channel_update_test.go
index 4ed019a8..0bc48e5c 100644
--- a/netann/channel_update_test.go
+++ b/netann/channel_update_test.go
@@ -103,6 +103,7 @@ func TestUpdateDisableFlag(t *testing.T) {
 	t.Parallel()
 
 	for _, tc := range updateDisableTests {
+		tc := tc
 		t.Run(tc.name, func(t *testing.T) {
 			// Create the initial update, the only fields we are
 			// concerned with in this test are the timestamp and the
@@ -127,7 +128,8 @@ func TestUpdateDisableFlag(t *testing.T) {
 			// disabled or enabled as prescribed in the test case.
 			err := netann.SignChannelUpdate(
 				tc.signer, pubKey, newUpdate,
-				netann.ChannelUpdateSetDisable(tc.disable),
+				netann.ChanUpdSetDisable(tc.disable),
+				netann.ChanUpdSetTimestamp,
 			)
 
 			var fail bool