From ed8a672bd39eeb0610c501501023d69af37a019d Mon Sep 17 00:00:00 2001
From: Olaoluwa Osuntokun <laolu32@gmail.com>
Date: Thu, 27 Feb 2025 13:56:38 -0800
Subject: [PATCH] protofsm: update StateMachine to meet new msgmux.Endpoint
 interface

---
 protofsm/state_machine.go      |  5 +++--
 protofsm/state_machine_test.go | 11 ++++++++---
 2 files changed, 11 insertions(+), 5 deletions(-)

diff --git a/protofsm/state_machine.go b/protofsm/state_machine.go
index bf61d727c..0e0343587 100644
--- a/protofsm/state_machine.go
+++ b/protofsm/state_machine.go
@@ -14,6 +14,7 @@ import (
 	"github.com/lightningnetwork/lnd/fn/v2"
 	"github.com/lightningnetwork/lnd/lnutils"
 	"github.com/lightningnetwork/lnd/lnwire"
+	"github.com/lightningnetwork/lnd/msgmux"
 )
 
 const (
@@ -250,7 +251,7 @@ func (s *StateMachine[Event, Env]) SendEvent(ctx context.Context, event Event) {
 
 // CanHandle returns true if the target message can be routed to the state
 // machine.
-func (s *StateMachine[Event, Env]) CanHandle(msg lnwire.Message) bool {
+func (s *StateMachine[Event, Env]) CanHandle(msg msgmux.PeerMsg) bool {
 	cfgMapper := s.cfg.MsgMapper
 	return fn.MapOptionZ(cfgMapper, func(mapper MsgMapper[Event]) bool {
 		return mapper.MapMsg(msg).IsSome()
@@ -267,7 +268,7 @@ func (s *StateMachine[Event, Env]) Name() string {
 // returned indicating that the message was processed. Otherwise, false is
 // returned.
 func (s *StateMachine[Event, Env]) SendMessage(ctx context.Context,
-	msg lnwire.Message) bool {
+	msg msgmux.PeerMsg) bool {
 
 	// If we have no message mapper, then return false as we can't process
 	// this message.
diff --git a/protofsm/state_machine_test.go b/protofsm/state_machine_test.go
index 8ac5ec603..ed18743b5 100644
--- a/protofsm/state_machine_test.go
+++ b/protofsm/state_machine_test.go
@@ -13,6 +13,7 @@ import (
 	"github.com/lightningnetwork/lnd/chainntnfs"
 	"github.com/lightningnetwork/lnd/fn/v2"
 	"github.com/lightningnetwork/lnd/lnwire"
+	"github.com/lightningnetwork/lnd/msgmux"
 	"github.com/stretchr/testify/mock"
 	"github.com/stretchr/testify/require"
 )
@@ -429,8 +430,12 @@ func TestStateMachineMsgMapper(t *testing.T) {
 
 	// The only thing we know how to map is the error message, which'll
 	// terminate the state machine.
-	wireError := &lnwire.Error{}
-	initMsg := &lnwire.Init{}
+	wireError := msgmux.PeerMsg{
+		Message: &lnwire.Error{},
+	}
+	initMsg := msgmux.PeerMsg{
+		Message: &lnwire.Init{},
+	}
 	dummyMapper.On("MapMsg", wireError).Return(
 		fn.Some(dummyEvents(&goToFin{})),
 	)
@@ -456,7 +461,7 @@ func TestStateMachineMsgMapper(t *testing.T) {
 
 	// First, we'll verify that the CanHandle method works as expected.
 	require.True(t, stateMachine.CanHandle(wireError))
-	require.False(t, stateMachine.CanHandle(&lnwire.Init{}))
+	require.False(t, stateMachine.CanHandle(initMsg))
 
 	// Next, we'll attempt to send the wire message into the state machine.
 	// We should transition to the final state.