From a5afcd21b73a2cda3ea2c6d45cf8636779b0e020 Mon Sep 17 00:00:00 2001
From: Elle Mouton <elle.mouton@gmail.com>
Date: Tue, 7 Nov 2023 09:41:41 +0200
Subject: [PATCH] netann: let SignChannelUpdate take ChannelUpdate interface

---
 netann/channel_update.go | 49 ++++++++++++++++++++++++++++++++--------
 1 file changed, 39 insertions(+), 10 deletions(-)

diff --git a/netann/channel_update.go b/netann/channel_update.go
index a81ccd504..a947aca43 100644
--- a/netann/channel_update.go
+++ b/netann/channel_update.go
@@ -67,7 +67,7 @@ func ChanUpdSetTimestamp(bestBlockHeight uint32) ChannelUpdateModifier {
 //
 // NOTE: This method modifies the given update.
 func SignChannelUpdate(signer keychain.MessageSignerRing,
-	keyLoc keychain.KeyLocator, update *lnwire.ChannelUpdate1,
+	keyLoc keychain.KeyLocator, update lnwire.ChannelUpdate,
 	mods ...ChannelUpdateModifier) error {
 
 	// Apply the requested changes to the channel update.
@@ -75,16 +75,45 @@ func SignChannelUpdate(signer keychain.MessageSignerRing,
 		modifier(update)
 	}
 
-	// Create the DER-encoded ECDSA signature over the message digest.
-	sig, err := SignAnnouncement(signer, keyLoc, update)
-	if err != nil {
-		return err
-	}
+	switch upd := update.(type) {
+	case *lnwire.ChannelUpdate1:
+		data, err := upd.DataToSign()
+		if err != nil {
+			return err
+		}
 
-	// Parse the DER-encoded signature into a fixed-size 64-byte array.
-	update.Signature, err = lnwire.NewSigFromSignature(sig)
-	if err != nil {
-		return err
+		sig, err := signer.SignMessage(keyLoc, data, true)
+		if err != nil {
+			return err
+		}
+
+		// Parse the DER-encoded signature into a fixed-size 64-byte
+		// array.
+		upd.Signature, err = lnwire.NewSigFromSignature(sig)
+		if err != nil {
+			return err
+		}
+
+	case *lnwire.ChannelUpdate2:
+		data, err := upd.DataToSign()
+		if err != nil {
+			return err
+		}
+
+		sig, err := signer.SignMessageSchnorr(
+			keyLoc, data, false, nil, upd.DigestTag(),
+		)
+		if err != nil {
+			return err
+		}
+
+		upd.Signature, err = lnwire.NewSigFromSignature(sig)
+		if err != nil {
+			return err
+		}
+	default:
+		return fmt.Errorf("unhandled implementation of "+
+			"ChannelUpdate: %T", update)
 	}
 
 	return nil