Merge pull request #982 from HathorNetwork/feat/arg-to-make-sync-v1-unavailable

jansegre · web-flow · commit 007eb6e0dfc4 · 2024-05-03T04:09:06.000+02:00
feat(cli): add argument to make sync-v1 unavailable
diff --git a/hathor/builder/__init__.py b/hathor/builder/__init__.py
@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from hathor.builder.builder import BuildArtifacts, Builder
+from hathor.builder.builder import BuildArtifacts, Builder, SyncSupportLevel
 from hathor.builder.cli_builder import CliBuilder
 from hathor.builder.resources_builder import ResourcesBuilder
 
@@ -21,4 +21,5 @@
     'Builder',
     'CliBuilder',
     'ResourcesBuilder',
+    'SyncSupportLevel',
 ]
diff --git a/hathor/builder/builder.py b/hathor/builder/builder.py
@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from enum import Enum
+from enum import Enum, IntEnum
 from typing import Any, Callable, NamedTuple, Optional, TypeAlias
 
 from structlog import get_logger
@@ -55,6 +55,34 @@
 logger = get_logger()
 
 
+class SyncSupportLevel(IntEnum):
+    UNAVAILABLE = 0  # not possible to enable at runtime
+    DISABLED = 1  # available but disabled by default, possible to enable at runtime
+    ENABLED = 2  # available and enabled by default, possible to disable at runtime
+
+    @classmethod
+    def add_factories(cls,
+                      p2p_manager: ConnectionsManager,
+                      sync_v1_support: 'SyncSupportLevel',
+                      sync_v2_support: 'SyncSupportLevel',
+                      ) -> None:
+        """Adds the sync factory to the manager according to the support level."""
+        from hathor.p2p.sync_v1.factory import SyncV11Factory
+        from hathor.p2p.sync_v2.factory import SyncV2Factory
+        from hathor.p2p.sync_version import SyncVersion
+
+        # sync-v1 support:
+        if sync_v1_support > cls.UNAVAILABLE:
+            p2p_manager.add_sync_factory(SyncVersion.V1_1, SyncV11Factory(p2p_manager))
+        if sync_v1_support is cls.ENABLED:
+            p2p_manager.enable_sync_version(SyncVersion.V1_1)
+        # sync-v2 support:
+        if sync_v2_support > cls.UNAVAILABLE:
+            p2p_manager.add_sync_factory(SyncVersion.V2, SyncV2Factory(p2p_manager))
+        if sync_v2_support is cls.ENABLED:
+            p2p_manager.enable_sync_version(SyncVersion.V2)
+
+
 class StorageType(Enum):
     MEMORY = 'memory'
     ROCKSDB = 'rocksdb'
@@ -147,8 +175,8 @@ def __init__(self) -> None:
         self._enable_tokens_index: bool = False
         self._enable_utxo_index: bool = False
 
-        self._enable_sync_v1: bool = True
-        self._enable_sync_v2: bool = False
+        self._sync_v1_support: SyncSupportLevel = SyncSupportLevel.UNAVAILABLE
+        self._sync_v2_support: SyncSupportLevel = SyncSupportLevel.UNAVAILABLE
 
         self._enable_stratum_server: Optional[bool] = None
 
@@ -168,6 +196,9 @@ def build(self) -> BuildArtifacts:
         if self._network is None:
             raise TypeError('you must set a network')
 
+        if SyncSupportLevel.ENABLED not in {self._sync_v1_support, self._sync_v2_support}:
+            raise TypeError('you must enable at least one sync version')
+
         settings = self._get_or_create_settings()
         reactor = self._get_reactor()
         pubsub = self._get_or_create_pubsub()
@@ -368,10 +399,6 @@ def _get_or_create_p2p_manager(self) -> ConnectionsManager:
         if self._p2p_manager:
             return self._p2p_manager
 
-        from hathor.p2p.sync_v1.factory import SyncV11Factory
-        from hathor.p2p.sync_v2.factory import SyncV2Factory
-        from hathor.p2p.sync_version import SyncVersion
-
         enable_ssl = True
         reactor = self._get_reactor()
         my_peer = self._get_peer_id()
