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

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