[relay] Negotiate the protocol version on the server side (#8231)

relay/metrics/realy.go

@@ -8,6 +8,8 @@ import (
 	log "github.com/sirupsen/logrus"
 	"go.opentelemetry.io/otel/attribute"
 	"go.opentelemetry.io/otel/metric"
+
+	"github.com/netbirdio/netbird/shared/relay/messages"
 )
 
 const (
@@ -17,16 +19,17 @@ const (
 type Metrics struct {
 	metric.Meter
 
-	TransferBytesSent  metric.Int64Counter
-	TransferBytesRecv  metric.Int64Counter
-	AuthenticationTime metric.Float64Histogram
-	PeerStoreTime      metric.Float64Histogram
-	peerReconnections  metric.Int64Counter
-	peers              metric.Int64UpDownCounter
-	peerActivityChan   chan string
-	peerLastActive     map[string]time.Time
-	mutexActivity      sync.Mutex
-	ctx                context.Context
+	TransferBytesSent    metric.Int64Counter
+	TransferBytesRecv    metric.Int64Counter
+	AuthenticationTime   metric.Float64Histogram
+	PeerStoreTime        metric.Float64Histogram
+	peerReconnections    metric.Int64Counter
+	peerProtocolVersions metric.Int64Counter
+	peers                metric.Int64UpDownCounter
+	peerActivityChan     chan string
+	peerLastActive       map[string]time.Time
+	mutexActivity        sync.Mutex
+	ctx                  context.Context
 }
 
 func NewMetrics(ctx context.Context, meter metric.Meter) (*Metrics, error) {
@@ -88,14 +91,22 @@ func NewMetrics(ctx context.Context, meter metric.Meter) (*Metrics, error) {
 		return nil, err
 	}
 
+	peerProtocolVersions, err := meter.Int64Counter("relay_peer_protocol_versions_total",
+		metric.WithDescription("Total number of peer connections by advertised and selected protocol version"),
+	)
+	if err != nil {
+		return nil, err
+	}
+
 	m := &Metrics{
-		Meter:              meter,
-		TransferBytesSent:  bytesSent,
-		TransferBytesRecv:  bytesRecv,
-		AuthenticationTime: authTime,
-		PeerStoreTime:      peerStoreTime,
-		peers:              peers,
-		peerReconnections:  peerReconnections,
+		Meter:                meter,
+		TransferBytesSent:    bytesSent,
+		TransferBytesRecv:    bytesRecv,
+		AuthenticationTime:   authTime,
+		PeerStoreTime:        peerStoreTime,
+		peers:                peers,
+		peerReconnections:    peerReconnections,
+		peerProtocolVersions: peerProtocolVersions,
 
 		ctx:              ctx,
 		peerActivityChan: make(chan string, 10),
@@ -151,6 +162,15 @@ func (m *Metrics) RecordPeerReconnection() {
 	m.peerReconnections.Add(m.ctx, 1)
 }
 
+// RecordProtocolVersion counts a new peer connection by the protocol version the client advertised
+// and the version the server selected for it
+func (m *Metrics) RecordProtocolVersion(advertised, selected messages.ProtocolVersion) {
+	m.peerProtocolVersions.Add(m.ctx, 1, metric.WithAttributes(
+		attribute.Int("advertised", int(advertised)),
+		attribute.Int("selected", int(selected)),
+	))
+}
+
 // PeerActivity increases the active connections
 func (m *Metrics) PeerActivity(peerID string) {
 	select {

relay/server/handshake.go

@@ -19,20 +19,9 @@ type Validator interface {
 	Validate(any) error
 }
 
-// preparedMsg contains the marshalled success response message
+// preparedMsg contains the marshalled success response message for every supported protocol version
 type preparedMsg struct {
-	responseAuthMsg []byte
-}
-
-func newPreparedMsg(instanceURL string) (*preparedMsg, error) {
-	ram, err := messages.MarshalAuthResponse(instanceURL)
-	if err != nil {
-		return nil, fmt.Errorf("failed to marshal auth response msg: %w", err)
-	}
-
-	return &preparedMsg{
-		responseAuthMsg: ram,
-	}, nil
+	responseAuthMsg map[messages.ProtocolVersion][]byte
 }
 
 type handshake struct {
@@ -40,7 +29,35 @@ type handshake struct {
 	validator   Validator
 	preparedMsg *preparedMsg
 
-	peerID *messages.PeerID
+	peerID            *messages.PeerID
+	advertisedVersion messages.ProtocolVersion
+	selectedVersion   messages.ProtocolVersion
+}
+
+func newPreparedMsg(instanceURL string) (*preparedMsg, error) {
+	responses := make(map[messages.ProtocolVersion][]byte, messages.CurrentProtocolVersion)
+	for version := messages.ProtocolV1; ; version++ {
+		ram, err := messages.NewEncoder(version).AuthResponse(instanceURL)
+		if err != nil {
+			return nil, fmt.Errorf("marshal auth response for version %d: %w", version, err)
+		}
+		responses[version] = ram
+		if version == messages.CurrentProtocolVersion {
+			break
+		}
+	}
+
+	return &preparedMsg{
+		responseAuthMsg: responses,
+	}, nil
+}
+
+func (p *preparedMsg) authResponse(version messages.ProtocolVersion) ([]byte, error) {
+	msg, ok := p.responseAuthMsg[version]
+	if !ok {
+		return nil, fmt.Errorf("no prepared auth response for protocol version %d", version)
+	}
+	return msg, nil
 }
 
 func (h *handshake) handshakeReceive(ctx context.Context) (*messages.PeerID, error) {
@@ -52,9 +69,9 @@ func (h *handshake) handshakeReceive(ctx context.Context) (*messages.PeerID, err
 
 	buf = buf[:n]
 
-	_, err = messages.ValidateVersion(buf)
+	advertised, err := messages.ReadVersion(buf)
 	if err != nil {
-		return nil, fmt.Errorf("validate version from %s: %w", h.conn.RemoteAddr(), err)
+		return nil, fmt.Errorf("read version from %s: %w", h.conn.RemoteAddr(), err)
 	}
 
 	msgType, err := messages.DetermineClientMessageType(buf)
@@ -70,12 +87,20 @@ func (h *handshake) handshakeReceive(ctx context.Context) (*messages.PeerID, err
 	if err != nil {
 		return peerID, err
 	}
+
 	h.peerID = peerID
+	h.advertisedVersion = advertised
+	h.selectedVersion = messages.SelectVersion(advertised, messages.CurrentProtocolVersion)
 	return peerID, nil
 }
 
 func (h *handshake) handshakeResponse(ctx context.Context) error {
-	if _, err := h.conn.Write(ctx, h.preparedMsg.responseAuthMsg); err != nil {
+	msg, err := h.preparedMsg.authResponse(h.selectedVersion)
+	if err != nil {
+		return err
+	}
+
+	if _, err := h.conn.Write(ctx, msg); err != nil {
 		return fmt.Errorf("handshake response write to %s (%s): %w", h.peerID, h.conn.RemoteAddr(), err)
 	}
 

relay/server/handshake_test.go

@@ -0,0 +1,378 @@
+package server
+
+import (
+	"context"
+	"encoding/binary"
+	"errors"
+	"net"
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"github.com/netbirdio/netbird/shared/relay/auth/allow"
+	"github.com/netbirdio/netbird/shared/relay/messages"
+)
+
+const (
+	testWaitTimeout = 2 * time.Second
+
+	// authTokenLenOffset is where the envelope's token length starts: header(2) + magic(4) + peer ID(36).
+	authTokenLenOffset = 2 + 4 + 36
+)
+
+// scriptedConn is a listener.Conn whose inbound frames are queued by the test and whose
+// outbound frames are recorded.
+type scriptedConn struct {
+	reads     chan []byte
+	closed    chan struct{}
+	closeOnce sync.Once
+	onWrite   func([]byte)
+
+	mu     sync.Mutex
+	writes [][]byte
+}
+
+type recordingValidator struct {
+	err error
+
+	mu       sync.Mutex
+	payloads [][]byte
+}
+
+func newScriptedConn() *scriptedConn {
+	return &scriptedConn{
+		reads:  make(chan []byte, 16),
+		closed: make(chan struct{}),
+	}
+}
+
+func (c *scriptedConn) Read(ctx context.Context, b []byte) (int, error) {
+	select {
+	case msg := <-c.reads:
+		return copy(b, msg), nil
+	case <-c.closed:
+		return 0, net.ErrClosed
+	case <-ctx.Done():
+		return 0, ctx.Err()
+	}
+}
+
+func (c *scriptedConn) Write(_ context.Context, b []byte) (int, error) {
+	select {
+	case <-c.closed:
+		return 0, net.ErrClosed
+	default:
+	}
+
+	frame := append([]byte(nil), b...)
+	c.mu.Lock()
+	c.writes = append(c.writes, frame)
+	c.mu.Unlock()
+
+	if c.onWrite != nil {
+		c.onWrite(frame)
+	}
+	return len(b), nil
+}
+
+func (c *scriptedConn) RemoteAddr() net.Addr {
+	return &net.TCPAddr{IP: net.IPv4(127, 0, 0, 1), Port: 1}
+}
+
+func (c *scriptedConn) Close() error {
+	c.closeOnce.Do(func() { close(c.closed) })
+	return nil
+}
+
+func (c *scriptedConn) Protocol() string {
+	return "test"
+}
+
+func (c *scriptedConn) deliver(t *testing.T, frame []byte) {
+	t.Helper()
+	select {
+	case c.reads <- frame:
+	case <-time.After(testWaitTimeout):
+		t.Fatal("scripted connection read queue is full")
+	}
+}
+
+func (c *scriptedConn) isClosed() bool {
+	select {
+	case <-c.closed:
+		return true
+	default:
+		return false
+	}
+}
+
+func (c *scriptedConn) frames() [][]byte {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	return append([][]byte(nil), c.writes...)
+}
+
+func (c *scriptedConn) waitFrames(t *testing.T, n int) [][]byte {
+	t.Helper()
+	deadline := time.Now().Add(testWaitTimeout)
+	for {
+		frames := c.frames()
+		if len(frames) >= n {
+			return frames
+		}
+		if time.Now().After(deadline) {
+			t.Fatalf("expected %d written frames, got %d", n, len(frames))
+		}
+		time.Sleep(5 * time.Millisecond)
+	}
+}
+
+func (c *scriptedConn) waitClosed(t *testing.T) {
+	t.Helper()
+	select {
+	case <-c.closed:
+	case <-time.After(testWaitTimeout):
+		t.Fatal("connection was not closed")
+	}
+}
+
+func (v *recordingValidator) Validate(credentials any) error {
+	v.mu.Lock()
+	defer v.mu.Unlock()
+
+	payload, _ := credentials.([]byte)
+	v.payloads = append(v.payloads, append([]byte(nil), payload...))
+	return v.err
+}
+
+func (v *recordingValidator) calls() [][]byte {
+	v.mu.Lock()
+	defer v.mu.Unlock()
+	return append([][]byte(nil), v.payloads...)
+}
+
+func TestRelay_AuthV1ResponseIsUnchanged(t *testing.T) {
+	validator := &recordingValidator{}
+	relay := newTestRelay(t, validator)
+	peerID := messages.HashID("v1-client")
+	token := []byte("v1-token")
+	auth, err := messages.MarshalAuthMsg(peerID, token)
+	require.NoError(t, err)
+
+	conn := newScriptedConn()
+	conn.deliver(t, auth)
+	relay.Accept(conn)
+
+	frames := conn.waitFrames(t, 1)
+	want, err := messages.MarshalAuthResponse(relay.instanceURL.String())
+	require.NoError(t, err)
+	assert.Equal(t, want, frames[0], "a v1 client must receive the same response bytes as today")
+	assert.Equal(t, [][]byte{token}, validator.calls(), "the validator must receive the v1 token bytes")
+	assert.False(t, conn.isClosed(), "an authenticated v1 connection must stay open")
+
+	item, ok := relay.store.Peer(peerID)
+	require.True(t, ok, "the peer must be stored after the response")
+	assert.Equal(t, messages.ProtocolV1, item.(*Peer).Version(), "a v1 client selects v1")
+}
+
+func TestRelay_AuthEnvelopeSelectsV1(t *testing.T) {
+	cases := []struct {
+		name    string
+		version messages.ProtocolVersion
+		ext     []byte
+	}{
+		{"v2 without extension", 2, nil},
+		{"v2 with json extension", 2, []byte(`{"cred2":"value"}`)},
+		{"v255 with opaque extension", 255, []byte("opaque-extension-bytes")},
+	}
+	for _, tc := range cases {
+		t.Run(tc.name, func(t *testing.T) {
+			validator := &recordingValidator{}
+			relay := newTestRelay(t, validator)
+			peerID := messages.HashID("future-client")
+			token := []byte("future-token")
+			auth, err := messages.MarshalAuthEnvelope(tc.version, peerID, token, tc.ext)
+			require.NoError(t, err)
+
+			conn := newScriptedConn()
+			conn.deliver(t, auth)
+			relay.Accept(conn)
+
+			frames := conn.waitFrames(t, 1)
+			want, err := messages.MarshalAuthResponse(relay.instanceURL.String())
+			require.NoError(t, err)
+			assert.Equal(t, want, frames[0], "the response must carry the selected v1 header")
+			assert.Equal(t, [][]byte{token}, validator.calls(), "the validator must receive the token without the extension")
+			assert.False(t, conn.isClosed(), "a higher advertised version must not be rejected")
+
+			item, ok := relay.store.Peer(peerID)
+			require.True(t, ok, "the peer must be stored after the response")
+			assert.Equal(t, messages.ProtocolV1, item.(*Peer).Version(), "the connection must use the selected v1")
+		})
+	}
+}
+
+func TestRelay_AuthRejectsInvalidFirstFrame(t *testing.T) {
+	peerID := messages.HashID("client")
+	validEnvelope, err := messages.MarshalAuthEnvelope(2, peerID, []byte("token"), []byte("ext"))
+	require.NoError(t, err)
+	validV1, err := messages.MarshalAuthMsg(peerID, []byte("token"))
+	require.NoError(t, err)
+	transportFirst, err := messages.MarshalTransportMsg(peerID, []byte("payload"))
+	require.NoError(t, err)
+
+	versionZero := append([]byte(nil), validV1...)
+	versionZero[0] = 0
+
+	badMagic := append([]byte(nil), validEnvelope...)
+	badMagic[2] ^= 0xFF
+
+	zeroTokenLen := append([]byte(nil), validEnvelope...)
+	binary.BigEndian.PutUint16(zeroTokenLen[authTokenLenOffset:], 0)
+
+	tokenLenPastEnd := append([]byte(nil), validEnvelope...)
+	binary.BigEndian.PutUint16(tokenLenPastEnd[authTokenLenOffset:], uint16(len(validEnvelope)))
+
+	cases := []struct {
+		name  string
+		frame []byte
+	}{
+		{"version zero", versionZero},
+		{"reserved hello type", []byte{byte(messages.ProtocolV1), byte(messages.MsgTypeHello)}},
+		{"unknown type before auth", []byte{byte(messages.ProtocolV1), 200}},
+		{"transport before auth", transportFirst},
+		{"bad magic in envelope", badMagic},
+		{"zero token length", zeroTokenLen},
+		{"token length past the frame", tokenLenPastEnd},
+		{"envelope truncated at the length field", validEnvelope[:authTokenLenOffset+1]},
+	}
+	for _, tc := range cases {
+		t.Run(tc.name, func(t *testing.T) {
+			validator := &recordingValidator{}
+			relay := newTestRelay(t, validator)
+
+			conn := newScriptedConn()
+			conn.deliver(t, tc.frame)
+			relay.Accept(conn)
+
+			assert.True(t, conn.isClosed(), "the connection must be closed")
+			assert.Empty(t, conn.frames(), "no response may be written")
+			assert.Empty(t, validator.calls(), "the validator must not run")
+			_, ok := relay.store.Peer(peerID)
+			assert.False(t, ok, "the peer must not be stored")
+		})
+	}
+}
+
+func TestRelay_AuthRejectsInvalidToken(t *testing.T) {
+	peerID := messages.HashID("client")
+	v1Auth, err := messages.MarshalAuthMsg(peerID, []byte("token"))
+	require.NoError(t, err)
+	envelopeAuth, err := messages.MarshalAuthEnvelope(2, peerID, []byte("token"), []byte(`{"ok":true}`))
+	require.NoError(t, err)
+
+	cases := []struct {
+		name  string
+		frame []byte
+	}{
+		{"v1 layout", v1Auth},
+		{"envelope with well-formed extension", envelopeAuth},
+	}
+	for _, tc := range cases {
+		t.Run(tc.name, func(t *testing.T) {
+			validator := &recordingValidator{err: errors.New("invalid token")}
+			relay := newTestRelay(t, validator)
+
+			conn := newScriptedConn()
+			conn.deliver(t, tc.frame)
+			relay.Accept(conn)
+
+			assert.True(t, conn.isClosed(), "the connection must be closed")
+			assert.Empty(t, conn.frames(), "no response may be written")
+			assert.Len(t, validator.calls(), 1, "the validator must run exactly once")
+			_, ok := relay.store.Peer(peerID)
+			assert.False(t, ok, "the peer must not be stored")
+		})
+	}
+}
+
+func TestRelay_AuthResponsePrecedesPublish(t *testing.T) {
+	relay := newTestRelay(t, &allow.Auth{})
+	peerID := messages.HashID("client")
+
+	onlineEvents := make(chan []messages.PeerID, 1)
+	watcher := relay.notifier.NewListener(func(ids []messages.PeerID) { onlineEvents <- ids }, func([]messages.PeerID) {})
+	t.Cleanup(func() { relay.notifier.RemoveListener(watcher) })
+	watcher.AddInterestedPeers([]messages.PeerID{peerID})
+
+	var storedAtResponse, notifiedAtResponse bool
+	conn := newScriptedConn()
+	conn.onWrite = func(frame []byte) {
+		if frame[1] != byte(messages.MsgTypeAuthResponse) {
+			return
+		}
+		_, storedAtResponse = relay.store.Peer(peerID)
+		notifiedAtResponse = len(onlineEvents) > 0
+	}
+
+	auth, err := messages.MarshalAuthMsg(peerID, []byte("token"))
+	require.NoError(t, err)
+	conn.deliver(t, auth)
+	relay.Accept(conn)
+
+	conn.waitFrames(t, 1)
+	assert.False(t, storedAtResponse, "the peer must not be in the store when the response is written")
+	assert.False(t, notifiedAtResponse, "no online notification may precede the response")
+
+	select {
+	case ids := <-onlineEvents:
+		assert.Equal(t, []messages.PeerID{peerID}, ids, "the peer must be published after the response")
+	case <-time.After(testWaitTimeout):
+		t.Fatal("the peer was not published after the response")
+	}
+}
+
+func TestRelay_ReconnectClosesPreviousConnection(t *testing.T) {
+	relay := newTestRelay(t, &allow.Auth{})
+	first, peerID := acceptPeer(t, relay, "client")
+
+	auth, err := messages.MarshalAuthMsg(peerID, []byte("token"))
+	require.NoError(t, err)
+	second := newScriptedConn()
+	second.deliver(t, auth)
+	relay.Accept(second)
+
+	second.waitFrames(t, 1)
+	first.waitClosed(t)
+	assert.False(t, second.isClosed(), "the reconnecting client must keep its new connection")
+	item, ok := relay.store.Peer(peerID)
+	require.True(t, ok, "the peer must stay stored after a reconnect")
+	assert.Same(t, second, item.(*Peer).conn, "the store must hold the new connection")
+}
+
+func newTestRelay(t *testing.T, validator Validator) *Relay {
+	t.Helper()
+	relay, err := NewRelay(Config{
+		ExposedAddress: "rel://127.0.0.1:1234",
+		AuthValidator:  validator,
+	})
+	require.NoError(t, err)
+	t.Cleanup(func() { relay.Shutdown(context.Background()) })
+	return relay
+}
+
+func acceptPeer(t *testing.T, relay *Relay, name string) (*scriptedConn, messages.PeerID) {
+	t.Helper()
+	peerID := messages.HashID(name)
+	auth, err := messages.MarshalAuthMsg(peerID, []byte("token"))
+	require.NoError(t, err)
+
+	conn := newScriptedConn()
+	conn.deliver(t, auth)
+	relay.Accept(conn)
+	conn.waitFrames(t, 1)
+	return conn, peerID
+}

relay/server/peer.go

@@ -27,6 +27,8 @@ type Peer struct {
 	metrics  *metrics.Metrics
 	log      *log.Entry
 	id       messages.PeerID
+	version  messages.ProtocolVersion
+	enc      messages.Encoder
 	conn     listener.Conn
 	connMu   sync.RWMutex
 	store    *store.Store
@@ -42,12 +44,14 @@ type Peer struct {
 }
 
 // NewPeer creates a new Peer instance and prepare custom logging
-func NewPeer(metrics *metrics.Metrics, id messages.PeerID, conn listener.Conn, store *store.Store, notifier *store.PeerNotifier) *Peer {
+func NewPeer(metrics *metrics.Metrics, id messages.PeerID, version messages.ProtocolVersion, conn listener.Conn, store *store.Store, notifier *store.PeerNotifier) *Peer {
 	ctx, cancel := context.WithCancel(context.Background())
 	p := &Peer{
 		metrics:   metrics,
 		log:       log.WithField("peer_id", id.String()),
 		id:        id,
+		version:   version,
+		enc:       messages.NewEncoder(version),
 		conn:      conn,
 		store:     store,
 		notifier:  notifier,
@@ -95,16 +99,15 @@ func (p *Peer) Work() {
 
 		msg := buf[:n]
 
-		_, err = messages.ValidateVersion(msg)
-		if err != nil {
-			p.log.Warnf("failed to validate protocol version: %s", err)
-			return
+		if err := messages.ValidateConnectionVersion(msg, p.version); err != nil {
+			p.log.Tracef("drop message: %s", err)
+			continue
 		}
 
 		msgType, err := messages.DetermineClientMessageType(msg)
 		if err != nil {
-			p.log.Errorf("failed to determine message type: %s", err)
-			return
+			p.log.Tracef("drop message: %s", err)
+			continue
 		}
 
 		p.handleMsgType(ctx, msgType, hc, n, msg)
@@ -115,6 +118,11 @@ func (p *Peer) ID() messages.PeerID {
 	return p.id
 }
 
+// Version returns the protocol version selected for this connection
+func (p *Peer) Version() messages.ProtocolVersion {
+	return p.version
+}
+
 func (p *Peer) handleMsgType(ctx context.Context, msgType messages.MsgType, hc *healthcheck.Sender, n int, msg []byte) {
 	switch msgType {
 	case messages.MsgTypeHealthCheck:
@@ -133,7 +141,7 @@ func (p *Peer) handleMsgType(ctx context.Context, msgType messages.MsgType, hc *
 	case messages.MsgTypeUnsubscribePeerState:
 		p.handleUnsubscribePeerState(msg)
 	default:
-		p.log.Warnf("received unexpected message type: %s", msgType)
+		p.log.Tracef("drop %s message: not allowed after authentication", msgType)
 	}
 }
 
@@ -149,7 +157,7 @@ func (p *Peer) Write(ctx context.Context, b []byte) (int, error) {
 func (p *Peer) CloseGracefully(ctx context.Context) {
 	p.connMu.Lock()
 	defer p.connMu.Unlock()
-	err := p.writeWithTimeout(ctx, messages.MarshalCloseMsg())
+	err := p.writeWithTimeout(ctx, p.enc.CloseMsg())
 	if err != nil {
 		p.log.Errorf("failed to send close message to peer: %s", p.String())
 	}
@@ -187,7 +195,7 @@ func (p *Peer) handleHealthcheckEvents(ctx context.Context, hc *healthcheck.Send
 	for {
 		select {
 		case <-hc.HealthCheck:
-			_, err := p.Write(ctx, messages.MarshalHealthcheck())
+			_, err := p.Write(ctx, p.enc.Healthcheck())
 			if err != nil {
 				p.log.Errorf("failed to send healthcheck message: %s", err)
 				return
@@ -220,7 +228,7 @@ func (p *Peer) handleTransportMsg(msg []byte) {
 	}
 	dp := item.(*Peer)
 
-	err = messages.UpdateTransportMsg(msg, p.id)
+	err = messages.UpdateTransportMsg(msg, dp.version, p.id)
 	if err != nil {
 		p.log.Errorf("failed to update transport message: %s", err)
 		return
@@ -267,7 +275,7 @@ func (p *Peer) handleUnsubscribePeerState(msg []byte) {
 }
 
 func (p *Peer) sendPeersOnline(peers []messages.PeerID) {
-	msgs, err := messages.MarshalPeersOnline(peers)
+	msgs, err := p.enc.PeersOnline(peers)
 	if err != nil {
 		p.log.Errorf("failed to marshal peer location message: %s", err)
 		return
@@ -284,7 +292,7 @@ func (p *Peer) sendPeersWentOffline(peers []messages.PeerID) {
 	p.notificationMutex.Lock()
 	defer p.notificationMutex.Unlock()
 
-	msgs, err := messages.MarshalPeersWentOffline(peers)
+	msgs, err := p.enc.PeersWentOffline(peers)
 	if err != nil {
 		p.log.Errorf("failed to marshal peer location message: %s", err)
 		return

relay/server/peer_test.go

@@ -0,0 +1,99 @@
+package server
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"github.com/netbirdio/netbird/shared/relay/auth/allow"
+	"github.com/netbirdio/netbird/shared/relay/messages"
+)
+
+func TestPeer_DropsFramesItCannotHandle(t *testing.T) {
+	secondAuth, err := messages.MarshalAuthMsg(messages.HashID("alice"), []byte("token"))
+	require.NoError(t, err)
+
+	cases := []struct {
+		name  string
+		frame []byte
+	}{
+		{"unknown message type", []byte{byte(messages.ProtocolV1), 200, 0xDE, 0xAD}},
+		{"version above the selected one", []byte{2, byte(messages.MsgTypeHealthCheck)}},
+		{"version zero", []byte{0, byte(messages.MsgTypeHealthCheck)}},
+		{"header shorter than two bytes", []byte{byte(messages.ProtocolV1)}},
+		{"reserved hello type", []byte{byte(messages.ProtocolV1), byte(messages.MsgTypeHello)}},
+		{"server to client peers online type", []byte{byte(messages.ProtocolV1), byte(messages.MsgTypePeersOnline)}},
+		{"auth after authentication", secondAuth},
+	}
+	for _, tc := range cases {
+		t.Run(tc.name, func(t *testing.T) {
+			relay := newTestRelay(t, &allow.Auth{})
+			alice, _ := acceptPeer(t, relay, "alice")
+			_, bobID := acceptPeer(t, relay, "bob")
+
+			alice.deliver(t, tc.frame)
+			subscribe, err := messages.MarshalSubPeerStateMsg([]messages.PeerID{bobID})
+			require.NoError(t, err)
+			alice.deliver(t, subscribe[0])
+
+			// The subscription answer proves the loop survived the dropped frame before it.
+			frames := alice.waitFrames(t, 2)
+			assert.False(t, alice.isClosed(), "a dropped frame must not close the connection")
+			online := frames[1]
+			assert.Equal(t, byte(messages.ProtocolV1), online[0], "notifications must carry the connection version")
+			assert.Equal(t, byte(messages.MsgTypePeersOnline), online[1], "the subscription must be answered")
+			ids, err := messages.UnmarshalPeersOnlineMsg(online)
+			require.NoError(t, err)
+			assert.Equal(t, []messages.PeerID{bobID}, ids, "the online peer must be reported")
+		})
+	}
+}
+
+func TestPeer_ClosesOnClientCloseMessage(t *testing.T) {
+	relay := newTestRelay(t, &allow.Auth{})
+	conn, peerID := acceptPeer(t, relay, "alice")
+
+	conn.deliver(t, messages.MarshalCloseMsg())
+
+	conn.waitClosed(t)
+	assert.Len(t, conn.frames(), 1, "only the auth response may be written")
+	assert.Eventually(t, func() bool {
+		_, ok := relay.store.Peer(peerID)
+		return !ok
+	}, testWaitTimeout, 5*time.Millisecond, "the peer must be removed from the store")
+}
+
+func TestPeer_ForwardsTransportWithSenderAndDestinationVersion(t *testing.T) {
+	relay := newTestRelay(t, &allow.Auth{})
+	alice, aliceID := acceptPeer(t, relay, "alice")
+	bob, bobID := acceptPeer(t, relay, "bob")
+
+	payload := []byte("hello bob")
+	transport, err := messages.MarshalTransportMsg(bobID, payload)
+	require.NoError(t, err)
+	alice.deliver(t, transport)
+
+	frames := bob.waitFrames(t, 2)
+	forwarded := frames[1]
+	assert.Equal(t, byte(messages.ProtocolV1), forwarded[0], "the forwarded frame must carry the destination's version")
+	assert.Equal(t, byte(messages.MsgTypeTransport), forwarded[1], "the forwarded frame must stay a transport message")
+	senderID, gotPayload, err := messages.UnmarshalTransportMsg(forwarded)
+	require.NoError(t, err)
+	assert.Equal(t, aliceID, *senderID, "the peer ID must be rewritten to the sender")
+	assert.Equal(t, payload, gotPayload, "the payload must be preserved")
+	assert.False(t, alice.isClosed(), "the sender must stay connected")
+}
+
+func TestPeer_GracefulCloseCarriesConnectionVersion(t *testing.T) {
+	relay := newTestRelay(t, &allow.Auth{})
+	conn, _ := acceptPeer(t, relay, "alice")
+
+	relay.Shutdown(context.Background())
+
+	frames := conn.waitFrames(t, 2)
+	assert.Equal(t, messages.NewEncoder(messages.ProtocolV1).CloseMsg(), frames[1], "the close message must carry the connection version")
+	assert.True(t, conn.isClosed(), "the connection must be closed after the close message")
+}

relay/server/relay.go

@@ -150,8 +150,17 @@ func (r *Relay) Accept(conn listener.Conn) {
 		return
 	}
 
-	peer := NewPeer(r.metrics, *peerID, conn, r.store, r.notifier)
-	peer.log.Infof("peer connected from: %s", conn.RemoteAddr())
+	if err := h.handshakeResponse(hsCtx); err != nil {
+		log.Errorf("failed to send handshake response, close connection: %s", err)
+		if cErr := conn.Close(); cErr != nil {
+			log.Errorf("failed to close connection, %s: %s", conn.RemoteAddr(), cErr)
+		}
+		return
+	}
+	r.metrics.RecordAuthenticationTime(time.Since(acceptTime))
+
+	peer := NewPeer(r.metrics, *peerID, h.selectedVersion, conn, r.store, r.notifier)
+	peer.log.Infof("peer connected from: %s, advertised protocol version: %d, selected: %d", conn.RemoteAddr(), h.advertisedVersion, h.selectedVersion)
 	storeTime := time.Now()
 	if isReconnection := r.store.AddPeer(peer); isReconnection {
 		r.metrics.RecordPeerReconnection()
@@ -161,6 +170,7 @@ func (r *Relay) Accept(conn listener.Conn) {
 	transport := conn.Protocol()
 	r.metrics.RecordPeerStoreTime(time.Since(storeTime))
 	r.metrics.PeerConnected(peer.String(), transport)
+	r.metrics.RecordProtocolVersion(h.advertisedVersion, h.selectedVersion)
 	go func() {
 		peer.Work()
 		if deleted := r.store.DeletePeer(peer); deleted {
@@ -169,12 +179,6 @@ func (r *Relay) Accept(conn listener.Conn) {
 		peer.log.Debugf("relay connection closed")
 		r.metrics.PeerDisconnected(peer.String(), transport)
 	}()
-
-	if err := h.handshakeResponse(hsCtx); err != nil {
-		log.Errorf("failed to send handshake response, close peer: %s", err)
-		peer.Close()
-	}
-	r.metrics.RecordAuthenticationTime(time.Since(acceptTime))
 }
 
 // Shutdown closes the relay server

relay/server/ws_version_test.go

@@ -0,0 +1,97 @@
+package server
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"github.com/coder/websocket"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"github.com/netbirdio/netbird/shared/relay"
+	"github.com/netbirdio/netbird/shared/relay/messages"
+)
+
+func TestServer_FutureClientOverWebSocket(t *testing.T) {
+	addr := freeAddress(t)
+	srv := newTestServer(t, addr)
+	errChan := make(chan error, 1)
+	go func() {
+		errChan <- srv.Listen(ListenerConfig{Address: addr})
+	}()
+	waitForListeners(t, srv, errChan)
+
+	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+	defer cancel()
+
+	instanceURL := srv.InstanceURL()
+	want, err := messages.MarshalAuthResponse(instanceURL.String())
+	require.NoError(t, err)
+
+	// A synthetic client advertising v2 with an extension area that a v1 server does not read.
+	future := dialRelayWebSocket(t, ctx, addr)
+	futureID := messages.HashID("future-client")
+	futureAuth, err := messages.MarshalAuthEnvelope(2, futureID, []byte("future-token"), []byte(`{"cred2":"value"}`))
+	require.NoError(t, err)
+	require.NoError(t, future.Write(ctx, websocket.MessageBinary, futureAuth))
+	assert.Equal(t, want, readFrame(t, ctx, future), "a v2 client must receive the v1 response")
+
+	legacy := dialRelayWebSocket(t, ctx, addr)
+	legacyID := messages.HashID("legacy-client")
+	legacyAuth, err := messages.MarshalAuthMsg(legacyID, []byte("legacy-token"))
+	require.NoError(t, err)
+	require.NoError(t, legacy.Write(ctx, websocket.MessageBinary, legacyAuth))
+	assert.Equal(t, want, readFrame(t, ctx, legacy), "a v1 client must receive the response it gets today")
+
+	toLegacy, err := messages.MarshalTransportMsg(legacyID, []byte("from future"))
+	require.NoError(t, err)
+	require.NoError(t, future.Write(ctx, websocket.MessageBinary, toLegacy))
+	assertTransportFrame(t, readFrame(t, ctx, legacy), futureID, "from future")
+
+	toFuture, err := messages.MarshalTransportMsg(futureID, []byte("from legacy"))
+	require.NoError(t, err)
+	require.NoError(t, legacy.Write(ctx, websocket.MessageBinary, toFuture))
+	assertTransportFrame(t, readFrame(t, ctx, future), legacyID, "from legacy")
+
+	// After the response the connection is v1. A frame with another version is dropped and the
+	// connection stays usable, which the transport message right after it proves.
+	require.NoError(t, future.Write(ctx, websocket.MessageBinary, []byte{2, byte(messages.MsgTypeHealthCheck)}))
+	toLegacyAgain, err := messages.MarshalTransportMsg(legacyID, []byte("after the dropped frame"))
+	require.NoError(t, err)
+	require.NoError(t, future.Write(ctx, websocket.MessageBinary, toLegacyAgain))
+	assertTransportFrame(t, readFrame(t, ctx, legacy), futureID, "after the dropped frame")
+
+	require.NoError(t, future.Write(ctx, websocket.MessageBinary, messages.MarshalCloseMsg()))
+	require.NoError(t, legacy.Write(ctx, websocket.MessageBinary, messages.MarshalCloseMsg()))
+}
+
+func dialRelayWebSocket(t *testing.T, ctx context.Context, addr string) *websocket.Conn {
+	t.Helper()
+	conn, resp, err := websocket.Dial(ctx, "ws://"+addr+relay.WebSocketURLPath, nil) //nolint:bodyclose
+	require.NoError(t, err)
+	if resp != nil && resp.Body != nil {
+		_ = resp.Body.Close()
+	}
+	t.Cleanup(func() { _ = conn.CloseNow() })
+	return conn
+}
+
+func readFrame(t *testing.T, ctx context.Context, conn *websocket.Conn) []byte {
+	t.Helper()
+	msgType, frame, err := conn.Read(ctx)
+	require.NoError(t, err)
+	require.Equal(t, websocket.MessageBinary, msgType, "relay frames are binary messages")
+	return frame
+}
+
+func assertTransportFrame(t *testing.T, frame []byte, sender messages.PeerID, payload string) {
+	t.Helper()
+	require.GreaterOrEqual(t, len(frame), 2, "a frame must carry the protocol header")
+	assert.Equal(t, byte(messages.ProtocolV1), frame[0], "the forwarded frame must carry the destination's version")
+	assert.Equal(t, byte(messages.MsgTypeTransport), frame[1], "the forwarded frame must be a transport message")
+	id, got, err := messages.UnmarshalTransportMsg(frame)
+	require.NoError(t, err)
+	assert.Equal(t, sender, *id, "the peer ID must be rewritten to the sender")
+	assert.Equal(t, payload, string(got), "the payload must be preserved")
+}

shared/relay/messages/message.go

@@ -2,16 +2,27 @@ package messages
 
 import (
 	"bytes"
+	"encoding/binary"
 	"errors"
 	"fmt"
+	"math"
 )
 
 const (
-	MaxHandshakeSize     = 212
+	// TODO: MaxHandshakeSize is the buffer size of the first read on every connection before
+	// authentication. Weigh it against memory use under many unauthenticated connections; a
+	// smaller limit may still fit the extension area.
+	MaxHandshakeSize     = 8192
 	MaxHandshakeRespSize = 8192
 	MaxMessageSize       = 8820
 
-	CurrentProtocolVersion = 1
+	ProtocolV1 ProtocolVersion = 1
+	// CurrentProtocolVersion is the highest protocol version this build supports.
+	CurrentProtocolVersion = ProtocolV1
+
+	// authEnvelopeMinVersion is the first advertised version whose Auth message carries a
+	// length-prefixed token followed by an extension area.
+	authEnvelopeMinVersion ProtocolVersion = 2
 
 	MsgTypeUnknown MsgType = 0
 	// MsgTypeHello and MsgTypeHelloResponse are the removed legacy handshake
@@ -38,10 +49,13 @@ const (
 
 	// auth message
 	sizeOfMagicByte     = 4
+	sizeOfTokenLen      = 2
 	headerSizeAuth      = sizeOfMagicByte + peerIDSize
 	offsetMagicByte     = sizeOfProtoHeader
 	offsetAuthPeerID    = sizeOfProtoHeader + sizeOfMagicByte
 	headerTotalSizeAuth = sizeOfProtoHeader + headerSizeAuth
+	offsetAuthTokenLen  = headerTotalSizeAuth
+	offsetAuthToken     = offsetAuthTokenLen + sizeOfTokenLen
 
 	// transport
 	headerSizeTransport      = peerIDSize
@@ -52,6 +66,10 @@ const (
 var (
 	ErrInvalidMessageLength = errors.New("invalid message length")
 	ErrUnsupportedVersion   = errors.New("unsupported version")
+	ErrInvalidTokenLength   = errors.New("invalid token length")
+	ErrReservedMsgType      = errors.New("reserved message type")
+	ErrMsgTypeDirection     = errors.New("message type not valid from client")
+	ErrUnknownMsgType       = errors.New("unknown message type")
 
 	magicHeader = []byte{0x21, 0x12, 0xA4, 0x42}
 
@@ -60,6 +78,19 @@ var (
 
 type MsgType byte
 
+// ProtocolVersion is the relay wire protocol version carried in the first byte of every frame.
+type ProtocolVersion uint8
+
+// Encoder marshals server-originated messages with the protocol version selected for one connection.
+type Encoder struct {
+	version ProtocolVersion
+}
+
+// NewEncoder returns an Encoder that stamps every message with the given protocol version.
+func NewEncoder(version ProtocolVersion) Encoder {
+	return Encoder{version: version}
+}
+
 func (m MsgType) String() string {
 	switch m {
 	case MsgTypeHello:
@@ -89,18 +120,99 @@ func (m MsgType) String() string {
 	}
 }
 
-// ValidateVersion checks if the given version is supported by the protocol
-func ValidateVersion(msg []byte) (int, error) {
+// Version returns the protocol version the Encoder stamps on its messages.
+func (e Encoder) Version() ProtocolVersion {
+	return e.version
+}
+
+// AuthResponse creates the response to a successful Auth message. The payload is the instance URL
+// of the server, which the clients use to choose a common Relay server.
+func (e Encoder) AuthResponse(address string) ([]byte, error) {
+	ab := []byte(address)
+	msg := make([]byte, sizeOfProtoHeader, sizeOfProtoHeader+len(ab))
+
+	msg[0] = byte(e.version)
+	msg[1] = byte(MsgTypeAuthResponse)
+
+	msg = append(msg, ab...)
+
+	if len(msg) > MaxHandshakeRespSize {
+		return nil, fmt.Errorf("invalid message length: %d", len(msg))
+	}
+
+	return msg, nil
+}
+
+// CloseMsg creates a close message for a graceful disconnect.
+func (e Encoder) CloseMsg() []byte {
+	return []byte{
+		byte(e.version),
+		byte(MsgTypeClose),
+	}
+}
+
+// Healthcheck creates a health check message.
+func (e Encoder) Healthcheck() []byte {
+	return []byte{
+		byte(e.version),
+		byte(MsgTypeHealthCheck),
+	}
+}
+
+// PeersOnline creates the peers online notification messages.
+func (e Encoder) PeersOnline(ids []PeerID) ([][]byte, error) {
+	return marshalPeerIDs(e.version, ids, byte(MsgTypePeersOnline))
+}
+
+// PeersWentOffline creates the peers went offline notification messages.
+func (e Encoder) PeersWentOffline(ids []PeerID) ([][]byte, error) {
+	return marshalPeerIDs(e.version, ids, byte(MsgTypePeersWentOffline))
+}
+
+// ReadVersion returns the protocol version carried in the frame header. Zero is invalid.
+func ReadVersion(msg []byte) (ProtocolVersion, error) {
 	if len(msg) < sizeOfProtoHeader {
 		return 0, ErrInvalidMessageLength
 	}
-	version := int(msg[0])
-	if version != CurrentProtocolVersion {
+	version := ProtocolVersion(msg[0])
+	if version == 0 {
 		return 0, fmt.Errorf("%d: %w", version, ErrUnsupportedVersion)
 	}
 	return version, nil
 }
 
+// SelectVersion returns the highest protocol version both sides support.
+func SelectVersion(advertised, supported ProtocolVersion) ProtocolVersion {
+	if advertised < supported {
+		return advertised
+	}
+	return supported
+}
+
+// ValidateVersion checks that the frame carries the protocol version of this build.
+func ValidateVersion(msg []byte) (int, error) {
+	version, err := ReadVersion(msg)
+	if err != nil {
+		return 0, err
+	}
+	if version != CurrentProtocolVersion {
+		return 0, fmt.Errorf("%d: %w", version, ErrUnsupportedVersion)
+	}
+	return int(version), nil
+}
+
+// ValidateConnectionVersion checks that the frame carries the version selected for the connection.
+func ValidateConnectionVersion(msg []byte, selected ProtocolVersion) error {
+	version, err := ReadVersion(msg)
+	if err != nil {
+		return err
+	}
+	if version != selected {
+		return fmt.Errorf("%d instead of selected %d: %w", version, selected, ErrUnsupportedVersion)
+	}
+	return nil
+}
+
 // DetermineClientMessageType determines the message type from the first the message
 func DetermineClientMessageType(msg []byte) (MsgType, error) {
 	if len(msg) < sizeOfProtoHeader {
@@ -117,8 +229,12 @@ func DetermineClientMessageType(msg []byte) (MsgType, error) {
 		MsgTypeSubscribePeerState,
 		MsgTypeUnsubscribePeerState:
 		return msgType, nil
+	case MsgTypeHello, MsgTypeHelloResponse:
+		return MsgTypeUnknown, fmt.Errorf("%d: %w", msgType, ErrReservedMsgType)
+	case MsgTypeAuthResponse, MsgTypePeersOnline, MsgTypePeersWentOffline:
+		return MsgTypeUnknown, fmt.Errorf("%d: %w", msgType, ErrMsgTypeDirection)
 	default:
-		return MsgTypeUnknown, fmt.Errorf("invalid msg type %d", msgType)
+		return MsgTypeUnknown, fmt.Errorf("%d: %w", msgType, ErrUnknownMsgType)
 	}
 }
 
@@ -162,8 +278,41 @@ func MarshalAuthMsg(peerID PeerID, authPayload []byte) ([]byte, error) {
 	return msg, nil
 }
 
-// UnmarshalAuthMsg extracts peerID and the auth payload from the message
+// MarshalAuthEnvelope creates the Auth message layout used by clients advertising protocol
+// version 2 or higher: the token is length-prefixed and followed by an extension area that
+// servers without support for it ignore.
+func MarshalAuthEnvelope(version ProtocolVersion, peerID PeerID, token, ext []byte) ([]byte, error) {
+	if version < authEnvelopeMinVersion {
+		return nil, fmt.Errorf("protocol version %d does not use the auth envelope", version)
+	}
+	if len(token) == 0 || len(token) > math.MaxUint16 {
+		return nil, fmt.Errorf("%d: %w", len(token), ErrInvalidTokenLength)
+	}
+
+	size := offsetAuthToken + len(token) + len(ext)
+	if size > MaxHandshakeSize {
+		return nil, fmt.Errorf("too large auth message: %d", size)
+	}
+
+	msg := make([]byte, size)
+	msg[0] = byte(version)
+	msg[1] = byte(MsgTypeAuth)
+	copy(msg[offsetMagicByte:], magicHeader)
+	copy(msg[offsetAuthPeerID:], peerID[:])
+	binary.BigEndian.PutUint16(msg[offsetAuthTokenLen:], uint16(len(token)))
+	copy(msg[offsetAuthToken:], token)
+	copy(msg[offsetAuthToken+len(token):], ext)
+	return msg, nil
+}
+
+// UnmarshalAuthMsg extracts the peer ID and the auth token from an Auth message of any
+// protocol version. The extension area of the envelope layout is not returned.
 func UnmarshalAuthMsg(msg []byte) (*PeerID, []byte, error) {
+	version, err := ReadVersion(msg)
+	if err != nil {
+		return nil, nil, err
+	}
+
 	if len(msg) < headerTotalSizeAuth {
 		return nil, nil, ErrInvalidMessageLength
 	}
@@ -174,7 +323,15 @@ func UnmarshalAuthMsg(msg []byte) (*PeerID, []byte, error) {
 	}
 
 	peerID := PeerID(msg[offsetAuthPeerID:headerTotalSizeAuth])
-	return &peerID, msg[headerTotalSizeAuth:], nil
+	if version < authEnvelopeMinVersion {
+		return &peerID, msg[headerTotalSizeAuth:], nil
+	}
+
+	token, err := unmarshalAuthEnvelopeToken(msg)
+	if err != nil {
+		return nil, nil, err
+	}
+	return &peerID, token, nil
 }
 
 // MarshalAuthResponse creates a response message to the auth.
@@ -182,19 +339,7 @@ func UnmarshalAuthMsg(msg []byte) (*PeerID, []byte, error) {
 // instance URL. This URL will be used by choose the common Relay server in case if the peers are in different Relay
 // servers.
 func MarshalAuthResponse(address string) ([]byte, error) {
-	ab := []byte(address)
-	msg := make([]byte, sizeOfProtoHeader, sizeOfProtoHeader+len(ab))
-
-	msg[0] = byte(CurrentProtocolVersion)
-	msg[1] = byte(MsgTypeAuthResponse)
-
-	msg = append(msg, ab...)
-
-	if len(msg) > MaxHandshakeRespSize {
-		return nil, fmt.Errorf("invalid message length: %d", len(msg))
-	}
-
-	return msg, nil
+	return NewEncoder(CurrentProtocolVersion).AuthResponse(address)
 }
 
 // UnmarshalAuthResponse it is a confirmation message to auth success
@@ -209,10 +354,7 @@ func UnmarshalAuthResponse(msg []byte) (string, error) {
 // The close message is used to close the connection gracefully between the client and the server. The server and the
 // client can send this message. After receiving this message, the server or client will close the connection.
 func MarshalCloseMsg() []byte {
-	return []byte{
-		byte(CurrentProtocolVersion),
-		byte(MsgTypeClose),
-	}
+	return NewEncoder(CurrentProtocolVersion).CloseMsg()
 }
 
 // MarshalTransportMsg creates a transport message.
@@ -252,13 +394,14 @@ func UnmarshalTransportID(buf []byte) (*PeerID, error) {
 	return &id, nil
 }
 
-// UpdateTransportMsg updates the peerID in the transport message.
-// With this function the server can reuse the given byte slice to update the peerID in the transport message. So do
-// need to allocate a new byte slice.
-func UpdateTransportMsg(msg []byte, peerID PeerID) error {
+// UpdateTransportMsg rewrites the transport message in place for forwarding: the header carries
+// the protocol version of the destination connection and the peer ID field the sender's ID.
+// The payload bytes are left untouched, so no new byte slice is allocated.
+func UpdateTransportMsg(msg []byte, version ProtocolVersion, peerID PeerID) error {
 	if len(msg) < offsetTransportID+peerIDSize {
 		return ErrInvalidMessageLength
 	}
+	msg[0] = byte(version)
 	copy(msg[offsetTransportID:], peerID[:])
 	return nil
 }
@@ -269,3 +412,20 @@ func UpdateTransportMsg(msg []byte, peerID PeerID) error {
 func MarshalHealthcheck() []byte {
 	return healthCheckMsg
 }
+
+func unmarshalAuthEnvelopeToken(msg []byte) ([]byte, error) {
+	if len(msg) < offsetAuthToken {
+		return nil, ErrInvalidMessageLength
+	}
+
+	tokenLen := int(binary.BigEndian.Uint16(msg[offsetAuthTokenLen:offsetAuthToken]))
+	if tokenLen == 0 {
+		return nil, fmt.Errorf("%d: %w", tokenLen, ErrInvalidTokenLength)
+	}
+
+	end := offsetAuthToken + tokenLen
+	if end > len(msg) {
+		return nil, fmt.Errorf("%d exceeds the message: %w", tokenLen, ErrInvalidTokenLength)
+	}
+	return msg[offsetAuthToken:end], nil
+}

shared/relay/messages/message_test.go

@@ -1,7 +1,11 @@
 package messages
 
 import (
+	"encoding/binary"
 	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 )
 
 func TestDetermineClientMessageTypeRejectsHello(t *testing.T) {
@@ -119,3 +123,245 @@ func TestMarshalHealthcheck(t *testing.T) {
 		t.Errorf("expected %d, got %d", MsgTypeHealthCheck, msgType)
 	}
 }
+
+func TestReadVersion(t *testing.T) {
+	_, err := ReadVersion([]byte{byte(ProtocolV1)})
+	assert.ErrorIs(t, err, ErrInvalidMessageLength, "a header needs two bytes")
+
+	_, err = ReadVersion([]byte{0, byte(MsgTypeAuth)})
+	assert.ErrorIs(t, err, ErrUnsupportedVersion, "version zero is invalid")
+
+	for _, version := range []ProtocolVersion{1, 2, 128, 255} {
+		got, err := ReadVersion([]byte{byte(version), byte(MsgTypeAuth)})
+		require.NoError(t, err)
+		assert.Equal(t, version, got, "the advertised version must be returned unchanged")
+	}
+}
+
+func TestSelectVersion(t *testing.T) {
+	assert.Equal(t, ProtocolV1, SelectVersion(1, 1), "equal versions select themselves")
+	assert.Equal(t, ProtocolV1, SelectVersion(2, 1), "a higher advertised version selects the server maximum")
+	assert.Equal(t, ProtocolV1, SelectVersion(255, 1), "the highest advertised version selects the server maximum")
+	assert.Equal(t, ProtocolVersion(2), SelectVersion(2, 3), "a lower advertised version selects itself")
+	assert.Equal(t, ProtocolVersion(2), SelectVersion(3, 2), "a higher advertised version selects the server maximum")
+}
+
+func TestValidateVersionAcceptsOnlyCurrent(t *testing.T) {
+	_, err := ValidateVersion([]byte{byte(CurrentProtocolVersion), byte(MsgTypeHealthCheck)})
+	require.NoError(t, err)
+
+	_, err = ValidateVersion([]byte{byte(CurrentProtocolVersion) + 1, byte(MsgTypeHealthCheck)})
+	assert.ErrorIs(t, err, ErrUnsupportedVersion, "the client build accepts only its own version")
+
+	_, err = ValidateVersion([]byte{0, byte(MsgTypeHealthCheck)})
+	assert.ErrorIs(t, err, ErrUnsupportedVersion, "version zero is invalid")
+}
+
+func TestValidateConnectionVersion(t *testing.T) {
+	require.NoError(t, ValidateConnectionVersion([]byte{1, byte(MsgTypeHealthCheck)}, ProtocolV1))
+	assert.ErrorIs(t, ValidateConnectionVersion([]byte{2, byte(MsgTypeHealthCheck)}, ProtocolV1), ErrUnsupportedVersion, "a frame must carry the selected version")
+	assert.ErrorIs(t, ValidateConnectionVersion([]byte{1, byte(MsgTypeHealthCheck)}, 2), ErrUnsupportedVersion, "a frame below the selected version is rejected too")
+	assert.ErrorIs(t, ValidateConnectionVersion([]byte{0, byte(MsgTypeHealthCheck)}, ProtocolV1), ErrUnsupportedVersion, "version zero is invalid")
+	assert.ErrorIs(t, ValidateConnectionVersion([]byte{1}, ProtocolV1), ErrInvalidMessageLength, "a header needs two bytes")
+}
+
+func TestDetermineClientMessageTypeClassification(t *testing.T) {
+	rejected := []struct {
+		name    string
+		msgType MsgType
+		wantErr error
+	}{
+		{"hello is reserved", MsgTypeHello, ErrReservedMsgType},
+		{"hello response is reserved", MsgTypeHelloResponse, ErrReservedMsgType},
+		{"auth response is server to client only", MsgTypeAuthResponse, ErrMsgTypeDirection},
+		{"peers online is server to client only", MsgTypePeersOnline, ErrMsgTypeDirection},
+		{"peers went offline is server to client only", MsgTypePeersWentOffline, ErrMsgTypeDirection},
+		{"zero is unknown", MsgTypeUnknown, ErrUnknownMsgType},
+		{"a future type is unknown", MsgType(200), ErrUnknownMsgType},
+	}
+	for _, tc := range rejected {
+		t.Run(tc.name, func(t *testing.T) {
+			got, err := DetermineClientMessageType([]byte{byte(ProtocolV1), byte(tc.msgType)})
+			assert.ErrorIs(t, err, tc.wantErr)
+			assert.Equal(t, MsgTypeUnknown, got, "a rejected type must not be returned")
+		})
+	}
+
+	accepted := []MsgType{MsgTypeAuth, MsgTypeTransport, MsgTypeClose, MsgTypeHealthCheck, MsgTypeSubscribePeerState, MsgTypeUnsubscribePeerState}
+	for _, msgType := range accepted {
+		got, err := DetermineClientMessageType([]byte{byte(ProtocolV1), byte(msgType)})
+		require.NoError(t, err)
+		assert.Equal(t, msgType, got, "a client message type must be returned")
+	}
+}
+
+func TestAuthEnvelopeRoundTrip(t *testing.T) {
+	peerID := HashID("future-client")
+	token := []byte("0123456789abcdef")
+	cases := []struct {
+		name    string
+		version ProtocolVersion
+		ext     []byte
+	}{
+		{"v2 without extension", 2, nil},
+		{"v2 with json extension", 2, []byte(`{"cred2":"value"}`)},
+		{"v255 with opaque extension", 255, []byte("opaque-extension-bytes")},
+	}
+	for _, tc := range cases {
+		t.Run(tc.name, func(t *testing.T) {
+			msg, err := MarshalAuthEnvelope(tc.version, peerID, token, tc.ext)
+			require.NoError(t, err)
+
+			version, err := ReadVersion(msg)
+			require.NoError(t, err)
+			assert.Equal(t, tc.version, version, "the header must carry the advertised version")
+
+			msgType, err := DetermineClientMessageType(msg)
+			require.NoError(t, err)
+			assert.Equal(t, MsgType(MsgTypeAuth), msgType, "the envelope is an auth message")
+
+			gotID, gotToken, err := UnmarshalAuthMsg(msg)
+			require.NoError(t, err)
+			assert.Equal(t, peerID, *gotID, "the peer ID must be preserved")
+			assert.Equal(t, token, gotToken, "the token must be returned without the extension bytes")
+		})
+	}
+}
+
+func TestMarshalAuthEnvelopeRejectsInvalidInput(t *testing.T) {
+	peerID := HashID("future-client")
+
+	_, err := MarshalAuthEnvelope(ProtocolV1, peerID, []byte("token"), nil)
+	assert.Error(t, err, "v1 does not use the envelope layout")
+
+	_, err = MarshalAuthEnvelope(2, peerID, nil, nil)
+	assert.ErrorIs(t, err, ErrInvalidTokenLength, "the token is mandatory")
+
+	_, err = MarshalAuthEnvelope(2, peerID, []byte("token"), make([]byte, MaxHandshakeSize))
+	assert.Error(t, err, "the envelope must fit the handshake size limit")
+}
+
+func TestUnmarshalAuthMsgV1LayoutUnchanged(t *testing.T) {
+	peerID := HashID("legacy-client")
+	// In the v1 layout the token runs to the end of the frame, so bytes that would be a
+	// length prefix in the envelope belong to the token here.
+	payload := []byte{0x00, 0x05, 'a', 'b', 'c', 'd', 'e', 'f', 'g'}
+	msg, err := MarshalAuthMsg(peerID, payload)
+	require.NoError(t, err)
+
+	gotID, gotToken, err := UnmarshalAuthMsg(msg)
+	require.NoError(t, err)
+	assert.Equal(t, peerID, *gotID, "the peer ID must be preserved")
+	assert.Equal(t, payload, gotToken, "the v1 token must be returned whole")
+}
+
+func TestUnmarshalAuthMsgRejectsMalformedFrames(t *testing.T) {
+	peerID := HashID("future-client")
+	valid, err := MarshalAuthEnvelope(2, peerID, []byte("token"), []byte("ext"))
+	require.NoError(t, err)
+
+	t.Run("zero token length", func(t *testing.T) {
+		msg := append([]byte(nil), valid...)
+		binary.BigEndian.PutUint16(msg[offsetAuthTokenLen:], 0)
+		_, _, err := UnmarshalAuthMsg(msg)
+		assert.ErrorIs(t, err, ErrInvalidTokenLength)
+	})
+
+	t.Run("token length past the end", func(t *testing.T) {
+		msg := append([]byte(nil), valid...)
+		binary.BigEndian.PutUint16(msg[offsetAuthTokenLen:], uint16(len(msg)))
+		_, _, err := UnmarshalAuthMsg(msg)
+		assert.ErrorIs(t, err, ErrInvalidTokenLength)
+	})
+
+	t.Run("truncated before the length field", func(t *testing.T) {
+		_, _, err := UnmarshalAuthMsg(valid[:headerTotalSizeAuth+1])
+		assert.ErrorIs(t, err, ErrInvalidMessageLength)
+	})
+
+	t.Run("version zero", func(t *testing.T) {
+		msg := append([]byte(nil), valid...)
+		msg[0] = 0
+		_, _, err := UnmarshalAuthMsg(msg)
+		assert.ErrorIs(t, err, ErrUnsupportedVersion)
+	})
+
+	t.Run("invalid magic", func(t *testing.T) {
+		msg := append([]byte(nil), valid...)
+		msg[offsetMagicByte] ^= 0xFF
+		_, _, err := UnmarshalAuthMsg(msg)
+		assert.Error(t, err)
+	})
+
+	t.Run("shorter than the fixed header", func(t *testing.T) {
+		_, _, err := UnmarshalAuthMsg(valid[:headerTotalSizeAuth-1])
+		assert.ErrorIs(t, err, ErrInvalidMessageLength)
+	})
+}
+
+func TestEncoderStampsVersion(t *testing.T) {
+	for _, version := range []ProtocolVersion{1, 2, 7} {
+		enc := NewEncoder(version)
+		assert.Equal(t, version, enc.Version())
+		assert.Equal(t, []byte{byte(version), byte(MsgTypeHealthCheck)}, enc.Healthcheck(), "healthcheck must carry the version")
+		assert.Equal(t, []byte{byte(version), byte(MsgTypeClose)}, enc.CloseMsg(), "close must carry the version")
+
+		resp, err := enc.AuthResponse("rel://relay.example")
+		require.NoError(t, err)
+		assert.Equal(t, byte(version), resp[0], "auth response must carry the version")
+		assert.Equal(t, byte(MsgTypeAuthResponse), resp[1])
+		addr, err := UnmarshalAuthResponse(resp)
+		require.NoError(t, err)
+		assert.Equal(t, "rel://relay.example", addr, "the auth response payload must be the address")
+
+		online, err := enc.PeersOnline([]PeerID{HashID("a")})
+		require.NoError(t, err)
+		require.Len(t, online, 1)
+		assert.Equal(t, byte(version), online[0][0], "peers online must carry the version")
+		assert.Equal(t, byte(MsgTypePeersOnline), online[0][1])
+
+		offline, err := enc.PeersWentOffline([]PeerID{HashID("a")})
+		require.NoError(t, err)
+		require.Len(t, offline, 1)
+		assert.Equal(t, byte(version), offline[0][0], "peers went offline must carry the version")
+		assert.Equal(t, byte(MsgTypePeersWentOffline), offline[0][1])
+	}
+}
+
+func TestPackageEncodersUseCurrentVersion(t *testing.T) {
+	current := NewEncoder(CurrentProtocolVersion)
+	assert.Equal(t, current.CloseMsg(), MarshalCloseMsg())
+	assert.Equal(t, current.Healthcheck(), MarshalHealthcheck())
+
+	want, err := current.AuthResponse("addr")
+	require.NoError(t, err)
+	got, err := MarshalAuthResponse("addr")
+	require.NoError(t, err)
+	assert.Equal(t, want, got)
+
+	ids := []PeerID{HashID("a"), HashID("b")}
+	wantOnline, err := current.PeersOnline(ids)
+	require.NoError(t, err)
+	gotOnline, err := MarshalPeersOnline(ids)
+	require.NoError(t, err)
+	assert.Equal(t, wantOnline, gotOnline)
+}
+
+func TestUpdateTransportMsgRewritesVersionAndSender(t *testing.T) {
+	dst := HashID("dst")
+	sender := HashID("sender")
+	payload := []byte("payload")
+	msg, err := MarshalTransportMsg(dst, payload)
+	require.NoError(t, err)
+
+	require.NoError(t, UpdateTransportMsg(msg, ProtocolVersion(2), sender))
+	assert.Equal(t, byte(2), msg[0], "the header must carry the destination version")
+	assert.Equal(t, byte(MsgTypeTransport), msg[1], "the message type must be preserved")
+
+	gotID, gotPayload, err := UnmarshalTransportMsg(msg)
+	require.NoError(t, err)
+	assert.Equal(t, sender, *gotID, "the peer ID must be rewritten to the sender")
+	assert.Equal(t, payload, gotPayload, "the payload must be preserved")
+
+	assert.ErrorIs(t, UpdateTransportMsg(msg[:headerTotalSizeTransport-1], ProtocolV1, sender), ErrInvalidMessageLength)
+}

shared/relay/messages/peer_state.go

@@ -5,7 +5,7 @@ import (
 )
 
 func MarshalSubPeerStateMsg(ids []PeerID) ([][]byte, error) {
-	return marshalPeerIDs(ids, byte(MsgTypeSubscribePeerState))
+	return marshalPeerIDs(CurrentProtocolVersion, ids, byte(MsgTypeSubscribePeerState))
 }
 
 func UnmarshalSubPeerStateMsg(buf []byte) ([]PeerID, error) {
@@ -13,7 +13,7 @@ func UnmarshalSubPeerStateMsg(buf []byte) ([]PeerID, error) {
 }
 
 func MarshalUnsubPeerStateMsg(ids []PeerID) ([][]byte, error) {
-	return marshalPeerIDs(ids, byte(MsgTypeUnsubscribePeerState))
+	return marshalPeerIDs(CurrentProtocolVersion, ids, byte(MsgTypeUnsubscribePeerState))
 }
 
 func UnmarshalUnsubPeerStateMsg(buf []byte) ([]PeerID, error) {
@@ -21,7 +21,7 @@ func UnmarshalUnsubPeerStateMsg(buf []byte) ([]PeerID, error) {
 }
 
 func MarshalPeersOnline(ids []PeerID) ([][]byte, error) {
-	return marshalPeerIDs(ids, byte(MsgTypePeersOnline))
+	return marshalPeerIDs(CurrentProtocolVersion, ids, byte(MsgTypePeersOnline))
 }
 
 func UnmarshalPeersOnlineMsg(buf []byte) ([]PeerID, error) {
@@ -29,7 +29,7 @@ func UnmarshalPeersOnlineMsg(buf []byte) ([]PeerID, error) {
 }
 
 func MarshalPeersWentOffline(ids []PeerID) ([][]byte, error) {
-	return marshalPeerIDs(ids, byte(MsgTypePeersWentOffline))
+	return marshalPeerIDs(CurrentProtocolVersion, ids, byte(MsgTypePeersWentOffline))
 }
 
 func UnMarshalPeersWentOffline(buf []byte) ([]PeerID, error) {
@@ -37,7 +37,7 @@ func UnMarshalPeersWentOffline(buf []byte) ([]PeerID, error) {
 }
 
 // marshalPeerIDs is a generic function to marshal peer IDs with a specific message type
-func marshalPeerIDs(ids []PeerID, msgType byte) ([][]byte, error) {
+func marshalPeerIDs(version ProtocolVersion, ids []PeerID, msgType byte) ([][]byte, error) {
 	if len(ids) == 0 {
 		return nil, fmt.Errorf("no list of peer ids provided")
 	}
@@ -54,7 +54,7 @@ func marshalPeerIDs(ids []PeerID, msgType byte) ([][]byte, error) {
 
 		totalSize := sizeOfProtoHeader + len(chunk)*peerIDSize
 		buf := make([]byte, totalSize)
-		buf[0] = byte(CurrentProtocolVersion)
+		buf[0] = byte(version)
 		buf[1] = msgType
 
 		offset := sizeOfProtoHeader