[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")
+}