From bd9e59dd90cab837ef807f0cfbe148fee7b414f2 Mon Sep 17 00:00:00 2001
From: jnovikov <johnnovikov0@gmail.com>
Date: Fri, 9 Feb 2024 21:26:53 +0000
Subject: [PATCH] Make checker more reliable

---
 checkers/cell/cell_lib.py |  7 +++++--
 checkers/cell/checker.py  | 29 ++++++++++++++++++++++-------
 services/cell/conf/rr.yml |  2 +-
 3 files changed, 28 insertions(+), 10 deletions(-)

diff --git a/checkers/cell/cell_lib.py b/checkers/cell/cell_lib.py
index 6bec37f..33914fa 100644
--- a/checkers/cell/cell_lib.py
+++ b/checkers/cell/cell_lib.py
@@ -93,7 +93,7 @@ def subscribe(self, ws: websockets.sync.connection.Connection, op_id: int, sheet
         return msg
         
     
-    def modify_cell(self, ws: websockets.sync.connection.Connection, op_id: int, sheet_id: str, cell: str, value: str, token: str):
+    def modify_cell(self, ws: websockets.sync.connection.Connection, op_id: int, sheet_id: str, cell: str, value: str, token: str, timeout: int = None):
         payload = {
             "id": op_id,
             "rpc": {
@@ -108,8 +108,11 @@ def modify_cell(self, ws: websockets.sync.connection.Connection, op_id: int, she
         }
 
         ws.send(json.dumps(payload))
+        msg = self.c.decode_json(ws.recv(timeout=timeout), 'Invalid JSON received from RPC')
+        if not msg:
+            raise TimeoutError("Got empty response")
 
-        msg = self.c.decode_json(ws.recv(), 'Invalid JSON received from RPC')
+        self.c.assert_eq(type(msg), dict, 'Invalid JSON received from RPC')
         self.c.assert_eq(msg.get('id'), op_id, 'Invalid id in RPC response')
         return msg
 
diff --git a/checkers/cell/checker.py b/checkers/cell/checker.py
index 1cefcf9..cd3f863 100755
--- a/checkers/cell/checker.py
+++ b/checkers/cell/checker.py
@@ -21,7 +21,7 @@
 
 class Checker(BaseChecker):
     vulns: int = 1
-    timeout: int = 15
+    timeout: int = 20
     uses_attack_data: bool = True
 
     def __init__(self, *args, **kwargs):
@@ -102,15 +102,30 @@ def check(self):
         
             new_cell_value = rnd_string(30)
             # Modify cell.
-            _modify_result = self.cm.modify_cell(ws_writer, 2, sid, 'B2', new_cell_value, modifyToken)
+            # Try to modify cell twice to solve the reliability issue.
+            for _attempt in range(0, 2):
+                try:
+                    _modify_result = self.cm.modify_cell(ws_writer, 2, sid, 'B2', new_cell_value, modifyToken, timeout=self.timeout // 3)
+                    break
+                except TimeoutError:
+                    continue
+
+            self.assert_neq(_modify_result, None, 'Failed to modify cell: no response received')
             self.assert_in('rpc', _modify_result.keys(), 'Invalid modify_cell response')
-            self.assert_eq(_modify_result.get('rpc', {}).get('data', {}).get('sheet', {}).get('title', None), sheet_name, 'Invalid sheet title in modify_cell response')
-            cell_values = [x.get('val', '') for x in _modify_result.get('rpc', {}).get('data').get('sheet', {}).get('cells', [])]
+            mb_sheet = _modify_result.get('rpc', {}).get('data', {}).get('sheet', {})
+            self.assert_eq(mb_sheet.get('title', None), sheet_name, 'Invalid sheet title in modify_cell response')
+            cell_values = [x.get('val', '') for x in mb_sheet.get('cells', [])]
             self.assert_in(new_cell_value, cell_values, 'New cell value not found in modify_cell response')
 
-            # Fet update from reader.
-            expected_update = ws_reader.recv()
-            update_data = self.decode_json(expected_update, 'Invalid JSON received from sheet update')
+            # Get update from reader.
+            # Try to get update twice to solve the reliability issue.
+            for _attempt in range(0, 2):
+                expected_update = ws_reader.recv(timeout=self.timeout // 3)
+                update_data = self.decode_json(expected_update, 'Invalid JSON received from sheet update')
+                if update_data:
+                    break
+
+            self.assert_eq(type(update_data), dict, 'Invalid JSON received from sheet update')
             update_data = update_data.get('push', {}).get('pub', {}).get('data', {})
             self.assert_eq(update_data.get('title', ''), sheet_name, 'Invalid sheet title returned from sheet update')
             cell_values = [x.get('val', '') for x in update_data.get('cells', [])]
diff --git a/services/cell/conf/rr.yml b/services/cell/conf/rr.yml
index eaf3e8a..ed2b3b0 100644
--- a/services/cell/conf/rr.yml
+++ b/services/cell/conf/rr.yml
@@ -20,7 +20,7 @@ http:
   pool:
     num_workers: 2
     supervisor:
-      max_worker_memory: 100
+      max_worker_memory: 128
 server:
   command: 'php app.php'
   relay: pipes