firewall: retry qdb.read/multiread() on InterruptedError in remaining call sites

Author: Harshit6-b

qubesagent/firewall.py

@@ -120,7 +120,12 @@ def run_user_script(self):
 
     def read_rules(self, target):
         """Read rules from QubesDB and return them as a list of dicts"""
-        entries = self.qdb.multiread('/qubes-firewall/{}/'.format(target))
+        while True:
+            try:
+                entries = self.qdb.multiread('/qubes-firewall/{}/'.format(target))
+                break
+            except InterruptedError:
+                continue
         assert isinstance(entries, dict)
         # drop full path
         entries = dict(((k.split('/')[3], v.decode())
@@ -193,7 +198,12 @@ def update_handled(self, addr):
         User applications may watch these paths for count increases to remain
         up to date with QubesDB changes.
         """
-        cnt = self.qdb.read('/qubes-firewall-handled/{}'.format(addr))
+        while True:
+            try:
+                cnt = self.qdb.read('/qubes-firewall-handled/{}'.format(addr))
+                break
+            except InterruptedError:
+                continue
         try:
             cnt = int(cnt)
         except (TypeError, ValueError):
@@ -289,7 +299,13 @@ def is_ip6(addr):
     def log_error(self, msg):
         self.log.error(msg)
 
-        user = (self.qdb.read('/default-user') or b'user').decode()
+        while True:
+            try:
+                user = (self.qdb.read('/default-user') or b'user').decode()
+                break
+            except InterruptedError:
+                continue
+            
         try:
             uid = pwd.getpwnam(user).pw_uid
         except KeyError:

qubesagent/test_firewall.py

@@ -594,4 +594,69 @@ def read_with_interrupt(key):
             self.obj, 4)
 
         self.assertEqual(result, ['10.137.0.1', '10.137.0.2'])
-        self.assertEqual(call_count[0], 2)
+        self.assertEqual(call_count[0], 2)
+
+    def test_read_rules_retries_on_interrupted_error(self):
+        """read_rules() must retry qdb.multiread() on InterruptedError."""
+        self.obj.qdb.entries['/qubes-firewall/10.137.0.2/policy'] = b'drop'
+        self.obj.qdb.entries['/qubes-firewall/10.137.0.2/0000'] = b'action=accept'
+
+        original_multiread = self.obj.qdb.multiread
+        call_count = [0]
+
+        def multiread_with_interrupt(path):
+            call_count[0] += 1
+            if call_count[0] == 1:
+                raise InterruptedError
+            return original_multiread(path)
+
+        self.obj.qdb.multiread = multiread_with_interrupt
+
+        result = qubesagent.firewall.FirewallWorker.read_rules(
+            self.obj, '10.137.0.2')
+
+        self.assertEqual(call_count[0], 2)
+        self.assertIsNotNone(result)
+
+    def test_update_handled_retries_on_interrupted_error(self):
+        """update_handled() must retry qdb.read() on InterruptedError."""
+        self.obj.qdb.entries['/qubes-firewall-handled/10.137.0.2'] = b'5'
+
+        original_read = self.obj.qdb.read
+        call_count = [0]
+
+        def read_with_interrupt(key):
+            call_count[0] += 1
+            if call_count[0] == 1:
+                raise InterruptedError
+            return original_read(key)
+
+        self.obj.qdb.read = read_with_interrupt
+
+        qubesagent.firewall.FirewallWorker.update_handled(
+            self.obj, '10.137.0.2')
+
+        self.assertEqual(call_count[0], 2)
+        self.assertEqual(
+            self.obj.qdb.entries['/qubes-firewall-handled/10.137.0.2'], '6')
+        
+    def test_log_error_retries_on_interrupted_error(self):
+        """log_error() must retry qdb.read('/default-user') on InterruptedError."""
+        self.obj.qdb.entries['/default-user'] = b'user'
+
+        original_read = self.obj.qdb.read
+        call_count = [0]
+
+        def read_with_interrupt(key):
+            call_count[0] += 1
+            if key == '/default-user' and call_count[0] == 1:
+                raise InterruptedError
+            return original_read(key)
+
+        self.obj.qdb.read = read_with_interrupt
+
+        with patch('subprocess.check_output'):
+            qubesagent.firewall.FirewallWorker.log_error(
+                self.obj, 'test error')
+
+        self.assertGreaterEqual(call_count[0], 2)