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

2026-10-25 22:09:08 +01:00 · 2026-10-21 12:27:44 +02:00
parent aa0f569677
commit 1e63b68c4d
10 changed files with 1128 additions and 91 deletions
@@ -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
+}