diff --git a/lntemp/rpc/lnd.go b/lntemp/rpc/lnd.go index f6b688e04..41c6e31b0 100644 --- a/lntemp/rpc/lnd.go +++ b/lntemp/rpc/lnd.go @@ -648,3 +648,37 @@ func (h *HarnessRPC) SubscribeChannelEvents() ChannelEventsClient { return client } + +type CustomMessageClient lnrpc.Lightning_SubscribeCustomMessagesClient + +// SubscribeCustomMessages creates a subscription client for custom messages. +func (h *HarnessRPC) SubscribeCustomMessages() (CustomMessageClient, + context.CancelFunc) { + + ctxt, cancel := context.WithCancel(h.runCtx) + + req := &lnrpc.SubscribeCustomMessagesRequest{} + + // SubscribeCustomMessages needs to have the context alive for the + // entire test case as the returned client will be used for send and + // receive events stream. Thus we use runCtx here instead of a timeout + // context. + stream, err := h.LN.SubscribeCustomMessages(ctxt, req) + h.NoError(err, "SubscribeCustomMessages") + + return stream, cancel +} + +// SendCustomMessage makes a RPC call to the node's SendCustomMessage and +// returns the response. +func (h *HarnessRPC) SendCustomMessage( + req *lnrpc.SendCustomMessageRequest) *lnrpc.SendCustomMessageResponse { + + ctxt, cancel := context.WithTimeout(h.runCtx, DefaultTimeout) + defer cancel() + + resp, err := h.LN.SendCustomMessage(ctxt, req) + h.NoError(err, "SendCustomMessage") + + return resp +} diff --git a/lntest/itest/list_on_test.go b/lntest/itest/list_on_test.go index a6c695a5d..3bd12f9d3 100644 --- a/lntest/itest/list_on_test.go +++ b/lntest/itest/list_on_test.go @@ -497,4 +497,8 @@ var allTestCasesTemp = []*lntemp.TestCase{ Name: "open channel fee policy", TestFunc: testOpenChannelUpdateFeePolicy, }, + { + Name: "custom message", + TestFunc: testCustomMessage, + }, } diff --git a/lntest/itest/lnd_custom_message.go b/lntest/itest/lnd_custom_message.go index 1d9c36d45..5faf425a1 100644 --- a/lntest/itest/lnd_custom_message.go +++ b/lntest/itest/lnd_custom_message.go @@ -1,13 +1,11 @@ package itest import ( - "context" "fmt" - "sync" "time" "github.com/lightningnetwork/lnd/lnrpc" - "github.com/lightningnetwork/lnd/lntest" + "github.com/lightningnetwork/lnd/lntemp" "github.com/lightningnetwork/lnd/lnwire" "github.com/stretchr/testify/require" ) @@ -15,16 +13,8 @@ import ( // testCustomMessage tests sending and receiving of overridden custom message // types (within the message type range usually reserved for protocol messages) // via the send and subscribe custom message APIs. -func testCustomMessage(net *lntest.NetworkHarness, t *harnessTest) { - ctx, cancel := context.WithCancel(context.Background()) - var wg sync.WaitGroup - - // At the end of our test, cancel our context and wait for all - // goroutines to exit. - defer func() { - cancel() - wg.Wait() - }() +func testCustomMessage(ht *lntemp.HarnessTest) { + alice, bob := ht.Alice, ht.Bob var ( overrideType1 uint32 = 554 @@ -34,28 +24,19 @@ func testCustomMessage(net *lntest.NetworkHarness, t *harnessTest) { // Update Alice to accept custom protocol messages with type 1 but do // not allow Bob to handle them yet. - net.Alice.Cfg.ExtraArgs = append( - net.Alice.Cfg.ExtraArgs, + extraArgs := []string{ fmt.Sprintf(msgOverrideArg, overrideType1), - ) - require.NoError(t.t, net.RestartNode(net.Alice, nil, nil)) - - // Wait for Alice's server to be active after the restart before we - // try to subscribe to our message stream. - require.NoError(t.t, net.Alice.WaitUntilServerActive()) + } + ht.RestartNodeWithExtraArgs(alice, extraArgs) // Subscribe Alice to custom messages before we send any, so that we // don't miss any. - msgClient, err := net.Alice.LightningClient.SubscribeCustomMessages( - ctx, &lnrpc.SubscribeCustomMessagesRequest{}, - ) - require.NoError(t.t, err, "alice could not subscribe") + msgClient, cancel := alice.RPC.SubscribeCustomMessages() + defer cancel() // Create a channel to receive custom messages on. messages := make(chan *lnrpc.CustomMessage) - wg.Add(1) go func() { - defer wg.Done() for { // If we fail to receive, just exit. The test should // fail elsewhere if it doesn't get a message that it @@ -69,120 +50,108 @@ func testCustomMessage(net *lntest.NetworkHarness, t *harnessTest) { // test is shutting down. select { case messages <- msg: - case <-ctx.Done(): + case <-ht.Context().Done(): return } } }() // Connect alice and bob so that they can exchange messages. - net.EnsureConnected(t.t, net.Alice, net.Bob) + ht.EnsureConnected(alice, bob) // Create a custom message that is within our allowed range. msgType := uint32(lnwire.CustomTypeStart + 1) msgData := []byte{1, 2, 3} // Send it from Bob to Alice. - ctxt, _ := context.WithTimeout(ctx, defaultTimeout) - _, err = net.Bob.LightningClient.SendCustomMessage( - ctxt, &lnrpc.SendCustomMessageRequest{ - Peer: net.Alice.PubKey[:], - Type: msgType, - Data: msgData, - }, - ) - require.NoError(t.t, err, "bob could not send") + bobMsg := &lnrpc.SendCustomMessageRequest{ + Peer: alice.PubKey[:], + Type: msgType, + Data: msgData, + } + bob.RPC.SendCustomMessage(bobMsg) - // Wait for Alice to receive the message. It should come through because - // it is within our allowed range. + // Wait for Alice to receive the message. It should come through + // because it is within our allowed range. select { case msg := <-messages: - // Check our type and data and (sanity) check the peer we got it - // from. - require.Equal(t.t, msgType, msg.Type, "first msg type wrong") - require.Equal(t.t, msgData, msg.Data, "first msg data wrong") - require.Equal(t.t, net.Bob.PubKey[:], msg.Peer, "first msg "+ + // Check our type and data and (sanity) check the peer we got + // it from. + require.Equal(ht, msgType, msg.Type, "first msg type wrong") + require.Equal(ht, msgData, msg.Data, "first msg data wrong") + require.Equal(ht, bob.PubKey[:], msg.Peer, "first msg "+ "peer wrong") case <-time.After(defaultTimeout): - t.t.Fatalf("alice did not receive first custom message: %v", + ht.Fatalf("alice did not receive first custom message: %v", msgType) } // Try to send a message from Bob to Alice which has a message type // outside of the custom type range and assert that it fails. - ctxt, _ = context.WithTimeout(ctx, defaultTimeout) - _, err = net.Bob.LightningClient.SendCustomMessage( - ctxt, &lnrpc.SendCustomMessageRequest{ - Peer: net.Alice.PubKey[:], - Type: overrideType1, - Data: msgData, - }, - ) - require.Error(t.t, err, "bob should not be able to send type 1") + bobMsg = &lnrpc.SendCustomMessageRequest{ + Peer: alice.PubKey[:], + Type: overrideType1, + Data: msgData, + } + _, err := bob.RPC.LN.SendCustomMessage(ht.Context(), bobMsg) + require.Error(ht, err, "bob should not be able to send type 1") // Now, restart Bob with the ability to send two different custom // protocol messages. - net.Bob.Cfg.ExtraArgs = append( - net.Bob.Cfg.ExtraArgs, + extraArgs = []string{ fmt.Sprintf(msgOverrideArg, overrideType1), fmt.Sprintf(msgOverrideArg, overrideType2), - ) - require.NoError(t.t, net.RestartNode(net.Bob, nil, nil)) + } + ht.RestartNodeWithExtraArgs(bob, extraArgs) // Make sure Bob and Alice are connected after his restart. - net.EnsureConnected(t.t, net.Alice, net.Bob) + ht.EnsureConnected(alice, bob) // Send a message from Bob to Alice with a type that Bob is allowed to // send, but Alice will not handle as a custom message. - ctxt, _ = context.WithTimeout(ctx, defaultTimeout) - _, err = net.Bob.LightningClient.SendCustomMessage( - ctxt, &lnrpc.SendCustomMessageRequest{ - Peer: net.Alice.PubKey[:], - Type: overrideType2, - Data: msgData, - }, - ) - require.NoError(t.t, err, "bob should be able to send type 2") + bobMsg = &lnrpc.SendCustomMessageRequest{ + Peer: alice.PubKey[:], + Type: overrideType2, + Data: msgData, + } + bob.RPC.SendCustomMessage(bobMsg) // Do a quick check that Alice did not receive this message in her - // stream. Note that this is an instant check, so could miss the message - // being received. We'll also check below that she didn't get it, this - // is just a sanity check. + // stream. Note that this is an instant check, so could miss the + // message being received. We'll also check below that she didn't get + // it, this is just a sanity check. select { case msg := <-messages: - t.t.Fatalf("unexpected message: %v", msg) + ht.Fatalf("unexpected message: %v", msg) default: } // Finally, send a custom message with a type that Bob is allowed to // send and Alice is configured to receive. - ctxt, _ = context.WithTimeout(ctx, defaultTimeout) - _, err = net.Bob.LightningClient.SendCustomMessage( - ctxt, &lnrpc.SendCustomMessageRequest{ - Peer: net.Alice.PubKey[:], - Type: overrideType1, - Data: msgData, - }, - ) - require.NoError(t.t, err, "bob should be able to send type 1") + bobMsg = &lnrpc.SendCustomMessageRequest{ + Peer: alice.PubKey[:], + Type: overrideType1, + Data: msgData, + } + bob.RPC.SendCustomMessage(bobMsg) // Wait to receive a message from Bob. This check serves to ensure that // our message type 1 was delivered, and assert that the preceding one - // was not (we could have missed it in our check above). When we receive - // the second message, we know that the first one did not go through, - // because we expect our messages to deliver in order. + // was not (we could have missed it in our check above). When we + // receive the second message, we know that the first one did not go + // through, because we expect our messages to deliver in order. select { case msg := <-messages: - // Check our type and data and (sanity) check the peer we got it - // from. - require.Equal(t.t, overrideType1, msg.Type, "second message "+ + // Check our type and data and (sanity) check the peer we got + // it from. + require.Equal(ht, overrideType1, msg.Type, "second message "+ "type") - require.Equal(t.t, msgData, msg.Data, "second message data") - require.Equal(t.t, net.Bob.PubKey[:], msg.Peer, "second "+ + require.Equal(ht, msgData, msg.Data, "second message data") + require.Equal(ht, bob.PubKey[:], msg.Peer, "second "+ "message peer") case <-time.After(defaultTimeout): - t.t.Fatalf("alice did not receive second custom message") + ht.Fatalf("alice did not receive second custom message") } } diff --git a/lntest/itest/lnd_test_list_on_test.go b/lntest/itest/lnd_test_list_on_test.go index 97d5282a4..0783688bd 100644 --- a/lntest/itest/lnd_test_list_on_test.go +++ b/lntest/itest/lnd_test_list_on_test.go @@ -8,8 +8,4 @@ var allTestCases = []*testCase{ name: "async bidirectional payments", test: testBidirectionalAsyncPayments, }, - { - name: "custom messaging", - test: testCustomMessage, - }, }