From 386c2508dc6a90d0cdc85efbc53478827a2f4a36 Mon Sep 17 00:00:00 2001
From: Lukas Pukenis <lukas.pukenis@nordsec.com>
Date: Fri, 31 Jan 2025 14:32:24 +0200
Subject: [PATCH] Force Pyro to timeout quickly and introduce retries

We observe that sometimes pyro remote calls are with
time gaps in between which should not happen. We guess
it might be due to network issues. However the packets
still arrive and might be due to TCP retransmission.

Pyro has a timeout property which is not exposed publicly.
This property is set and looped over if timeout error is
returned to fix the issue.

Signed-off-by: Lukas Pukenis <lukas.pukenis@nordsec.com>
---
 .unreleased/LLT-5975_pyro_timeouts     |  0
 nat-lab/tests/uniffi/libtelio_proxy.py | 35 ++++++++++++++++++--------
 2 files changed, 24 insertions(+), 11 deletions(-)
 create mode 100644 .unreleased/LLT-5975_pyro_timeouts

diff --git a/.unreleased/LLT-5975_pyro_timeouts b/.unreleased/LLT-5975_pyro_timeouts
new file mode 100644
index 000000000..e69de29bb
diff --git a/nat-lab/tests/uniffi/libtelio_proxy.py b/nat-lab/tests/uniffi/libtelio_proxy.py
index 6c76be788..10973ad09 100644
--- a/nat-lab/tests/uniffi/libtelio_proxy.py
+++ b/nat-lab/tests/uniffi/libtelio_proxy.py
@@ -36,14 +36,23 @@ def __init__(self, name: str, object_uri: str, features: libtelio.Features):
         self._features = features
 
     def _handle_remote_error(self, f):
+        print(
+            f"[{self._name}]: [{datetime.time()}]: handle_remote_error outer: uri: {self._uri}"
+        )
         with Proxy(self._uri) as remote:
-            fn_res = f(remote)
-            if fn_res is None:
-                return None
-            (res, err) = fn_res
-            if err is not None:
-                raise Exception(err)
-            return res
+            print(f"[{self._name}]: [{datetime.time()}]: handle_remote_error inner")
+            try:
+                fn_res = f(remote)
+                if fn_res is None:
+                    return None
+                (res, err) = fn_res
+                if err is not None:
+                    print(
+                        f"[{self._name}]: [{datetime.now()}]: Pyro error: {err}")
+                    raise Exception(err)
+                return res
+            except Pyro5.errors.TimeoutError:
+                print(f"[{self._name}]: [{datetime.time()}]: Pyro5 timeout")
 
     @move_to_async_thread
     def shutdown(self, container_or_vm_name: Optional[str] = None):
@@ -90,9 +99,11 @@ def create(self):
     def next_event(self) -> libtelio.Event:
         try:
             ev = self._handle_remote_error(lambda r: r.next_event())
-            print(f"[{self._name}]:{datetime.now()} proxy::next_event returned: {ev}")
+            print(
+                f"[{self._name}]:{datetime.now()} proxy::next_event returned: {ev}")
         except Exception as e:
-            print(f"[{self._name}]:{datetime.now()} proxy::next_event: exception {e}")
+            print(
+                f"[{self._name}]:{datetime.now()} proxy::next_event: exception {e}")
             raise
 
         return ev
@@ -134,7 +145,8 @@ def connect_to_exit_node(self, public_key, allowed_ips, endpoint):
     @move_to_async_thread
     def connect_to_exit_node_pq(self, public_key, allowed_ips, endpoint):
         self._handle_remote_error(
-            lambda r: r.connect_to_exit_node_pq(public_key, allowed_ips, endpoint)
+            lambda r: r.connect_to_exit_node_pq(
+                public_key, allowed_ips, endpoint)
         )
 
     @move_to_async_thread
@@ -143,7 +155,8 @@ def disconnect_from_exit_nodes(self):
 
     @move_to_async_thread
     def enable_magic_dns(self, forward_servers):
-        self._handle_remote_error(lambda r: r.enable_magic_dns(forward_servers))
+        self._handle_remote_error(
+            lambda r: r.enable_magic_dns(forward_servers))
 
     @move_to_async_thread
     def disable_magic_dns(self):