Merge pull request #5 from logos-co/test-dos-robustness

Test/spam protection

src/libs/common.py

@@ -40,5 +40,15 @@ def to_app_id(n: int) -> list:
     return list(n.to_bytes(32, byteorder="big"))
 
 
+def random_divide_k(n, k):
+    if n < k:
+        raise ValueError(f"n={n} must be at least k={k} to split into {k} parts")
+    cuts = sorted(random.sample(range(1, n), k - 1))
+    parts = [cuts[0]] + [cuts[i] - cuts[i - 1] for i in range(1, len(cuts))] + [n - cuts[-1]]
+    return parts
+
+
 def generate_random_bytes(n=31):
+    if n < 0:
+        raise ValueError("Input must be an unsigned integer (non-negative)")
     return os.urandom(n)

src/steps/da.py

@@ -85,28 +85,37 @@ class StepsDataAvailability(StepsCommon):
         return executor
 
     @allure.step
-    @retry(stop=stop_after_delay(65), wait=wait_fixed(1), reraise=True)
-    def disperse_data(self, data, app_id, index, utf8=True, padding=True):
-        response = []
-        request = prepare_dispersal_request(data, app_id, index, utf8=utf8, padding=padding)
-        executor = self.find_executor_node()
-        try:
-            response = executor.send_dispersal_request(request)
-        except Exception as ex:
-            assert "Bad Request" in str(ex) or "Internal Server Error" in str(ex)
+    def disperse_data(self, data, app_id, index, timeout_duration=65, utf8=True, padding=True):
+        @retry(stop=stop_after_delay(timeout_duration), wait=wait_fixed(1), reraise=True)
+        def disperse(my_self=self):
+            response = []
+            request = prepare_dispersal_request(data, app_id, index, utf8=utf8, padding=padding)
+            executor = my_self.find_executor_node()
+            try:
+                response = executor.send_dispersal_request(request)
+            except Exception as ex:
+                assert "Bad Request" in str(ex) or "Internal Server Error" in str(ex)
 
-        assert response.status_code == 200, "Send dispersal finished with unexpected response code"
+            assert hasattr(response, "status_code"), "Missing status_code"
+            assert response.status_code in (200, 429), "Unexpected status code"
+
+            return response
+
+        return disperse()
 
     @allure.step
-    @retry(stop=stop_after_delay(45), wait=wait_fixed(1), reraise=True)
-    def get_data_range(self, node, app_id, start, end):
-        response = []
-        query = prepare_get_range_request(app_id, start, end)
-        try:
-            response = node.send_get_data_range_request(query)
-        except Exception as ex:
-            assert "Bad Request" in str(ex) or "Internal Server Error" in str(ex)
+    def get_data_range(self, node, app_id, start, end, timeout_duration=45):
+        @retry(stop=stop_after_delay(timeout_duration), wait=wait_fixed(1), reraise=True)
+        def get_range():
+            response = []
+            query = prepare_get_range_request(app_id, start, end)
+            try:
+                response = node.send_get_data_range_request(query)
+            except Exception as ex:
+                assert "Bad Request" in str(ex) or "Internal Server Error" in str(ex)
 
-        assert response_contains_data(response), "Get data range response is empty"
+            assert response_contains_data(response), "Get data range response is empty"
 
-        return response
+            return response
+
+        return get_range()

tests/dos_robustness/test_spam_protection.py

@@ -0,0 +1,91 @@
+import random
+import pytest
+
+from src.libs.common import delay, to_app_id, to_index, random_divide_k, generate_random_bytes
+from src.libs.custom_logger import get_custom_logger
+from src.steps.da import StepsDataAvailability
+from src.test_data import DATA_TO_DISPERSE
+
+logger = get_custom_logger(__name__)
+
+
+@pytest.mark.usefixtures("setup_2_node_cluster")
+class TestSpamProtection(StepsDataAvailability):
+    main_nodes = []
+
+    def test_spam_protection_valid_uploads(self):
+        num_samples = len(DATA_TO_DISPERSE)
+        missing_dispersals = num_samples
+        for i in range(num_samples):
+            try:
+                response = self.disperse_data(DATA_TO_DISPERSE[i], to_app_id(1), to_index(0), timeout_duration=0)
+                if response.status_code == 200:
+                    missing_dispersals -= 1
+            except Exception as ex:
+                raise Exception(f"Dispersal #{i+1} was not successful with error {ex}")
+
+            delay(0.1)
+
+        assert missing_dispersals == 0, f"{missing_dispersals} dispersals were not successful"
+
+    def test_spam_protection_single_burst(self):
+        rate_limit = 1000
+        spam_per_burst = rate_limit + 10
+        successful_dispersals = 0
+        for i in range(spam_per_burst):
+            try:
+                response = self.disperse_data(DATA_TO_DISPERSE[0], to_app_id(1), to_index(0), timeout_duration=0)
+                if response.status_code == 429:
+                    break
+                else:
+                    successful_dispersals += 1
+            except Exception as ex:
+                raise Exception(f"Dispersal #{i+1} was not successful with error {ex}")
+
+        assert successful_dispersals <= rate_limit, "All consecutive dispersals were successful without any constraint"
+
+    def test_spam_protection_random_bytes_single_burst(self):
+        rate_limit = 1000
+        spam_per_burst = rate_limit + 10
+
+        n = random.randint(1, 256)
+        data_to_disperse = generate_random_bytes(n)
+
+        successful_dispersals = 0
+        for i in range(spam_per_burst):
+            try:
+                response = self.disperse_data(data_to_disperse, to_app_id(1), to_index(0), timeout_duration=0, utf8=False)
+                if response.status_code == 429:
+                    break
+                else:
+                    successful_dispersals += 1
+            except Exception as ex:
+                raise Exception(f"Dispersal #{i+1} was not successful with error {ex}")
+
+        assert successful_dispersals <= rate_limit, "All consecutive dispersals were successful without any constraint"
+
+    def test_spam_protection_multiple_bursts(self):
+        time_interval = 60
+        rate_limit = 1000
+        bursts = 3
+        spam_per_burst = int((rate_limit + 10) / bursts)
+
+        waiting_intervals = random_divide_k(time_interval, bursts)
+
+        logger.debug(f"Waiting intervals: {waiting_intervals}")
+
+        successful_dispersals = 0
+        for b in range(bursts):
+            for i in range(spam_per_burst):
+                try:
+                    response = self.disperse_data(DATA_TO_DISPERSE[7], to_app_id(1), to_index(0), timeout_duration=0)
+                    if response.status_code == 429:
+                        break
+                    else:
+                        successful_dispersals += 1
+                except Exception as ex:
+                    raise Exception(f"Dispersal #{i+1} was not successful with error {ex}")
+
+            delay(waiting_intervals[b])
+
+        assert successful_dispersals <= 1000, "All dispersals were successful without any constraint"

tests/networking_privacy/test_networking_privacy.py

@@ -23,7 +23,7 @@ class TestNetworkingPrivacy(StepsDataAvailability):
                 self.disperse_data(DATA_TO_DISPERSE[7], to_app_id(1), to_index(0))
                 successful_dispersals += 1
             except Exception as ex:
-                logger.warning(f"Dispersal #{i} was not successful with error {ex}")
+                logger.warning(f"Dispersal #{i+1} was not successful with error {ex}")
 
             if successful_dispersals == 10:
                 break