@@ -387,12 +414,7 @@ def _get_or_create_p2p_manager(self) -> ConnectionsManager:
             whitelist_only=False,
             rng=self._rng,
         )
-        self._p2p_manager.add_sync_factory(SyncVersion.V1_1, SyncV11Factory(self._p2p_manager))
-        self._p2p_manager.add_sync_factory(SyncVersion.V2, SyncV2Factory(self._p2p_manager))
-        if self._enable_sync_v1:
-            self._p2p_manager.enable_sync_version(SyncVersion.V1_1)
-        if self._enable_sync_v2:
-            self._p2p_manager.enable_sync_version(SyncVersion.V2)
+        SyncSupportLevel.add_factories(self._p2p_manager, self._sync_v1_support, self._sync_v2_support)
         return self._p2p_manager
 
     def _get_or_create_indexes_manager(self) -> IndexesManager:
@@ -698,34 +720,34 @@ def set_network(self, network: str) -> 'Builder':
         self._network = network
         return self
 
-    def set_enable_sync_v1(self, enable_sync_v1: bool) -> 'Builder':
+    def set_sync_v1_support(self, support_level: SyncSupportLevel) -> 'Builder':
         self.check_if_can_modify()
-        self._enable_sync_v1 = enable_sync_v1
+        self._sync_v1_support = support_level
         return self
 
-    def set_enable_sync_v2(self, enable_sync_v2: bool) -> 'Builder':
+    def set_sync_v2_support(self, support_level: SyncSupportLevel) -> 'Builder':
         self.check_if_can_modify()
-        self._enable_sync_v2 = enable_sync_v2
+        self._sync_v2_support = support_level
         return self
 
     def enable_sync_v1(self) -> 'Builder':
         self.check_if_can_modify()
-        self._enable_sync_v1 = True
+        self._sync_v1_support = SyncSupportLevel.ENABLED
         return self
 
     def disable_sync_v1(self) -> 'Builder':
         self.check_if_can_modify()
-        self._enable_sync_v1 = False
+        self._sync_v1_support = SyncSupportLevel.DISABLED
         return self
 
     def enable_sync_v2(self) -> 'Builder':
         self.check_if_can_modify()
-        self._enable_sync_v2 = True
+        self._sync_v2_support = SyncSupportLevel.ENABLED
         return self
 
     def disable_sync_v2(self) -> 'Builder':
         self.check_if_can_modify()
-        self._enable_sync_v2 = False
+        self._sync_v2_support = SyncSupportLevel.DISABLED
         return self
 
     def set_full_verification(self, full_verification: bool) -> 'Builder':
diff --git a/hathor/builder/cli_builder.py b/hathor/builder/cli_builder.py
@@ -49,9 +49,10 @@
 
 
 class SyncChoice(Enum):
-    V1_ONLY = auto()
-    V2_ONLY = auto()
-    BRIDGE = auto()
+    V1_DEFAULT = auto()  # v1 enabled, v2 disabled but can be enabled in runtime
+    V2_DEFAULT = auto()  # v2 enabled, v1 disabled but can be enabled in runtime
+    BRIDGE_DEFAULT = auto()  # both enabled, either can be disabled in runtime
+    V2_ONLY = auto()  # v1 is unavailable, it cannot be enabled in runtime
 
 
 class CliBuilder:
@@ -70,15 +71,13 @@ def check_or_raise(self, condition: bool, message: str) -> None:
 
     def create_manager(self, reactor: Reactor) -> HathorManager:
         import hathor
+        from hathor.builder import SyncSupportLevel
         from hathor.conf.get_settings import get_global_settings, get_settings_source
         from hathor.daa import TestMode
         from hathor.event.storage import EventMemoryStorage, EventRocksDBStorage, EventStorage
         from hathor.event.websocket.factory import EventWebsocketFactory
         from hathor.p2p.netfilter.utils import add_peer_id_blacklist
         from hathor.p2p.peer_discovery import BootstrapPeerDiscovery, DNSPeerDiscovery
