From dc32ca61f89dc89ebea0393c3b1083ee79439237 Mon Sep 17 00:00:00 2001
From: Oliver Gugger <gugger@gmail.com>
Date: Wed, 6 Jul 2022 21:16:56 +0200
Subject: [PATCH] rpcperms: add replaceProtoMsg

Because of the way the gRPC Receive() method is designed, we need a way
to replace a proto message with the content of another one without
replacing the original instance itself (e.g. overwrite all values in the
existing struct instance).
---
 rpcperms/middleware_handler.go      | 26 +++++++++
 rpcperms/middleware_handler_test.go | 90 +++++++++++++++++++++++++++++
 2 files changed, 116 insertions(+)
 create mode 100644 rpcperms/middleware_handler_test.go

diff --git a/rpcperms/middleware_handler.go b/rpcperms/middleware_handler.go
index 4663bec8a..0a890d2c3 100644
--- a/rpcperms/middleware_handler.go
+++ b/rpcperms/middleware_handler.go
@@ -529,3 +529,29 @@ func parseProto(typeName string, serialized []byte) (proto.Message, error) {
 
 	return msg.Interface(), nil
 }
+
+// replaceProtoMsg replaces the given target message with the content of the
+// replacement message.
+func replaceProtoMsg(target interface{}, replacement interface{}) error {
+	targetMsg, ok := target.(proto.Message)
+	if !ok {
+		return fmt.Errorf("target is not a proto message: %v", target)
+	}
+
+	replacementMsg, ok := replacement.(proto.Message)
+	if !ok {
+		return fmt.Errorf("replacement is not a proto message: %v",
+			replacement)
+	}
+
+	if targetMsg.ProtoReflect().Type() !=
+		replacementMsg.ProtoReflect().Type() {
+
+		return fmt.Errorf("replacement message is of wrong type")
+	}
+
+	proto.Reset(targetMsg)
+	proto.Merge(targetMsg, replacementMsg)
+
+	return nil
+}
diff --git a/rpcperms/middleware_handler_test.go b/rpcperms/middleware_handler_test.go
new file mode 100644
index 000000000..aa8715bf0
--- /dev/null
+++ b/rpcperms/middleware_handler_test.go
@@ -0,0 +1,90 @@
+package rpcperms
+
+import (
+	"encoding/json"
+	"testing"
+
+	"github.com/lightningnetwork/lnd/lnrpc"
+	"github.com/stretchr/testify/require"
+)
+
+// TestReplaceProtoMsg makes sure the proto message replacement works as
+// expected.
+func TestReplaceProtoMsg(t *testing.T) {
+	testCases := []struct {
+		name        string
+		original    interface{}
+		replacement interface{}
+		expectedErr string
+	}{{
+		name: "simple content replacement",
+		original: &lnrpc.Invoice{
+			Memo:  "This is a memo string",
+			Value: 123456,
+		},
+		replacement: &lnrpc.Invoice{
+			Memo:  "This is the replaced string",
+			Value: 654321,
+		},
+	}, {
+		name: "replace with empty message",
+		original: &lnrpc.Invoice{
+			Memo:  "This is a memo string",
+			Value: 123456,
+		},
+		replacement: &lnrpc.Invoice{},
+	}, {
+		name: "replace with fewer fields",
+		original: &lnrpc.Invoice{
+			Memo:  "This is a memo string",
+			Value: 123456,
+		},
+		replacement: &lnrpc.Invoice{
+			Value: 654321,
+		},
+	}, {
+		name: "wrong replacement type",
+		original: &lnrpc.Invoice{
+			Memo:  "This is a memo string",
+			Value: 123456,
+		},
+		replacement: &lnrpc.AddInvoiceResponse{},
+		expectedErr: "replacement message is of wrong type",
+	}, {
+		name:     "wrong original type",
+		original: &interceptRequest{},
+		replacement: &lnrpc.Invoice{
+			Memo:  "This is the replaced string",
+			Value: 654321,
+		},
+		expectedErr: "target is not a proto message",
+	}}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(tt *testing.T) {
+			err := replaceProtoMsg(tc.original, tc.replacement)
+
+			if tc.expectedErr != "" {
+				require.Error(tt, err)
+				require.Contains(
+					tt, err.Error(), tc.expectedErr,
+				)
+
+				return
+			}
+
+			require.NoError(tt, err)
+			jsonEqual(tt, tc.replacement, tc.original)
+		})
+	}
+}
+
+func jsonEqual(t *testing.T, expected, actual interface{}) {
+	expectedJSON, err := json.Marshal(expected)
+	require.NoError(t, err)
+
+	actualJSON, err := json.Marshal(actual)
+	require.NoError(t, err)
+
+	require.JSONEq(t, string(expectedJSON), string(actualJSON))
+}