From e9bc10ef59b5023f7f3337a98740c06c9884ff2b Mon Sep 17 00:00:00 2001
From: boatbomber <zack@boatbomber.com>
Date: Mon, 9 Oct 2023 19:13:03 -0400
Subject: [PATCH] Give half second grace for handshake

---
 package.json   |  2 +-
 src/server.lua | 31 ++++++++++++++++++++++++-------
 wally.lock     |  2 +-
 wally.toml     |  2 +-
 4 files changed, 27 insertions(+), 10 deletions(-)

diff --git a/package.json b/package.json
index 35c2dfe..9fea5ad 100644
--- a/package.json
+++ b/package.json
@@ -1,6 +1,6 @@
 {
   "name": "@boatbomber/encryptednet",
-  "version": "1.0.3",
+  "version": "1.0.4",
   "license": "MIT",
   "repository": {
     "type": "git",
diff --git a/src/server.lua b/src/server.lua
index c2a0708..7139053 100644
--- a/src/server.lua
+++ b/src/server.lua
@@ -26,6 +26,23 @@ Players.PlayerRemoving:Connect(function(Player)
 	PlayerData[Player] = nil
 end)
 
+local function getPlayerData(Player, shouldYield)
+	local playerData = PlayerData[Player]
+	if shouldYield and playerData == nil then
+		-- Perhaps they're in middle of the handshake?
+		-- Poll for half a second
+		for _=1, 15 do
+			task.wait(1/30)
+			playerData = PlayerData[Player]
+			if playerData then
+				break
+			end
+		end
+	end
+
+	return playerData
+end
+
 return function(Remote)
 	local Wrapper = setmetatable({}, { __index = Remote })
 
@@ -33,7 +50,7 @@ return function(Remote)
 
 	function Wrapper:Connect(callback)
 		Remote:Connect(function(Player, encryptedData, signature)
-			local playerData = PlayerData[Player]
+			local playerData = getPlayerData(Player, false) -- No yield since they can't have sent valid data without a handshake first
 			if not playerData then
 				return
 			end
@@ -59,7 +76,7 @@ return function(Remote)
 	end
 
 	function Wrapper:SendToPlayer(Player, ...)
-		local playerData = PlayerData[Player]
+		local playerData = getPlayerData(Player, true) -- Yield in case we're sending before they're ready
 		if not playerData then
 			return
 		end
@@ -81,13 +98,13 @@ return function(Remote)
 				continue
 			end
 
-			Wrapper:SendToPlayer(Player, ...)
+			task.spawn(Wrapper.SendToPlayer, Wrapper, Player, ...)
 		end
 	end
 
 	function Wrapper:SendToAllPlayers(...)
 		for _, Player in ipairs(Players:GetPlayers()) do
-			Wrapper:SendToPlayer(Player, ...)
+			task.spawn(Wrapper.SendToPlayer, Wrapper, Player, ...)
 		end
 	end
 
@@ -100,14 +117,14 @@ return function(Remote)
 				continue
 			end
 
-			Wrapper:SendToPlayer(Player, ...)
+			task.spawn(Wrapper.SendToPlayer, Wrapper, Player, ...)
 		end
 	end
 
 	-- AsyncFunction
 
 	function Wrapper:CallPlayerAsync(Player, ...)
-		local playerData = PlayerData[Player]
+		local playerData = getPlayerData(Player, true) -- Yield in case we're calling before they're ready
 		if not playerData then
 			return
 		end
@@ -125,7 +142,7 @@ return function(Remote)
 
 	function Wrapper:SetCallback(callback)
 		Remote:SetCallback(function(Player, encryptedData, signature)
-			local playerData = PlayerData[Player]
+			local playerData = getPlayerData(Player, false) -- No yield since they can't have sent valid data without a handshake first
 			if not playerData then
 				return
 			end
diff --git a/wally.lock b/wally.lock
index fdcbddd..02d3dd8 100644
--- a/wally.lock
+++ b/wally.lock
@@ -9,5 +9,5 @@ dependencies = []
 
 [[package]]
 name = "boatbomber/encryptednet"
-version = "1.0.3"
+version = "1.0.4"
 dependencies = []
diff --git a/wally.toml b/wally.toml
index e1ecc7a..80d1ce2 100644
--- a/wally.toml
+++ b/wally.toml
@@ -1,7 +1,7 @@
 [package]
 name = "boatbomber/encryptednet"
 description = "Authenticated encryption of Roblox networking with ECDH key exchanges and ChaCha20 ciphering."
-version = "1.0.3"
+version = "1.0.4"
 license = "MIT"
 authors = ["boatbomber (https://boatbomber.com)"]
 registry = "https://github.com/upliftgames/wally-index"