-        from hathor.p2p.sync_v1.factory import SyncV11Factory
-        from hathor.p2p.sync_v2.factory import SyncV2Factory
-        from hathor.p2p.sync_version import SyncVersion
         from hathor.storage import RocksDBStorage
         from hathor.transaction.storage import (
             TransactionCacheStorage,
@@ -177,34 +176,39 @@ def create_manager(self, reactor: Reactor) -> HathorManager:
 
         sync_choice: SyncChoice
         if self._args.sync_bridge:
-            sync_choice = SyncChoice.BRIDGE
+            sync_choice = SyncChoice.BRIDGE_DEFAULT
         elif self._args.sync_v1_only:
-            sync_choice = SyncChoice.V1_ONLY
+            sync_choice = SyncChoice.V1_DEFAULT
         elif self._args.sync_v2_only:
             self.log.warn('--sync-v2-only is the default, this parameter has no effect')
+            sync_choice = SyncChoice.V2_DEFAULT
+        elif self._args.x_remove_sync_v1:
             sync_choice = SyncChoice.V2_ONLY
         elif self._args.x_sync_bridge:
             self.log.warn('--x-sync-bridge is deprecated and will be removed, use --sync-bridge instead')
-            sync_choice = SyncChoice.BRIDGE
+            sync_choice = SyncChoice.BRIDGE_DEFAULT
         elif self._args.x_sync_v2_only:
             self.log.warn('--x-sync-v2-only is deprecated and will be removed, use --sync-v2-only instead')
-            sync_choice = SyncChoice.V2_ONLY
+            sync_choice = SyncChoice.V2_DEFAULT
         else:
             # XXX: this is the default behavior when no parameter is given
-            sync_choice = SyncChoice.V2_ONLY
+            sync_choice = SyncChoice.V2_DEFAULT
 
-        enable_sync_v1: bool
-        enable_sync_v2: bool
+        sync_v1_support: SyncSupportLevel
+        sync_v2_support: SyncSupportLevel
         match sync_choice:
-            case SyncChoice.V1_ONLY:
-                enable_sync_v1 = True
-                enable_sync_v2 = False
+            case SyncChoice.V1_DEFAULT:
+                sync_v1_support = SyncSupportLevel.ENABLED
+                sync_v2_support = SyncSupportLevel.DISABLED
+            case SyncChoice.V2_DEFAULT:
+                sync_v1_support = SyncSupportLevel.DISABLED
+                sync_v2_support = SyncSupportLevel.ENABLED
+            case SyncChoice.BRIDGE_DEFAULT:
+                sync_v1_support = SyncSupportLevel.ENABLED
+                sync_v2_support = SyncSupportLevel.ENABLED
             case SyncChoice.V2_ONLY:
-                enable_sync_v1 = False
-                enable_sync_v2 = True
-            case SyncChoice.BRIDGE:
-                enable_sync_v1 = True
-                enable_sync_v2 = True
+                sync_v1_support = SyncSupportLevel.UNAVAILABLE
+                sync_v2_support = SyncSupportLevel.ENABLED
 
         pubsub = PubSubManager(reactor)
 
@@ -293,12 +297,7 @@ def create_manager(self, reactor: Reactor) -> HathorManager:
             whitelist_only=False,
             rng=Random(),
         )
-        p2p_manager.add_sync_factory(SyncVersion.V1_1, SyncV11Factory(p2p_manager))
-        p2p_manager.add_sync_factory(SyncVersion.V2, SyncV2Factory(p2p_manager))
-        if enable_sync_v1:
-            p2p_manager.enable_sync_version(SyncVersion.V1_1)
-        if enable_sync_v2:
-            p2p_manager.enable_sync_version(SyncVersion.V2)
+        SyncSupportLevel.add_factories(p2p_manager, sync_v1_support, sync_v2_support)
 
         vertex_handler = VertexHandler(
             reactor=reactor,
diff --git a/hathor/cli/run_node.py b/hathor/cli/run_node.py
@@ -131,6 +131,8 @@ def create_parser(cls) -> ArgumentParser:
                                help='Enable running both sync protocols.')
         sync_args.add_argument('--sync-v1-only', action='store_true', help='Disable support for running sync-v2.')
         sync_args.add_argument('--sync-v2-only', action='store_true', help='Disable support for running sync-v1.')
