StatelessConnectHandlerComponent.h

Go to the documentation of this file.

// Copyright Epic Games, Inc. All Rights Reserved.
 
#pragma once
 
#include "PacketHandler.h"
#include "IPAddress.h"
 
 
DECLARE_LOG_CATEGORY_EXTERN(LogHandshake, Log, All);
 
 
// Forward Declarations
class UNetConnection;
class UNetDriver;
enum class EEngineNetworkRuntimeFeatures : uint16;
 
 
// Defines
#define SECRET_BYTE_SIZE 64
#define SECRET_COUNT 2
#define COOKIE_BYTE_SIZE 20
 
 
namespace UE::Net
{
    enum class EHandshakeVersion : uint8
    {
        Original                = 0,    // The original/unversioned handshake protocol
        Randomized              = 1,    // The version of the handshake protocol with randomization, versioning and debug/diagnostic tweaks
        NetCLVersion            = 2,    // Added Network CL version, for optional extra-early client rejection
 
        SessionClientId         = 3,    // Added SessionId (index of server game session) and ClientId (clientside index of NetConnection)
        NetCLUpgradeMessage     = 4,    // Added stateless handshake level NMT_Upgrade message, to allow Network CL checks to be permanently enabled
 
        Latest                  = NetCLUpgradeMessage
    };
 
    static constexpr EHandshakeVersion HandshakeCompatibilityBreaks[] = { EHandshakeVersion::SessionClientId };
 
 
    enum class EHandshakePacketType : uint8
    {
        InitialPacket       = 0,
        Challenge           = 1,
        Response            = 2,
        Ack                 = 3,
        RestartHandshake    = 4,
        RestartResponse     = 5,
        VersionUpgrade      = 6,
 
        Last = VersionUpgrade
    };
 
    enum class EHandshakePacketModifier : uint8
    {
        None                = 0x00,
        RestartHandshake    = 0x01      // Restart handshake packet
    };
 
    ENUM_CLASS_FLAGS(EHandshakePacketModifier);
 
 
    enum class EHandshakeFailureReason
    {
        None,
        WrongVersion        // Client/Server have incompatible versions
    };
 
    struct FStatelessHandshakeFailureInfo
    {
        EHandshakeFailureReason FailureReason = EHandshakeFailureReason::None;
 
        uint32 RemoteNetworkVersion = 0;
 
        EEngineNetworkRuntimeFeatures RemoteNetworkFeatures;
 
 
    public:
        FStatelessHandshakeFailureInfo();
    };
 
    using FHandshakeFailureFunc = TUniqueFunction<void(FStatelessHandshakeFailureInfo HandshakeFailureInfo)>;
}
 
 
class StatelessConnectHandlerComponent : public HandlerComponent
{
    using EHandshakeVersion = UE::Net::EHandshakeVersion;
    using EHandshakePacketType = UE::Net::EHandshakePacketType;
    using EHandshakePacketModifier = UE::Net::EHandshakePacketModifier;
    using FHandshakeFailureFunc = UE::Net::FHandshakeFailureFunc;
 
private:
    struct FCommonSendToClientParams
    {
        const TSharedPtr<const FInternetAddr>& ClientAddress;
        
        EHandshakeVersion HandshakeVersion;
 
        uint32 ClientID;
 
 
    public:
        FCommonSendToClientParams(const TSharedPtr<const FInternetAddr>& InClientAddress, EHandshakeVersion InHandshakeVersion, uint32 InClientID);
    };
 
public:
    ENGINE_API StatelessConnectHandlerComponent();
 
    ENGINE_API virtual void CountBytes(FArchive& Ar) const override;
 
    virtual bool IsValid() const override { return true; }
 
    ENGINE_API virtual void NotifyHandshakeBegin() override;
 
    ENGINE_API void InitFromConnectionless(StatelessConnectHandlerComponent* InConnectionlessHandler);
 
    ENGINE_API void SetHandshakeFailureCallback(FHandshakeFailureFunc&& InHandshakeFailureFunc);
 
private:
    void SendInitialPacket(EHandshakeVersion HandshakeVersion);
 
    void SendConnectChallenge(FCommonSendToClientParams CommonParams, uint8 ClientSentHandshakePacketCount);
 
    void SendChallengeResponse(EHandshakeVersion HandshakeVersion, uint8 InSecretId, double InTimestamp, uint8 InCookie[COOKIE_BYTE_SIZE]);
 
    void SendChallengeAck(FCommonSendToClientParams CommonParams, uint8 ClientSentHandshakePacketCount, uint8 InCookie[COOKIE_BYTE_SIZE]);
 
    void SendRestartHandshakeRequest(FCommonSendToClientParams CommonParams);
 
    void SendVersionUpgradeMessage(FCommonSendToClientParams CommonParams);
 
 
    void BeginHandshakePacket(FBitWriter& HandshakePacket, EHandshakePacketType HandshakePacketType, EHandshakeVersion HandshakeVersion,
                                uint8 SentHandshakePacketCount_LocalOrRemote, uint32 ClientID,
                                EHandshakePacketModifier HandshakePacketModifier=EHandshakePacketModifier::None);
 
 
    void SendToServer(EHandshakeVersion HandshakeVersion, EHandshakePacketType PacketType, FBitWriter& Packet);
 
