Fix gRPC connection issue with retry logic

Handeled `StatusCode.UNAVAILABLE` error with a retry loop and 1s delay
Added requests and angr to `pyproject.toml` to fix dependency issues

Signed-off-by: Amalia Ionescu <[email protected]>

Author: AmaliaI12

automatic_exploit_generation/exploiters/zeratool/zeratool.py

@@ -1,5 +1,6 @@
 import pickle
 import typing
+import time
 
 import docker
 import grpc
@@ -73,6 +74,7 @@ def _run_exploitation_in_container(
             ports={"13000/tcp": 13000},
             publish_all_ports=True,
         )
+        container.reload()
         container_ip = container.attrs["NetworkSettings"]["IPAddress"]
         exploit = self._request_exploitation_to_grpc_service(
             container_ip, overflow_only, format_only, win_funcs
@@ -103,8 +105,25 @@ def _request_exploitation_to_grpc_service(
             format_only=format_only,
             serialized_win_funcs=serialized_win_funcs_arg,
         )
-        response = stub.Exploit(query)
-        exploit = pickle.loads(response.pickledExploit)
+
+        num_retries = 0
+        retry_sleep_time = 1
+        exploit = None
+        while num_retries < 30:
+            try:
+                response = stub.Exploit(query)
+                exploit = pickle.loads(response.pickledExploit)
+                break
+            except grpc.RpcError as rpc_error:
+                if rpc_error.code() == grpc.StatusCode.CANCELLED:
+                    pass
+                elif rpc_error.code() == grpc.StatusCode.UNAVAILABLE:
+                    pass
+                else:
+                    print(f"Received unknown RPC error: code={rpc_error.code()} message={rpc_error.details()}")
+                    break
+                time.sleep(retry_sleep_time)
+                num_retries += 1
 
         return exploit
 
@@ -127,9 +146,9 @@ def _exploit(self) -> Exploit:
             exploit = self._run_exploitation_in_container(
                 overflow_only, format_only, win_funcs_used
             )
+
             if exploit:
                 return exploit
-
         return None
 
     def _get_win_funcs_names(

pyproject.toml

@@ -15,6 +15,8 @@ grpcio = "^1.54.2"
 protobuf = "^4.23.2"
 click = "^8.1.3"
 hexdump = "^3.3"
+requests = "2.31.0"
+angr = "^9.2.52"
 commons = {path = "../commons"}
 zeratool_lib = {path = "../zeratool_lib"}