+        sync_args.add_argument('--x-remove-sync-v1', action='store_true', help='Make sync-v1 unavailable, thus '
+                               'impossible to be enable in runtime.')
         sync_args.add_argument('--x-sync-v2-only', action='store_true', help=SUPPRESS)  # old argument
         sync_args.add_argument('--x-sync-bridge', action='store_true', help=SUPPRESS)  # old argument
         parser.add_argument('--x-localhost-only', action='store_true', help='Only connect to peers on localhost')
diff --git a/hathor/cli/run_node_args.py b/hathor/cli/run_node_args.py
@@ -65,6 +65,7 @@ class RunNodeArgs(BaseModel, extra=Extra.allow):
     enable_crash_api: bool
     x_sync_bridge: bool
     x_sync_v2_only: bool
+    x_remove_sync_v1: bool
     sync_bridge: bool
     sync_v1_only: bool
     sync_v2_only: bool
diff --git a/tests/simulation/base.py b/tests/simulation/base.py
@@ -1,5 +1,6 @@
 from typing import Optional
 
+from hathor.builder import SyncSupportLevel
 from hathor.manager import HathorManager
 from hathor.simulator import Simulator
 from hathor.types import VertexId
@@ -41,13 +42,15 @@ def create_peer(  # type: ignore[override]
                                                       'the test class or pass `enable_sync_v2` by argument')
             enable_sync_v2 = self._enable_sync_v2
         assert enable_sync_v1 or enable_sync_v2, 'enable at least one sync version'
+        sync_v1_support = SyncSupportLevel.ENABLED if enable_sync_v1 else SyncSupportLevel.DISABLED
+        sync_v2_support = SyncSupportLevel.ENABLED if enable_sync_v2 else SyncSupportLevel.DISABLED
         if simulator is None:
             simulator = self.simulator
 
         builder = simulator.get_default_builder() \
             .set_peer_id(self.get_random_peer_id_from_pool(rng=simulator.rng)) \
             .set_soft_voided_tx_ids(soft_voided_tx_ids) \
-            .set_enable_sync_v1(enable_sync_v1) \
-            .set_enable_sync_v2(enable_sync_v2)
+            .set_sync_v1_support(sync_v1_support) \
+            .set_sync_v2_support(sync_v2_support)
 
         return simulator.create_peer(builder)
diff --git a/tests/simulation/test_simulator_itself.py b/tests/simulation/test_simulator_itself.py
@@ -1,5 +1,6 @@
 import pytest
 
+from hathor.builder import SyncSupportLevel
 from hathor.manager import HathorManager
 from hathor.p2p.peer_id import PeerId
 from hathor.simulator import FakeConnection, Simulator
@@ -54,11 +55,13 @@ def create_simulator_peer(
                                                       'the test class or pass `enable_sync_v2` by argument')
             enable_sync_v2 = self._enable_sync_v2
         assert enable_sync_v1 or enable_sync_v2, 'enable at least one sync version'
+        sync_v1_support = SyncSupportLevel.ENABLED if enable_sync_v1 else SyncSupportLevel.DISABLED
+        sync_v2_support = SyncSupportLevel.ENABLED if enable_sync_v2 else SyncSupportLevel.DISABLED
 
         builder = simulator.get_default_builder() \
             .set_peer_id(self.get_random_peer_id_from_pool()) \
-            .set_enable_sync_v1(enable_sync_v1) \
-            .set_enable_sync_v2(enable_sync_v2)
+            .set_sync_v1_support(sync_v1_support) \
+            .set_sync_v2_support(sync_v2_support)
 
         return simulator.create_peer(builder)
 
diff --git a/tests/unittest.py b/tests/unittest.py
@@ -87,6 +87,8 @@ class TestBuilder(Builder):
     def __init__(self) -> None:
         super().__init__()
         self.set_network('testnet')
+        # default builder has sync-v2 enabled for tests
+        self.enable_sync_v2()
 
     def build(self) -> BuildArtifacts:
         artifacts = super().build()