    void SendToClient(FCommonSendToClientParams CommonParams, EHandshakePacketType PacketType, FBitWriter& Packet);
 
 
    void CapHandshakePacket(FBitWriter& HandshakePacket, EHandshakeVersion HandshakeVersion);
 
public:
    inline bool HasPassedChallenge(const TSharedPtr<const FInternetAddr>& Address, bool& bOutRestartedHandshake) const
    {
        bOutRestartedHandshake = bRestartedHandshake;
 
        return (LastChallengeSuccessAddress.IsValid() && Address.IsValid()) ? *LastChallengeSuccessAddress == *Address : false;
    }
 
 
    inline void GetChallengeSequence(int32& OutServerSequence, int32& OutClientSequence) const
    {
        OutServerSequence = LastServerSequence;
        OutClientSequence = LastClientSequence;
    }
 
    inline bool DoesRestartedHandshakeMatch(StatelessConnectHandlerComponent& NetConnComponent) const
    {
        return FMemory::Memcmp(AuthorisedCookie, NetConnComponent.AuthorisedCookie, COOKIE_BYTE_SIZE) == 0;
    }
 
    inline void ResetChallengeData()
    {
        LastChallengeSuccessAddress = nullptr;
        bRestartedHandshake = false;
        LastServerSequence = 0;
        LastClientSequence = 0;
        FMemory::Memzero(AuthorisedCookie, COOKIE_BYTE_SIZE);
    }
 
 
    ENGINE_API void SetDriver(UNetDriver* InDriver);
 
 
protected:
    ENGINE_API virtual void Initialize() override;
 
    ENGINE_API virtual void Incoming(FBitReader& Packet) override;
 
    ENGINE_API virtual void Outgoing(FBitWriter& Packet, FOutPacketTraits& Traits) override;
 
    ENGINE_API virtual void IncomingConnectionless(FIncomingPacketRef PacketRef) override;
 
    virtual bool CanReadUnaligned() const override
    {
        return true;
    }
 
    ENGINE_API virtual int32 GetReservedPacketBits() const override;
 
    ENGINE_API virtual void Tick(float DeltaTime) override;
 
private:
 
    struct FParsedHandshakeData
    {
        EHandshakeVersion RemoteMinVersion = EHandshakeVersion::Latest;
 
        EHandshakeVersion RemoteCurVersion = EHandshakeVersion::Latest;
 
        uint32 RemoteNetworkVersion = 0;
 
        EEngineNetworkRuntimeFeatures RemoteNetworkFeatures;
 
        UE::Net::EHandshakePacketType HandshakePacketType = UE::Net::EHandshakePacketType::InitialPacket;
 
        uint8 RemoteSentHandshakePacketCount = 0;
 
        bool bRestartHandshake = false;
 
        uint8 SecretId = 0;
 
        double Timestamp = 0.0;
 
        uint8 Cookie[COOKIE_BYTE_SIZE] = {};
 
        uint8 OrigCookie[COOKIE_BYTE_SIZE] = {};
 
 
    public:
        FParsedHandshakeData();
    };
 
 
    bool ParseHandshakePacket(FBitReader& Packet, FParsedHandshakeData& OutResult) const;
 
    bool ParseHandshakePacketOriginal(FBitReader& Packet, FParsedHandshakeData& OutResult) const;
 
    bool CheckVersion(const FParsedHandshakeData& HandshakeData, EHandshakeVersion& OutTargetVersion) const;
 
    void GenerateCookie(const TSharedPtr<const FInternetAddr>& ClientAddress, uint8 SecretId, double Timestamp, uint8 (&OutCookie)[COOKIE_BYTE_SIZE]) const;
 
    void UpdateSecret();
 
    int32 GetAdjustedSizeBits(int32 InSizeBits, EHandshakeVersion HandshakeVersion) const;
 
#if !UE_BUILD_SHIPPING
    bool TrackValidationLogs();
#endif
 
 
private:
    UNetDriver* Driver;
 
    FBitWriter OldPacket;
 
    TArray<uint8> HandshakeSecret[SECRET_COUNT];
 
    uint8 ActiveSecret;
 
    double LastSecretUpdateTimestamp;
 
    double LastInitTimestamp = 0.0;
 
    TSharedPtr<const FInternetAddr> LastChallengeSuccessAddress;
 
    int32 LastServerSequence;
 
    int32 LastClientSequence;
 
    EHandshakeVersion MinClientHandshakeVersion;
 
#if !UE_BUILD_SHIPPING
    double LastValidationLogPeriodStart = 0.0;
 
    uint32 ValidationLogCounter = 0;
#endif
 
    double LastUpgradeMessagePeriodStart = 0.0;
 
    uint32 UpgradeMessageCounter = 0;
 
    FHandshakeFailureFunc HandshakeFailureCallback;
 
 
    double LastClientSendTimestamp;
 
 
    double LastChallengeTimestamp;
 
    double LastRestartPacketTimestamp;
 
    uint8 LastSecretId;
 
    double LastTimestamp;
 
    uint8 LastCookie[COOKIE_BYTE_SIZE];
 
 
    bool bRestartedHandshake;
 
    uint8 AuthorisedCookie[COOKIE_BYTE_SIZE];
 
    TBitArray<> MagicHeader;
 
    uint32 MagicHeaderUint = 0;
 
    uint8 SentHandshakePacketCount = 0;
 
    EHandshakeVersion LastRemoteHandshakeVersion;
 
    uint32 CachedGlobalNetTravelCount = 0;
 
    uint32 CachedClientID = 0;
 
    static constexpr uint32 SessionIDSizeBits = 2;
 
    static constexpr uint32 ClientIDSizeBits = 3;
};