From c26bee22b2c44fa3a3f1d20a407b3b28578a5a4f Mon Sep 17 00:00:00 2001
From: Tom Weiland <tomdweiland@gmail.com>
Date: Wed, 4 Oct 2023 17:28:40 -0700
Subject: [PATCH] Fix connection attempts being rejected incorrectly

Closes #1
---
 .../RiptideSteamTransport/Transport/SteamClient.cs |  6 ++++++
 .../Transport/SteamConnection.cs                   |  2 ++
 .../RiptideSteamTransport/Transport/SteamPeer.cs   | 11 +++--------
 .../RiptideSteamTransport/Transport/SteamServer.cs | 14 ++++++++++++++
 4 files changed, 25 insertions(+), 8 deletions(-)

diff --git a/Assets/RiptideSteamTransport/Transport/SteamClient.cs b/Assets/RiptideSteamTransport/Transport/SteamClient.cs
index c122b08..d6d9497 100644
--- a/Assets/RiptideSteamTransport/Transport/SteamClient.cs
+++ b/Assets/RiptideSteamTransport/Transport/SteamClient.cs
@@ -14,6 +14,7 @@ public class SteamClient : SteamPeer, IClient
     {
         public event EventHandler Connected;
         public event EventHandler ConnectionFailed;
+        public event EventHandler<DataReceivedEventArgs> DataReceived;
         public event EventHandler<DisconnectedEventArgs> Disconnected;
 
         private const string LocalHostName = "localhost";
@@ -193,6 +194,11 @@ protected virtual void OnConnectionFailed()
             ConnectionFailed?.Invoke(this, EventArgs.Empty);
         }
 
+        protected override void OnDataReceived(byte[] dataBuffer, int amount, SteamConnection fromConnection)
+        {
+            DataReceived?.Invoke(this, new DataReceivedEventArgs(dataBuffer, amount, fromConnection));
+        }
+
         protected virtual void OnDisconnected(DisconnectReason reason)
         {
             Disconnected?.Invoke(this, new DisconnectedEventArgs(steamConnection, reason));
diff --git a/Assets/RiptideSteamTransport/Transport/SteamConnection.cs b/Assets/RiptideSteamTransport/Transport/SteamConnection.cs
index 6f1ac66..10cd244 100644
--- a/Assets/RiptideSteamTransport/Transport/SteamConnection.cs
+++ b/Assets/RiptideSteamTransport/Transport/SteamConnection.cs
@@ -14,6 +14,8 @@ public class SteamConnection : Connection, IEquatable<SteamConnection>
         public readonly CSteamID SteamId;
         public readonly HSteamNetConnection SteamNetConnection;
 
+        internal bool DidReceiveConnect;
+
         private readonly SteamPeer peer;
 
         internal SteamConnection(CSteamID steamId, HSteamNetConnection steamNetConnection, SteamPeer peer)
diff --git a/Assets/RiptideSteamTransport/Transport/SteamPeer.cs b/Assets/RiptideSteamTransport/Transport/SteamPeer.cs
index f8f5a6f..49501a1 100644
--- a/Assets/RiptideSteamTransport/Transport/SteamPeer.cs
+++ b/Assets/RiptideSteamTransport/Transport/SteamPeer.cs
@@ -10,13 +10,11 @@
 
 namespace Riptide.Transports.Steam
 {
-    public class SteamPeer
+    public abstract class SteamPeer
     {
-        /// <summary>The name to use when logging messages via <see cref="RiptideLogger"/>.</summary>
+        /// <summary>The name to use when logging messages via <see cref="Utils.RiptideLogger"/>.</summary>
         public const string LogName = "STEAM";
 
-        public event EventHandler<DataReceivedEventArgs> DataReceived;
-
         protected const int MaxMessages = 256;
 
         private readonly byte[] receiveBuffer;
@@ -66,9 +64,6 @@ internal void Send(byte[] dataBuffer, int numBytes, HSteamNetConnection toConnec
             handle.Free();
         }
 
-        protected virtual void OnDataReceived(byte[] dataBuffer, int amount, SteamConnection fromConnection)
-        {
-            DataReceived?.Invoke(this, new DataReceivedEventArgs(dataBuffer, amount, fromConnection));
-        }
+        protected abstract void OnDataReceived(byte[] dataBuffer, int amount, SteamConnection fromConnection);
     }
 }
diff --git a/Assets/RiptideSteamTransport/Transport/SteamServer.cs b/Assets/RiptideSteamTransport/Transport/SteamServer.cs
index f566a42..f5989c7 100644
--- a/Assets/RiptideSteamTransport/Transport/SteamServer.cs
+++ b/Assets/RiptideSteamTransport/Transport/SteamServer.cs
@@ -13,6 +13,7 @@ namespace Riptide.Transports.Steam
     public class SteamServer : SteamPeer, IServer
     {
         public event EventHandler<ConnectedEventArgs> Connected;
+        public event EventHandler<DataReceivedEventArgs> DataReceived;
         public event EventHandler<DisconnectedEventArgs> Disconnected;
 
         public ushort Port { get; private set; }
@@ -134,6 +135,19 @@ protected internal virtual void OnConnected(Connection connection)
             Connected?.Invoke(this, new ConnectedEventArgs(connection));
         }
 
+        protected override void OnDataReceived(byte[] dataBuffer, int amount, SteamConnection fromConnection)
+        {
+            if ((MessageHeader)dataBuffer[0] == MessageHeader.Connect)
+            {
+                if (fromConnection.DidReceiveConnect)
+                    return;
+
+                fromConnection.DidReceiveConnect = true;
+            }
+
+            DataReceived?.Invoke(this, new DataReceivedEventArgs(dataBuffer, amount, fromConnection));
+        }
+
         protected virtual void OnDisconnected(CSteamID steamId, DisconnectReason reason)
         {
             if (connections.TryGetValue(steamId, out SteamConnection connection))