lnwire: allow overriding of protocol messages outside of custom range

Add the ability to specify messages < CustomRangeStart that will still
be treated like custom messages by lnd (rather than unknown ones). This
allows code external to lnd to handle protocol messages that are *not
yet known to lnd*.

lnwire/custom.go

@@ -2,14 +2,66 @@ package lnwire
 
 import (
 	"bytes"
-	"errors"
+	"fmt"
 	"io"
+	"sync"
 )
 
 // CustomTypeStart is the start of the custom type range for peer messages as
 // defined in BOLT 01.
 const CustomTypeStart MessageType = 32768
 
+var (
+	// customTypeOverride contains a set of message types < CustomTypeStart
+	// that lnd allows to be treated as custom messages. This allows us to
+	// override messages reserved for the protocol level and treat them as
+	// custom messages. This set of message types is stored as a global so
+	// that we do not need to pass around state when accounting for this
+	// set of messages in message creation.
+	//
+	// Note: This global is protected by the customTypeOverride mutex.
+	customTypeOverride map[MessageType]struct{}
+
+	// customTypeOverrideMtx manages concurrent access to
+	// customTypeOverride.
+	customTypeOverrideMtx sync.RWMutex
+)
+
+// SetCustomOverrides validates that the set of override types are outside of
+// the custom message range (there's no reason to override messages that are
+// already within the range), and updates the customTypeOverride global to hold
+// this set of message types. Note that this function will completely overwrite
+// the set of overrides, so should be called with the full set of types.
+func SetCustomOverrides(overrideTypes []uint16) error {
+	customTypeOverrideMtx.Lock()
+	defer customTypeOverrideMtx.Unlock()
+
+	customTypeOverride = make(map[MessageType]struct{}, len(overrideTypes))
+
+	for _, t := range overrideTypes {
+		msgType := MessageType(t)
+
+		if msgType >= CustomTypeStart {
+			return fmt.Errorf("can't override type: %v, already "+
+				"in custom range", t)
+		}
+
+		customTypeOverride[msgType] = struct{}{}
+	}
+
+	return nil
+}
+
+// IsCustomOverride returns a bool indicating whether the message type is one
+// of the protocol messages that we override for custom use.
+func IsCustomOverride(t MessageType) bool {
+	customTypeOverrideMtx.RLock()
+	defer customTypeOverrideMtx.RUnlock()
+
+	_, ok := customTypeOverride[t]
+
+	return ok
+}
 
 // Custom represents an application-defined wire message.
 type Custom struct {
@@ -21,10 +73,11 @@ type Custom struct {
 // interface.
 var _ Message = (*Custom)(nil)
 
-// NewCustom instanties a new custom message.
+// NewCustom instantiates a new custom message.
 func NewCustom(msgType MessageType, data []byte) (*Custom, error) {
-	if msgType < CustomTypeStart {
-		return nil, errors.New("msg type not in custom range")
+	if msgType < CustomTypeStart && !IsCustomOverride(msgType) {
+		return nil, fmt.Errorf("msg type: %d not in custom range: %v "+
+			"and not overridden", msgType, CustomTypeStart)
 	}
 
 	return &Custom{

lnwire/message.go

@@ -237,9 +237,16 @@ func makeEmptyMessage(msgType MessageType) (Message, error) {
 	case MsgGossipTimestampRange:
 		msg = &GossipTimestampRange{}
 	default:
-		if msgType < CustomTypeStart {
+		// If the message is not within our custom range and has not
+		// specifically been overridden, return an unknown message.
+		//
+		// Note that we do not allow custom message overrides to replace
+		// known message types, only protocol messages that are not yet
+		// known to lnd.
+		if msgType < CustomTypeStart && !IsCustomOverride(msgType) {
 			return nil, &UnknownMessage{msgType}
 		}
+
 		msg = &Custom{
 			Type: msgType,
 		}