diff --git a/lnwire/custom.go b/lnwire/custom.go index c995d1171..232a8be52 100644 --- a/lnwire/custom.go +++ b/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{ diff --git a/lnwire/message.go b/lnwire/message.go index ac57c8a29..2e71173e5 100644 --- a/lnwire/message.go +++ b/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, }