1  package msgmux
  2  
  3  import (
  4  	"context"
  5  	"testing"
  6  
  7  	"github.com/lightningnetwork/lnd/lnwire"
  8  	"github.com/stretchr/testify/mock"
  9  	"github.com/stretchr/testify/require"
 10  )
 11  
 12  type mockEndpoint struct {
 13  	mock.Mock
 14  }
 15  
 16  func (m *mockEndpoint) Name() string {
 17  	args := m.Called()
 18  
 19  	return args.String(0)
 20  }
 21  
 22  func (m *mockEndpoint) CanHandle(msg PeerMsg) bool {
 23  	args := m.Called(msg)
 24  
 25  	return args.Bool(0)
 26  }
 27  
 28  func (m *mockEndpoint) SendMessage(ctx context.Context, msg PeerMsg) bool {
 29  	args := m.Called(ctx, msg)
 30  
 31  	return args.Bool(0)
 32  }
 33  
 34  // TestMessageRouterOperation tests the basic operation of the message router:
 35  // add new endpoints, route to them, remove, them, etc.
 36  func TestMessageRouterOperation(t *testing.T) {
 37  	ctx := t.Context()
 38  	msgRouter := NewMultiMsgRouter()
 39  	msgRouter.Start(ctx)
 40  	defer msgRouter.Stop()
 41  
 42  	openChanMsg := PeerMsg{
 43  		Message: &lnwire.OpenChannel{},
 44  	}
 45  	commitSigMsg := PeerMsg{
 46  		Message: &lnwire.CommitSig{},
 47  	}
 48  
 49  	errorMsg := PeerMsg{
 50  		Message: &lnwire.Error{},
 51  	}
 52  
 53  	// For this test, we'll have two endpoints, each with distinct names.
 54  	// One endpoint will only handle OpenChannel, while the other will
 55  	// handle the CommitSig message.
 56  	fundingEndpoint := &mockEndpoint{}
 57  	fundingEndpointName := "funding"
 58  	fundingEndpoint.On("Name").Return(fundingEndpointName)
 59  	fundingEndpoint.On("CanHandle", openChanMsg).Return(true)
 60  	fundingEndpoint.On("CanHandle", errorMsg).Return(false)
 61  	fundingEndpoint.On("CanHandle", commitSigMsg).Return(false)
 62  	fundingEndpoint.On("SendMessage", ctx, openChanMsg).Return(true)
 63  
 64  	commitEndpoint := &mockEndpoint{}
 65  	commitEndpointName := "commit"
 66  	commitEndpoint.On("Name").Return(commitEndpointName)
 67  	commitEndpoint.On("CanHandle", commitSigMsg).Return(true)
 68  	commitEndpoint.On("CanHandle", openChanMsg).Return(false)
 69  	commitEndpoint.On("CanHandle", errorMsg).Return(false)
 70  	commitEndpoint.On("SendMessage", ctx, commitSigMsg).Return(true)
 71  
 72  	t.Run("add endpoints", func(t *testing.T) {
 73  		// First, we'll add the funding endpoint to the router.
 74  		require.NoError(t, msgRouter.RegisterEndpoint(fundingEndpoint))
 75  
 76  		endpoints, err := msgRouter.endpoints().Unpack()
 77  		require.NoError(t, err)
 78  
 79  		// There should be a single endpoint registered.
 80  		require.Len(t, endpoints, 1)
 81  
 82  		// The name of the registered endpoint should be "funding".
 83  		require.Equal(
 84  			t, "funding", endpoints[fundingEndpointName].Name(),
 85  		)
 86  	})
 87  
 88  	t.Run("duplicate endpoint reject", func(t *testing.T) {
 89  		// Next, we'll attempt to add the funding endpoint again. This
 90  		// should return an ErrDuplicateEndpoint error.
 91  		require.ErrorIs(
 92  			t, msgRouter.RegisterEndpoint(fundingEndpoint),
 93  			ErrDuplicateEndpoint,
 94  		)
 95  	})
 96  
 97  	t.Run("route to endpoint", func(t *testing.T) {
 98  		// Next, we'll add our other endpoint, then attempt to route a
 99  		// message.
100  		require.NoError(t, msgRouter.RegisterEndpoint(commitEndpoint))
101  
102  		// If we try to route a message none of the endpoints know of,
103  		// we should get an error.
104  		require.ErrorIs(
105  			t, msgRouter.RouteMsg(errorMsg), ErrUnableToRouteMsg,
106  		)
107  
108  		fundingEndpoint.AssertCalled(t, "CanHandle", errorMsg)
109  		commitEndpoint.AssertCalled(t, "CanHandle", errorMsg)
110  
111  		// Next, we'll route the open channel message. Only the
112  		// fundingEndpoint should be used.
113  		require.NoError(t, msgRouter.RouteMsg(openChanMsg))
114  
115  		fundingEndpoint.AssertCalled(t, "CanHandle", openChanMsg)
116  		commitEndpoint.AssertCalled(t, "CanHandle", openChanMsg)
117  
118  		fundingEndpoint.AssertCalled(t, "SendMessage", ctx, openChanMsg)
119  		commitEndpoint.AssertNotCalled(
120  			t, "SendMessage", ctx, openChanMsg,
121  		)
122  
123  		// We'll do the same for the commit sig message.
124  		require.NoError(t, msgRouter.RouteMsg(commitSigMsg))
125  
126  		fundingEndpoint.AssertCalled(t, "CanHandle", commitSigMsg)
127  		commitEndpoint.AssertCalled(t, "CanHandle", commitSigMsg)
128  
129  		commitEndpoint.AssertCalled(t, "SendMessage", ctx, commitSigMsg)
130  		fundingEndpoint.AssertNotCalled(
131  			t, "SendMessage", ctx, commitSigMsg,
132  		)
133  	})
134  
135  	t.Run("remove endpoints", func(t *testing.T) {
136  		// Finally, we'll remove both endpoints.
137  		require.NoError(
138  			t, msgRouter.UnregisterEndpoint(fundingEndpointName),
139  		)
140  		require.NoError(
141  			t, msgRouter.UnregisterEndpoint(commitEndpointName),
142  		)
143  
144  		endpoints, err := msgRouter.endpoints().Unpack()
145  		require.NoError(t, err)
146  
147  		// There should be no endpoints registered.
148  		require.Len(t, endpoints, 0)
149  
150  		// Trying to route a message should fail.
151  		require.ErrorIs(
152  			t, msgRouter.RouteMsg(openChanMsg),
153  			ErrUnableToRouteMsg,
154  		)
155  		require.ErrorIs(
156  			t, msgRouter.RouteMsg(commitSigMsg),
157  			ErrUnableToRouteMsg,
158  		)
159  	})
160  
161  	commitEndpoint.AssertExpectations(t)
162  	fundingEndpoint.AssertExpectations(t)
163  }