From 52e42fc1071fa075eb59c3b43a48d26d082bf51c Mon Sep 17 00:00:00 2001
From: Oliver Gugger <gugger@gmail.com>
Date: Wed, 4 May 2022 18:31:47 +0200
Subject: [PATCH] lnrpc+itest: implement MuSig2Cleanup RPC

---
 lnrpc/signrpc/signer_server.go   | 22 ++++++++++++++++++++++
 lntest/itest/lnd_taproot_test.go | 20 +++++++++++++++++++-
 2 files changed, 41 insertions(+), 1 deletion(-)

diff --git a/lnrpc/signrpc/signer_server.go b/lnrpc/signrpc/signer_server.go
index 9523e12aa..ba98f276d 100644
--- a/lnrpc/signrpc/signer_server.go
+++ b/lnrpc/signrpc/signer_server.go
@@ -92,6 +92,10 @@ var (
 			Entity: "signer",
 			Action: "generate",
 		}},
+		"/signrpc.Signer/MuSig2Cleanup": {{
+			Entity: "signer",
+			Action: "generate",
+		}},
 	}
 
 	// DefaultSignerMacFilename is the default name of the signer macaroon
@@ -940,6 +944,24 @@ func (s *Server) MuSig2CombineSig(_ context.Context,
 	}, nil
 }
 
+// MuSig2Cleanup removes a session from memory to free up resources.
+func (s *Server) MuSig2Cleanup(_ context.Context,
+	in *MuSig2CleanupRequest) (*MuSig2CleanupResponse, error) {
+
+	// Check session ID length.
+	sessionID, err := parseMuSig2SessionID(in.SessionId)
+	if err != nil {
+		return nil, fmt.Errorf("error parsing session ID: %v", err)
+	}
+
+	err = s.cfg.Signer.MuSig2Cleanup(sessionID)
+	if err != nil {
+		return nil, fmt.Errorf("error cleaning up session: %v", err)
+	}
+
+	return &MuSig2CleanupResponse{}, nil
+}
+
 // parseRawKeyBytes checks that the provided raw public key is valid and returns
 // the public key. A nil public key is returned if the length of the rawKeyBytes
 // is zero.
diff --git a/lntest/itest/lnd_taproot_test.go b/lntest/itest/lnd_taproot_test.go
index e63be8ac1..589b75e8e 100644
--- a/lntest/itest/lnd_taproot_test.go
+++ b/lntest/itest/lnd_taproot_test.go
@@ -906,11 +906,29 @@ func testTaprootMuSig2CombinedLeafKeySpend(ctxt context.Context, t *harnessTest,
 		ctxt, &signrpc.MuSig2SignRequest{
 			SessionId:     sessResp3.SessionId,
 			MessageDigest: sigHash,
-			Cleanup:       true,
 		},
 	)
 	require.NoError(t.t, err)
 
+	// We manually clean up session 3, just to make sure that works as well.
+	_, err = alice.SignerClient.MuSig2Cleanup(
+		ctxt, &signrpc.MuSig2CleanupRequest{
+			SessionId: sessResp3.SessionId,
+		},
+	)
+	require.NoError(t.t, err)
+
+	// A second call to that cleaned up session should now fail with a
+	// specific error.
+	_, err = alice.SignerClient.MuSig2Sign(
+		ctxt, &signrpc.MuSig2SignRequest{
+			SessionId:     sessResp3.SessionId,
+			MessageDigest: sigHash,
+		},
+	)
+	require.Error(t.t, err)
+	require.Contains(t.t, err.Error(), "not found")
+
 	// Luckily only one of the signers needs to combine the signature, so
 	// let's do that now.
 	combineReq1, err := alice.SignerClient.MuSig2CombineSig(