Allow injecting the underlay type into messages. Fixes #222

Author: plorenz

impl.go

@@ -20,21 +20,23 @@ import (
 	"container/heap"
 	"crypto/x509"
 	"fmt"
-	"github.com/michaelquigley/pfxlog"
-	"github.com/openziti/foundation/v2/concurrenz"
-	"github.com/openziti/foundation/v2/info"
-	"github.com/openziti/foundation/v2/sequence"
-	"github.com/pkg/errors"
 	"io"
 	"net"
 	"sync"
 	"sync/atomic"
 	"time"
+
+	"github.com/michaelquigley/pfxlog"
+	"github.com/openziti/foundation/v2/concurrenz"
+	"github.com/openziti/foundation/v2/info"
+	"github.com/openziti/foundation/v2/sequence"
+	"github.com/pkg/errors"
 )
 
 const (
-	flagClosed    = 0
-	flagRxStarted = 1
+	flagClosed             = 0
+	flagRxStarted          = 1
+	flagInjectUnderlayType = 2
 )
 
 var connectionSeq = sequence.NewSequence()

message.go

@@ -21,10 +21,11 @@ import (
 	"context"
 	"encoding/binary"
 	"fmt"
-	"github.com/michaelquigley/pfxlog"
-	"github.com/pkg/errors"
 	"io"
 	"time"
+
+	"github.com/michaelquigley/pfxlog"
+	"github.com/pkg/errors"
 )
 
 /**
@@ -50,6 +51,7 @@ const (
 	IsGroupedHeader                 = 9
 	GroupSecretHeader               = 10
 	IsFirstGroupConnection          = 11
+	UnderlayTypeHeader              = 12
 
 	// Headers in the range 128-255 inclusive will be reflected when creating replies
 	ReflectedHeaderBitMask = 1 << 7

multi.go

@@ -40,6 +40,8 @@ type MultiChannelConfig struct {
 	UnderlayHandler UnderlayHandler
 	BindHandler     BindHandler
 	Underlay        Underlay
+
+	InjectUnderlayTypeIntoMessages bool
 }
 
 type senderContextImpl struct {
@@ -110,6 +112,8 @@ func NewMultiChannel(config *MultiChannelConfig) (MultiChannel, error) {
 		underlayHandler: config.UnderlayHandler,
 	}
 
+	impl.flags.Set(flagInjectUnderlayType, config.InjectUnderlayTypeIntoMessages)
+
 	impl.ownerId = config.Underlay.Id()
 	impl.certs.Store(config.Underlay.Certificates())
 	impl.headers.Store(config.Underlay.Headers())
@@ -495,6 +499,9 @@ func (self *multiChannelImpl) Rxer(underlay Underlay, notifier *CloseNotifier) {
 	log.Debug("started")
 	defer log.Debug("exited")
 
+	underlayType := GetUnderlayType(underlay)
+	injectType := self.flags.IsSet(flagInjectUnderlayType)
+
 	for {
 		m, err := underlay.Rx()
 		if err != nil {
@@ -508,6 +515,9 @@ func (self *multiChannelImpl) Rxer(underlay Underlay, notifier *CloseNotifier) {
 			return
 		}
 
+		if injectType {
+			m.Headers.PutStringHeader(UnderlayTypeHeader, underlayType)
+		}
 		self.Rx(m)
 	}
 }