sponsor process attach hardening (#208)

* make sure to clear sponsor slot if pickup does not happen

1 files changed, 9 insertions, 6 deletions

diff --git a/src/zenutil/zenserverprocess.cpp b/src/zenutil/zenserverprocess.cpp
index 7f053747e..4003cef48 100644
--- a/src/zenutil/zenserverprocess.cpp
+++ b/src/zenutil/zenserverprocess.cpp
@@ -430,17 +430,22 @@ ZenServerState::ZenServerEntry::AddSponsorProcess(uint32_t PidToAdd, uint64_t Ti
 {
 	uint32_t ServerPid	   = Pid.load();
 	auto	 WaitForPickup = [&](uint32_t AddedSlotIndex) {
+		if (Timeout == 0)
+		{
+			return true;
+		}
 		Stopwatch Timer;
 		while (SponsorPids[AddedSlotIndex] == PidToAdd)
 		{
-			// Sponsor processes are checked every second, so 2 second wait time should be enough
-			if (Timer.GetElapsedTimeMs() > 2000)
+			if (Timer.GetElapsedTimeMs() > Timeout)
 			{
+				SponsorPids[AddedSlotIndex].compare_exchange_strong(PidToAdd, 0);
 				return false;
 			}
 			std::error_code _;
 			if (!IsProcessRunning(ServerPid, _))
 			{
+				SponsorPids[AddedSlotIndex].compare_exchange_strong(PidToAdd, 0);
 				return false;
 			}
 			Sleep(100);
@@ -451,16 +456,14 @@ ZenServerState::ZenServerEntry::AddSponsorProcess(uint32_t PidToAdd, uint64_t Ti
 	{
 		if (SponsorPids[SponsorIndex].load(std::memory_order_relaxed) == PidToAdd)
 		{
-			return Timeout == 0 ? true : WaitForPickup(SponsorIndex);
+			return WaitForPickup(SponsorIndex);
 		}
 		uint32_t Expected = 0;
 		if (SponsorPids[SponsorIndex].compare_exchange_strong(Expected, PidToAdd))
 		{
-			// Success!
-			return Timeout == 0 ? true : WaitForPickup(SponsorIndex);
+			return WaitForPickup(SponsorIndex);
 		}
 	}
-
 	return false;
 }