[relay] Improve relay messages (#2574)

Co-authored-by: Zoltán Papp <zoltan.pmail@gmail.com>

relay/messages/message.go

@@ -7,12 +7,21 @@ import (
 )
 
 const (
-	MsgTypeUnknown       MsgType = 0
-	MsgTypeHello         MsgType = 1
+	MaxHandshakeSize     = 212
+	MaxHandshakeRespSize = 8192
+
+	CurrentProtocolVersion = 1
+
+	MsgTypeUnknown MsgType = 0
+	// Deprecated: Use MsgTypeAuth instead.
+	MsgTypeHello MsgType = 1
+	// Deprecated: Use MsgTypeAuthResponse instead.
 	MsgTypeHelloResponse MsgType = 2
 	MsgTypeTransport     MsgType = 3
 	MsgTypeClose         MsgType = 4
 	MsgTypeHealthCheck   MsgType = 5
+	MsgTypeAuth                  = 6
+	MsgTypeAuthResponse          = 7
 
 	SizeOfVersionByte = 1
 	SizeOfMsgType     = 1
@@ -22,12 +31,12 @@ const (
 	sizeOfMagicByte = 4
 
 	headerSizeTransport = IDSize
+
 	headerSizeHello     = sizeOfMagicByte + IDSize
 	headerSizeHelloResp = 0
 
-	MaxHandshakeSize = 8192
-
-	CurrentProtocolVersion = 1
+	headerSizeAuth     = sizeOfMagicByte + IDSize
+	headerSizeAuthResp = 0
 )
 
 var (
@@ -47,6 +56,10 @@ func (m MsgType) String() string {
 		return "hello"
 	case MsgTypeHelloResponse:
 		return "hello response"
+	case MsgTypeAuth:
+		return "auth"
+	case MsgTypeAuthResponse:
+		return "auth response"
 	case MsgTypeTransport:
 		return "transport"
 	case MsgTypeClose:
@@ -58,10 +71,6 @@ func (m MsgType) String() string {
 	}
 }
 
-type HelloResponse struct {
-	InstanceAddress string
-}
-
 // ValidateVersion checks if the given version is supported by the protocol
 func ValidateVersion(msg []byte) (int, error) {
 	if len(msg) < SizeOfVersionByte {
@@ -84,6 +93,7 @@ func DetermineClientMessageType(msg []byte) (MsgType, error) {
 	switch msgType {
 	case
 		MsgTypeHello,
+		MsgTypeAuth,
 		MsgTypeTransport,
 		MsgTypeClose,
 		MsgTypeHealthCheck:
@@ -103,6 +113,7 @@ func DetermineServerMessageType(msg []byte) (MsgType, error) {
 	switch msgType {
 	case
 		MsgTypeHelloResponse,
+		MsgTypeAuthResponse,
 		MsgTypeTransport,
 		MsgTypeClose,
 		MsgTypeHealthCheck:
@@ -112,6 +123,7 @@ func DetermineServerMessageType(msg []byte) (MsgType, error) {
 	}
 }
 
+// Deprecated: Use MarshalAuthMsg instead.
 // MarshalHelloMsg initial hello message
 // The Hello message is the first message sent by a client after establishing a connection with the Relay server. This
 // message is used to authenticate the client with the server. The authentication is done using an HMAC method.
@@ -135,6 +147,7 @@ func MarshalHelloMsg(peerID []byte, additions []byte) ([]byte, error) {
 	return msg, nil
 }
 
+// Deprecated: Use UnmarshalAuthMsg instead.
 // UnmarshalHelloMsg extracts peerID and the additional data from the hello message. The Additional data is used to
 // authenticate the client with the server.
 func UnmarshalHelloMsg(msg []byte) ([]byte, []byte, error) {
@@ -148,6 +161,7 @@ func UnmarshalHelloMsg(msg []byte) ([]byte, []byte, error) {
 	return msg[sizeOfMagicByte:headerSizeHello], msg[headerSizeHello:], nil
 }
 
+// Deprecated: Use MarshalAuthResponse instead.
 // MarshalHelloResponse creates a response message to the hello message.
 // In case of success connection the server response with a Hello Response message. This message contains the server's
 // instance URL. This URL will be used by choose the common Relay server in case if the peers are in different Relay
@@ -163,6 +177,7 @@ func MarshalHelloResponse(additionalData []byte) ([]byte, error) {
 	return msg, nil
 }
 
+// Deprecated: Use UnmarshalAuthResponse instead.
 // UnmarshalHelloResponse extracts the additional data from the hello response message.
 func UnmarshalHelloResponse(msg []byte) ([]byte, error) {
 	if len(msg) < headerSizeHelloResp {
@@ -171,6 +186,69 @@ func UnmarshalHelloResponse(msg []byte) ([]byte, error) {
 	return msg, nil
 }
 
+// MarshalAuthMsg initial authentication message
+// The Auth message is the first message sent by a client after establishing a connection with the Relay server. This
+// message is used to authenticate the client with the server. The authentication is done using an HMAC method.
+// The protocol does not limit to use HMAC, it can be any other method. If the authentication failed the server will
+// close the network connection without any response.
+func MarshalAuthMsg(peerID []byte, authPayload []byte) ([]byte, error) {
+	if len(peerID) != IDSize {
+		return nil, fmt.Errorf("invalid peerID length: %d", len(peerID))
+	}
+
+	msg := make([]byte, SizeOfProtoHeader+sizeOfMagicByte, SizeOfProtoHeader+headerSizeAuth+len(authPayload))
+
+	msg[0] = byte(CurrentProtocolVersion)
+	msg[1] = byte(MsgTypeAuth)
+
+	copy(msg[SizeOfProtoHeader:SizeOfProtoHeader+sizeOfMagicByte], magicHeader)
+
+	msg = append(msg, peerID...)
+	msg = append(msg, authPayload...)
+
+	return msg, nil
+}
+
+// UnmarshalAuthMsg extracts peerID and the auth payload from the message
+func UnmarshalAuthMsg(msg []byte) ([]byte, []byte, error) {
+	if len(msg) < headerSizeAuth {
+		return nil, nil, ErrInvalidMessageLength
+	}
+	if !bytes.Equal(msg[:sizeOfMagicByte], magicHeader) {
+		return nil, nil, errors.New("invalid magic header")
+	}
+
+	return msg[sizeOfMagicByte:headerSizeAuth], msg[headerSizeAuth:], nil
+}
+
+// MarshalAuthResponse creates a response message to the auth.
+// In case of success connection the server response with a AuthResponse message. This message contains the server's
+// 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+headerSizeAuthResp+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
+}
+
+// UnmarshalAuthResponse it is a confirmation message to auth success
+func UnmarshalAuthResponse(msg []byte) (string, error) {
+	if len(msg) < headerSizeAuthResp+1 {
+		return "", ErrInvalidMessageLength
+	}
+	return string(msg), nil
+}
+
 // MarshalCloseMsg creates a close message.
 // 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.