From 302cdaac70d0f279585a2b7cf35b2a2b4e7f9447 Mon Sep 17 00:00:00 2001
From: Ilija Lazoroski <ilija.la@live.com>
Date: Thu, 1 Sep 2022 09:57:46 +0200
Subject: [PATCH] Agent: Use event.wait() when running TCPRelay

---
 monkey/infection_monkey/tcp_relay.py | 3 +--
 vulture_allowlist.py                 | 3 +++
 2 files changed, 4 insertions(+), 2 deletions(-)

diff --git a/monkey/infection_monkey/tcp_relay.py b/monkey/infection_monkey/tcp_relay.py
index dd317bdd3..bb396de86 100644
--- a/monkey/infection_monkey/tcp_relay.py
+++ b/monkey/infection_monkey/tcp_relay.py
@@ -46,8 +46,7 @@ class TCPRelay(Thread):
         )
         proxy.start()
 
-        while not self._stopped.is_set():
-            sleep(0.001)
+        self._stopped.wait()
 
         self._wait_for_users_to_disconnect()
 
diff --git a/vulture_allowlist.py b/vulture_allowlist.py
index 5ab45088a..da14d4a42 100644
--- a/vulture_allowlist.py
+++ b/vulture_allowlist.py
@@ -286,6 +286,9 @@ event
 deserialize
 serialized_event
 
+# TODO: Remove after #2231 is closed
+relay_users
+
 # pydantic base models
 underscore_attrs_are_private
 extra