lnrpc: receive custom message

peer/brontide.go

@@ -94,6 +94,11 @@ type newChannelMsg struct {
 	err     chan error
 }
 
+type customMsg struct {
+	peer [33]byte
+	msg  lnwire.Custom
+}
+
 // closeMsg is a wrapper struct around any wire messages that deal with the
 // cooperative channel closure negotiation process. This struct includes the
 // raw channel ID targeted along with the original message.
@@ -318,6 +323,10 @@ type Config struct {
 	// that is accumulated before signing a new commitment.
 	ChannelCommitBatchSize uint32
 
+	// HandleCustomMessage is called whenever a custom message is received
+	// from the peer.
+	HandleCustomMessage func(peer [33]byte, msg *lnwire.Custom) error
+
 	// Quit is the server's quit channel. If this is closed, we halt operation.
 	Quit chan struct{}
 }
@@ -1449,6 +1458,13 @@ out:
 
 			discStream.AddMsg(msg)
 
+		case *lnwire.Custom:
+			err := p.handleCustomMessage(msg)
+			if err != nil {
+				p.storeError(err)
+				peerLog.Errorf("peer: %v, %v", p, err)
+			}
+
 		default:
 			// If the message we received is unknown to us, store
 			// the type to track the failure.
@@ -1486,6 +1502,17 @@ out:
 	peerLog.Tracef("readHandler for peer %v done", p)
 }
 
+// handleCustomMessage handles the given custom message if a handler is
+// registered.
+func (p *Brontide) handleCustomMessage(msg *lnwire.Custom) error {
+	if p.cfg.HandleCustomMessage == nil {
+		return fmt.Errorf("no custom message handler for "+
+			"message type %v", uint16(msg.MsgType()))
+	}
+
+	return p.cfg.HandleCustomMessage(p.PubKey(), msg)
+}
+
 // isActiveChannel returns true if the provided channel id is active, otherwise
 // returns false.
 func (p *Brontide) isActiveChannel(chanID lnwire.ChannelID) bool {
@@ -1686,6 +1713,8 @@ func messageSummary(msg lnwire.Message) string {
 			time.Unix(int64(msg.FirstTimestamp), 0),
 			msg.TimestampRange)
 
+	case *lnwire.Custom:
+		return fmt.Sprintf("type=%d", msg.Type)
 	}
 
 	return ""
@@ -1714,8 +1743,15 @@ func (p *Brontide) logWireMessage(msg lnwire.Message, read bool) {
 			preposition = "from"
 		}
 
+		var msgType string
+		if msg.MsgType() < lnwire.CustomTypeStart {
+			msgType = msg.MsgType().String()
+		} else {
+			msgType = "custom"
+		}
+
 		return fmt.Sprintf("%v %v%s %v %s", summaryPrefix,
-			msg.MsgType(), summary, preposition, p)
+			msgType, summary, preposition, p)
 	}))
 
 	switch m := msg.(type) {

peer/brontide_test.go

@@ -1071,6 +1071,8 @@ func TestPeerCustomMessage(t *testing.T) {
 
 	mockConn := newMockConn(t, 1)
 
+	receivedCustomChan := make(chan *customMsg)
+
 	remoteKey := [33]byte{8}
 
 	notifier := &mock.ChainNotifier{
@@ -1092,6 +1094,15 @@ func TestPeerCustomMessage(t *testing.T) {
 		ReadPool:                      readPool,
 		Conn:                          mockConn,
 		ChainNotifier:                 notifier,
+		HandleCustomMessage: func(
+			peer [33]byte, msg *lnwire.Custom) error {
+
+			receivedCustomChan <- &customMsg{
+				peer: peer,
+				msg:  *msg,
+			}
+			return nil
+		},
 	})
 
 	// Set up the init sequence.
@@ -1117,7 +1128,9 @@ func TestPeerCustomMessage(t *testing.T) {
 	require.NoError(t, alicePeer.Start())
 
 	// Send a custom message.
-	customMsg, err := lnwire.NewCustom(lnwire.MessageType(40000), []byte{1, 2, 3})
+	customMsg, err := lnwire.NewCustom(
+		lnwire.MessageType(40000), []byte{1, 2, 3},
+	)
 	require.NoError(t, err)
 
 	require.NoError(t, alicePeer.SendMessageLazy(false, customMsg))
@@ -1125,4 +1138,18 @@ func TestPeerCustomMessage(t *testing.T) {
 	// Verify that it is passed down to the noise layer correctly.
 	writtenMsg := <-mockConn.writtenMessages
 	require.Equal(t, []byte{0x9c, 0x40, 0x1, 0x2, 0x3}, writtenMsg)
+
+	// Receive a custom message.
+	receivedCustomMsg, err := lnwire.NewCustom(
+		lnwire.MessageType(40001), []byte{4, 5, 6},
+	)
+	require.NoError(t, err)
+
+	receivedData := []byte{0x9c, 0x41, 0x4, 0x5, 0x6}
+	mockConn.readMessages <- receivedData
+
+	// Verify that it is propagated up to the custom message handler.
+	receivedCustom := <-receivedCustomChan
+	require.Equal(t, remoteKey, receivedCustom.peer)
+	require.Equal(t, receivedCustomMsg, &receivedCustom.msg)
 }