sponsor process attach hardening (#208)

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

3 files changed, 11 insertions, 6 deletions

diff --git a/src/zen/cmds/up_cmd.cpp b/src/zen/cmds/up_cmd.cpp
index c65dc5d5a..ac2f42a86 100644
--- a/src/zen/cmds/up_cmd.cpp
+++ b/src/zen/cmds/up_cmd.cpp
@@ -176,6 +176,7 @@ AttachCommand::Run(const ZenCliOptions& GlobalOptions, int argc, char** argv)
 		return 1;
 	}
 
+	// Sponsor processes are checked every second, so 2 second wait time should be enough
 	if (!Entry->AddSponsorProcess(m_OwnerPid, 2000))
 	{
 		ZEN_WARN("unable to add sponsor process to running zen server instance");
diff --git a/src/zenserver/main.cpp b/src/zenserver/main.cpp
index 25a5d5cf5..2fb01ebf1 100644
--- a/src/zenserver/main.cpp
+++ b/src/zenserver/main.cpp
@@ -138,6 +138,7 @@ ZenEntryPoint::Run()
 					Entry->Pid.load(),
 					m_ServerOptions.OwnerPid);
 
+				// Sponsor processes are checked every second, so 2 second wait time should be enough
 				if (Entry->AddSponsorProcess(m_ServerOptions.OwnerPid, 2000))
 				{
 					std::exit(0);
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;
 }