From 95ff670c0e6c7aa32d384ad6ce0649c6aad6a14f Mon Sep 17 00:00:00 2001
From: Oliver Gugger <gugger@gmail.com>
Date: Wed, 4 May 2022 18:31:44 +0200
Subject: [PATCH] multi: add MuSig2Cleanup method to MuSig2 signer

---
 input/musig2.go                 |  3 +++
 input/test_utils.go             |  5 +++++
 lntest/mock/signer.go           | 10 ++++++++++
 lnwallet/btcwallet/signer.go    | 18 ++++++++++++++++++
 lnwallet/rpcwallet/rpcwallet.go | 18 ++++++++++++++++++
 watchtower/wtmock/signer.go     |  5 +++++
 6 files changed, 59 insertions(+)

diff --git a/input/musig2.go b/input/musig2.go
index a7a0d78e4..2fc689d77 100644
--- a/input/musig2.go
+++ b/input/musig2.go
@@ -58,6 +58,9 @@ type MuSig2Signer interface {
 	// returned.
 	MuSig2CombineSig(MuSig2SessionID,
 		[]*musig2.PartialSignature) (*schnorr.Signature, bool, error)
+
+	// MuSig2Cleanup removes a session from memory to free up resources.
+	MuSig2Cleanup(MuSig2SessionID) error
 }
 
 // MuSig2SessionInfo is a struct for keeping track of a signing session
diff --git a/input/test_utils.go b/input/test_utils.go
index 531a0d672..99b8e050f 100644
--- a/input/test_utils.go
+++ b/input/test_utils.go
@@ -178,6 +178,11 @@ func (m *MockSigner) MuSig2CombineSig(MuSig2SessionID,
 	return nil, false, nil
 }
 
+// MuSig2Cleanup removes a session from memory to free up resources.
+func (m *MockSigner) MuSig2Cleanup(MuSig2SessionID) error {
+	return nil
+}
+
 // findKey searches through all stored private keys and returns one
 // corresponding to the hashed pubkey if it can be found. The public key may
 // either correspond directly to the private key or to the private key with a
diff --git a/lntest/mock/signer.go b/lntest/mock/signer.go
index 7cfd0bd26..4dd465cd8 100644
--- a/lntest/mock/signer.go
+++ b/lntest/mock/signer.go
@@ -94,6 +94,11 @@ func (d *DummySigner) MuSig2CombineSig(input.MuSig2SessionID,
 	return nil, false, nil
 }
 
+// MuSig2Cleanup removes a session from memory to free up resources.
+func (d *DummySigner) MuSig2Cleanup(input.MuSig2SessionID) error {
+	return nil
+}
+
 // SingleSigner is an implementation of the Signer interface that signs
 // everything with a single private key.
 type SingleSigner struct {
@@ -227,3 +232,8 @@ func (s *SingleSigner) MuSig2CombineSig(input.MuSig2SessionID,
 
 	return nil, false, nil
 }
+
+// MuSig2Cleanup removes a session from memory to free up resources.
+func (s *SingleSigner) MuSig2Cleanup(input.MuSig2SessionID) error {
+	return nil
+}
diff --git a/lnwallet/btcwallet/signer.go b/lnwallet/btcwallet/signer.go
index a52805b7d..d1c88509c 100644
--- a/lnwallet/btcwallet/signer.go
+++ b/lnwallet/btcwallet/signer.go
@@ -704,6 +704,24 @@ func (b *BtcWallet) MuSig2CombineSig(sessionID input.MuSig2SessionID,
 	return finalSig, session.HaveAllSigs, nil
 }
 
+// MuSig2Cleanup removes a session from memory to free up resources.
+func (b *BtcWallet) MuSig2Cleanup(sessionID input.MuSig2SessionID) error {
+	// We hold the lock during the whole operation, we don't want any
+	// interference with calls that might come through in parallel for the
+	// same session.
+	b.musig2SessionsMtx.Lock()
+	defer b.musig2SessionsMtx.Unlock()
+
+	_, ok := b.musig2Sessions[sessionID]
+	if !ok {
+		return fmt.Errorf("session with ID %x not found", sessionID[:])
+	}
+
+	delete(b.musig2Sessions, sessionID)
+
+	return nil
+}
+
 // A compile time check to ensure that BtcWallet implements the Signer
 // interface.
 var _ input.Signer = (*BtcWallet)(nil)
diff --git a/lnwallet/rpcwallet/rpcwallet.go b/lnwallet/rpcwallet/rpcwallet.go
index 0463c4d7c..1c727a140 100644
--- a/lnwallet/rpcwallet/rpcwallet.go
+++ b/lnwallet/rpcwallet/rpcwallet.go
@@ -782,6 +782,24 @@ func (r *RPCKeyRing) MuSig2CombineSig(sessionID input.MuSig2SessionID,
 	return finalSig, resp.HaveAllSignatures, nil
 }
 
+// MuSig2Cleanup removes a session from memory to free up resources.
+func (r *RPCKeyRing) MuSig2Cleanup(sessionID input.MuSig2SessionID) error {
+	req := &signrpc.MuSig2CleanupRequest{
+		SessionId: sessionID[:],
+	}
+
+	ctxt, cancel := context.WithTimeout(context.Background(), r.rpcTimeout)
+	defer cancel()
+
+	_, err := r.signerClient.MuSig2Cleanup(ctxt, req)
+	if err != nil {
+		return fmt.Errorf("error cleaning up MuSig2 session in remote "+
+			"signer instance: %v", err)
+	}
+
+	return nil
+}
+
 // remoteSign signs the input specified in signDesc of the given transaction tx
 // using the remote signing instance.
 func (r *RPCKeyRing) remoteSign(tx *wire.MsgTx, signDesc *input.SignDescriptor,
diff --git a/watchtower/wtmock/signer.go b/watchtower/wtmock/signer.go
index 94ba8171b..b41d8ceb5 100644
--- a/watchtower/wtmock/signer.go
+++ b/watchtower/wtmock/signer.go
@@ -108,6 +108,11 @@ func (s *MockSigner) MuSig2CombineSig(input.MuSig2SessionID,
 	return nil, false, nil
 }
 
+// MuSig2Cleanup removes a session from memory to free up resources.
+func (s *MockSigner) MuSig2Cleanup(input.MuSig2SessionID) error {
+	return nil
+}
+
 // AddPrivKey records the passed privKey in the MockSigner's registry of keys it
 // can sign with in the future. A unique key locator is returned, allowing the
 // caller to sign with this key when presented via an input.SignDescriptor.