From ba198d40d868893fe8f8faa1e684c780da7fa0f3 Mon Sep 17 00:00:00 2001 From: unconst Date: Sat, 17 Jun 2023 14:38:45 -0600 Subject: [PATCH 1/5] initial commit --- bittensor/_axon/__init__.py | 36 ++------------------ bittensor/_synapse/synapse.py | 14 ++++++-- bittensor/_synapse/text_prompting/synapse.py | 16 +++++---- 3 files changed, 23 insertions(+), 43 deletions(-) diff --git a/bittensor/_axon/__init__.py b/bittensor/_axon/__init__.py index 01d656f0a9..02c182d28c 100644 --- a/bittensor/_axon/__init__.py +++ b/bittensor/_axon/__init__.py @@ -29,7 +29,7 @@ from dataclasses import dataclass from substrateinterface import Keypair import bittensor.utils.networking as net -from typing import Callable, Dict, Optional, Tuple, Union +from typing import Dict, Optional, Tuple class axon: """ Axon object for serving synapse receptors. """ @@ -60,7 +60,6 @@ def __init__( max_workers: Optional[int] = None, server: "grpc._server._Server" = None, maximum_concurrent_rpcs: Optional[int] = None, - blacklist: Optional[Callable] = None, ) -> "bittensor.Axon": r"""Creates a new bittensor.Axon object from passed arguments. Args: @@ -80,8 +79,6 @@ def __init__( Used to create the threadpool if not passed, specifies the number of active threads servicing requests. maximum_concurrent_rpcs (:type:`Optional[int]`, `optional`): Maximum allowed concurrently processed RPCs. - blacklist (:obj:`Optional[callable]`, `optional`): - function to blacklist requests. """ self.metagraph = metagraph self.wallet = wallet @@ -111,7 +108,6 @@ def __init__( self.external_ip = self.config.axon.external_ip if self.config.axon.external_ip != None else bittensor.utils.networking.get_external_ip() self.external_port = self.config.axon.external_port if self.config.axon.external_port != None else self.config.axon.port self.full_address = str(self.config.axon.ip) + ":" + str(self.config.axon.port) - self.blacklist = blacklist self.started = False # Build priority thread pool @@ -119,16 +115,14 @@ def __init__( # Build interceptor. self.receiver_hotkey = self.wallet.hotkey.ss58_address - self.auth_interceptor = AuthInterceptor( - receiver_hotkey=self.receiver_hotkey, blacklist=self.blacklist - ) + self.auth_interceptor = AuthInterceptor(receiver_hotkey=self.receiver_hotkey) # Build grpc server if server is None: self.thread_pool = futures.ThreadPoolExecutor(max_workers=self.config.axon.max_workers) self.server = grpc.server( self.thread_pool, - interceptors=(self.auth_interceptor,), + interceptors=(self.auth_interceptor), maximum_concurrent_rpcs=self.config.axon.maximum_concurrent_rpcs, options=[("grpc.keepalive_time_ms", 100000), ("grpc.keepalive_timeout_ms", 500000)], ) @@ -280,18 +274,14 @@ class AuthInterceptor(grpc.ServerInterceptor): def __init__( self, receiver_hotkey: str, - blacklist: Callable = None, ): r"""Creates a new server interceptor that authenticates incoming messages from passed arguments. Args: receiver_hotkey(str): the SS58 address of the hotkey which should be targeted by RPCs - black_list (Function, `optional`): - black list function that prevents certain pubkeys from sending messages """ super().__init__() self.nonces = {} - self.blacklist = blacklist self.receiver_hotkey = receiver_hotkey @@ -348,25 +338,8 @@ def check_signature( raise Exception("Signature mismatch") self.nonces[endpoint_key] = nonce - def black_list_checking(self, hotkey: str, method: str): - r"""Tries to call to blacklist function in the miner and checks if it should blacklist the pubkey""" - if self.blacklist is None: - return - - request_type = { - "/Bittensor/Forward": bittensor.proto.RequestType.FORWARD, - "/Bittensor/Backward": bittensor.proto.RequestType.BACKWARD, - }.get(method) - if request_type is None: - raise Exception("Unknown request type") - - failed, error_message = self.blacklist(hotkey, request_type) - if failed: - raise Exception(str(error_message)) - def intercept_service(self, continuation, handler_call_details): r"""Authentication between bittensor nodes. Intercepts messages and checks them""" - method = handler_call_details.method metadata = dict(handler_call_details.invocation_metadata) try: @@ -382,9 +355,6 @@ def intercept_service(self, continuation, handler_call_details): nonce, sender_hotkey, signature, receptor_uuid ) - # blacklist checking - self.black_list_checking(sender_hotkey, method) - return continuation(handler_call_details) except Exception as e: diff --git a/bittensor/_synapse/synapse.py b/bittensor/_synapse/synapse.py index 72e2d9ef3a..ef113d39f6 100644 --- a/bittensor/_synapse/synapse.py +++ b/bittensor/_synapse/synapse.py @@ -35,13 +35,22 @@ class SynapseCall( ABC ): def __init__( self, synapse: 'bittensor.Synapse', - request_proto: object + request_proto: object, + context: grpc.ServicerContext, ): + metadata = dict(context.invocation_metadata()) + ( + _, + sender_hotkey, + _, + _, + ) = synapse.axon.auth_interceptor.parse_signature(metadata) + self.completed = False self.start_time = time.time() self.timeout = request_proto.timeout self.src_version = request_proto.version - self.src_hotkey = request_proto.hotkey + self.src_hotkey = sender_hotkey self.dest_hotkey = synapse.axon.wallet.hotkey.ss58_address self.dest_version = bittensor.__version_as_int__ self.return_code: bittensor.proto.ReturnCode = bittensor.proto.ReturnCode.Success @@ -171,4 +180,3 @@ def apply( self, call: SynapseCall ) -> object: call.end() call.log_outbound() return call._get_response_proto() - diff --git a/bittensor/_synapse/text_prompting/synapse.py b/bittensor/_synapse/text_prompting/synapse.py index e2f544ac5f..9bdd3975b3 100644 --- a/bittensor/_synapse/text_prompting/synapse.py +++ b/bittensor/_synapse/text_prompting/synapse.py @@ -34,8 +34,9 @@ def __init__( synapse: "bittensor.TextPromptingSynapseMulti", request_proto: bittensor.proto.MultiForwardTextPromptingRequest, multi_forward_callback: Callable, + context: grpc.ServicerContext, ): - super().__init__( synapse = synapse, request_proto = request_proto ) + super().__init__( synapse = synapse, request_proto = request_proto, context = context ) self.messages: List[ Dict[str, str] ] = request_proto.messages self.formatted_messages = [ json.loads(message) for message in self.messages ] self.multi_forward_callback = multi_forward_callback @@ -67,8 +68,9 @@ def __init__( synapse: "TextPromptingSynapse", request_proto: bittensor.proto.ForwardTextPromptingRequest, forward_callback: Callable, + context: grpc.ServicerContext, ): - super().__init__( synapse = synapse, request_proto = request_proto ) + super().__init__( synapse = synapse, request_proto = request_proto, context = context ) self.messages = request_proto.messages self.formatted_messages = [ json.loads(message) for message in self.messages ] self.forward_callback = forward_callback @@ -99,8 +101,9 @@ def __init__( synapse: "TextPromptingSynapse", request_proto: bittensor.proto.BackwardTextPromptingRequest, backward_callback: Callable, + context: grpc.ServicerContext, ): - super().__init__( synapse = synapse, request_proto = request_proto ) + super().__init__( synapse = synapse, request_proto = request_proto, context = context ) self.formatted_messages = [ message for message in request_proto.messages ] self.formatted_rewards = torch.tensor( [ request_proto.rewards ], dtype = torch.float32 ) self.completion = request_proto.response @@ -140,17 +143,16 @@ def multi_forward( self, messages: List[Dict[str, str]] ) -> List[ str ]: ... def backward( self, messages: List[Dict[str, str]], response: str, rewards: torch.FloatTensor ) -> str: ... def Forward( self, request: bittensor.proto.ForwardTextPromptingRequest, context: grpc.ServicerContext ) -> bittensor.proto.ForwardTextPromptingResponse: - call = SynapseForward( self, request, self.forward ) + call = SynapseForward( self, request, self.forward, context ) bittensor.logging.trace( 'Forward: {} '.format( call ) ) return self.apply( call = call ) def MultiForward( self, request: bittensor.proto.MultiForwardTextPromptingRequest, context: grpc.ServicerContext ) -> bittensor.proto.MultiForwardTextPromptingResponse: - call = SynapseForwardMulti( self, request, self.multi_forward ) + call = SynapseForwardMulti( self, request, self.multi_forward, context ) bittensor.logging.trace( 'MultiForward: {} '.format( call ) ) return self.apply( call = call ) def Backward( self, request: bittensor.proto.BackwardTextPromptingRequest, context: grpc.ServicerContext ) -> bittensor.proto.BackwardTextPromptingResponse: - call = SynapseBackward( self, request, self.backward ) + call = SynapseBackward( self, request, self.backward, context ) bittensor.logging.trace( 'Backward: {}'.format( call ) ) return self.apply( call = call ) - From 2583f08c5b71a00926e1bedb8f0ea99a2edab980 Mon Sep 17 00:00:00 2001 From: unconst Date: Sat, 17 Jun 2023 14:40:54 -0600 Subject: [PATCH 2/5] fix auth tuple --- bittensor/_axon/__init__.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/bittensor/_axon/__init__.py b/bittensor/_axon/__init__.py index 02c182d28c..d73fb3941c 100644 --- a/bittensor/_axon/__init__.py +++ b/bittensor/_axon/__init__.py @@ -122,7 +122,7 @@ def __init__( self.thread_pool = futures.ThreadPoolExecutor(max_workers=self.config.axon.max_workers) self.server = grpc.server( self.thread_pool, - interceptors=(self.auth_interceptor), + interceptors=(self.auth_interceptor,), maximum_concurrent_rpcs=self.config.axon.maximum_concurrent_rpcs, options=[("grpc.keepalive_time_ms", 100000), ("grpc.keepalive_timeout_ms", 500000)], ) From f98a55f764aa01659a881345dc61c25f4fbb4fcc Mon Sep 17 00:00:00 2001 From: unconst Date: Sat, 17 Jun 2023 14:53:57 -0600 Subject: [PATCH 3/5] fix tests that do not have the hotkey signature --- bittensor/_synapse/synapse.py | 24 ++++++++++++------- bittensor/_synapse/text_prompting/synapse.py | 6 ++--- .../bittensor_tests/test_synapse.py | 9 +++++-- 3 files changed, 26 insertions(+), 13 deletions(-) diff --git a/bittensor/_synapse/synapse.py b/bittensor/_synapse/synapse.py index ef113d39f6..4b08507a9b 100644 --- a/bittensor/_synapse/synapse.py +++ b/bittensor/_synapse/synapse.py @@ -36,15 +36,23 @@ def __init__( self, synapse: 'bittensor.Synapse', request_proto: object, - context: grpc.ServicerContext, + context: grpc.ServicerContext = None, ): - metadata = dict(context.invocation_metadata()) - ( - _, - sender_hotkey, - _, - _, - ) = synapse.axon.auth_interceptor.parse_signature(metadata) + # Optionally check context. + # If the context is not sent then this call does not check the signature. + # This allows spoofing attacks. + if context: + metadata = dict(context.invocation_metadata()) + ( + _, + sender_hotkey, + _, + _, + ) = synapse.axon.auth_interceptor.parse_signature(metadata) + else: + # If no context is sent, then the call is not authenticated. + # we use the hotkey from the request instead. + sender_hotkey = request_proto.hotkey self.completed = False self.start_time = time.time() diff --git a/bittensor/_synapse/text_prompting/synapse.py b/bittensor/_synapse/text_prompting/synapse.py index 9bdd3975b3..21f47219a6 100644 --- a/bittensor/_synapse/text_prompting/synapse.py +++ b/bittensor/_synapse/text_prompting/synapse.py @@ -34,7 +34,7 @@ def __init__( synapse: "bittensor.TextPromptingSynapseMulti", request_proto: bittensor.proto.MultiForwardTextPromptingRequest, multi_forward_callback: Callable, - context: grpc.ServicerContext, + context: grpc.ServicerContext = None ): super().__init__( synapse = synapse, request_proto = request_proto, context = context ) self.messages: List[ Dict[str, str] ] = request_proto.messages @@ -68,7 +68,7 @@ def __init__( synapse: "TextPromptingSynapse", request_proto: bittensor.proto.ForwardTextPromptingRequest, forward_callback: Callable, - context: grpc.ServicerContext, + context: grpc.ServicerContext = None ): super().__init__( synapse = synapse, request_proto = request_proto, context = context ) self.messages = request_proto.messages @@ -101,7 +101,7 @@ def __init__( synapse: "TextPromptingSynapse", request_proto: bittensor.proto.BackwardTextPromptingRequest, backward_callback: Callable, - context: grpc.ServicerContext, + context: grpc.ServicerContext = None ): super().__init__( synapse = synapse, request_proto = request_proto, context = context ) self.formatted_messages = [ message for message in request_proto.messages ] diff --git a/tests/unit_tests/bittensor_tests/test_synapse.py b/tests/unit_tests/bittensor_tests/test_synapse.py index 44b4cd7580..5f8ce43def 100644 --- a/tests/unit_tests/bittensor_tests/test_synapse.py +++ b/tests/unit_tests/bittensor_tests/test_synapse.py @@ -94,14 +94,19 @@ def test_text_prompting_synapse_backward(): def test_text_prompting_synapse_blacklist(): synapse = get_synapse() - request = bittensor.proto.ForwardTextPromptingRequest() + request = bittensor.proto.ForwardTextPromptingRequest( + hotkey = "5DD26kC2kxajmwfbbZmVmxhrY9VeeyR1Gpzy9i8wxLUg6zxm" + ) + context = MagicMock() call = bittensor._synapse.text_prompting.synapse.SynapseForward( synapse, request, synapse.forward ) blacklist = synapse.blacklist( call ) assert blacklist == False def test_text_prompting_synapse_priority(): synapse = get_synapse() - request = bittensor.proto.ForwardTextPromptingRequest() + request = bittensor.proto.ForwardTextPromptingRequest( + hotkey = "5DD26kC2kxajmwfbbZmVmxhrY9VeeyR1Gpzy9i8wxLUg6zxm" + ) call = bittensor._synapse.text_prompting.synapse.SynapseForward( synapse, request, synapse.forward ) priority = synapse.priority( call ) assert priority == 0.0 From b370326be2d5951fdb51eddcbe0e59a52520d747 Mon Sep 17 00:00:00 2001 From: unconst Date: Sat, 17 Jun 2023 15:02:31 -0600 Subject: [PATCH 4/5] fix to tests with mock instead of allowing None value --- bittensor/_synapse/synapse.py | 26 +++++++----------- bittensor/_synapse/text_prompting/synapse.py | 6 ++--- .../bittensor_tests/test_synapse.py | 27 +++++++++++++------ 3 files changed, 31 insertions(+), 28 deletions(-) diff --git a/bittensor/_synapse/synapse.py b/bittensor/_synapse/synapse.py index 4b08507a9b..5a4dbd4225 100644 --- a/bittensor/_synapse/synapse.py +++ b/bittensor/_synapse/synapse.py @@ -36,24 +36,16 @@ def __init__( self, synapse: 'bittensor.Synapse', request_proto: object, - context: grpc.ServicerContext = None, + context: grpc.ServicerContext, ): - # Optionally check context. - # If the context is not sent then this call does not check the signature. - # This allows spoofing attacks. - if context: - metadata = dict(context.invocation_metadata()) - ( - _, - sender_hotkey, - _, - _, - ) = synapse.axon.auth_interceptor.parse_signature(metadata) - else: - # If no context is sent, then the call is not authenticated. - # we use the hotkey from the request instead. - sender_hotkey = request_proto.hotkey - + metadata = dict(context.invocation_metadata()) + ( + _, + sender_hotkey, + _, + _, + ) = synapse.axon.auth_interceptor.parse_signature(metadata) + self.completed = False self.start_time = time.time() self.timeout = request_proto.timeout diff --git a/bittensor/_synapse/text_prompting/synapse.py b/bittensor/_synapse/text_prompting/synapse.py index 21f47219a6..77af418f22 100644 --- a/bittensor/_synapse/text_prompting/synapse.py +++ b/bittensor/_synapse/text_prompting/synapse.py @@ -34,7 +34,7 @@ def __init__( synapse: "bittensor.TextPromptingSynapseMulti", request_proto: bittensor.proto.MultiForwardTextPromptingRequest, multi_forward_callback: Callable, - context: grpc.ServicerContext = None + context: grpc.ServicerContext ): super().__init__( synapse = synapse, request_proto = request_proto, context = context ) self.messages: List[ Dict[str, str] ] = request_proto.messages @@ -68,7 +68,7 @@ def __init__( synapse: "TextPromptingSynapse", request_proto: bittensor.proto.ForwardTextPromptingRequest, forward_callback: Callable, - context: grpc.ServicerContext = None + context: grpc.ServicerContext ): super().__init__( synapse = synapse, request_proto = request_proto, context = context ) self.messages = request_proto.messages @@ -101,7 +101,7 @@ def __init__( synapse: "TextPromptingSynapse", request_proto: bittensor.proto.BackwardTextPromptingRequest, backward_callback: Callable, - context: grpc.ServicerContext = None + context: grpc.ServicerContext ): super().__init__( synapse = synapse, request_proto = request_proto, context = context ) self.formatted_messages = [ message for message in request_proto.messages ] diff --git a/tests/unit_tests/bittensor_tests/test_synapse.py b/tests/unit_tests/bittensor_tests/test_synapse.py index 5f8ce43def..6b45db57a3 100644 --- a/tests/unit_tests/bittensor_tests/test_synapse.py +++ b/tests/unit_tests/bittensor_tests/test_synapse.py @@ -94,20 +94,31 @@ def test_text_prompting_synapse_backward(): def test_text_prompting_synapse_blacklist(): synapse = get_synapse() - request = bittensor.proto.ForwardTextPromptingRequest( - hotkey = "5DD26kC2kxajmwfbbZmVmxhrY9VeeyR1Gpzy9i8wxLUg6zxm" - ) + request = bittensor.proto.ForwardTextPromptingRequest() + + # Mock the signature checking of the context. context = MagicMock() - call = bittensor._synapse.text_prompting.synapse.SynapseForward( synapse, request, synapse.forward ) + context.invocation_metadata.return_value = {} + synapse.axon = MagicMock() + synapse.axon.auth_interceptor = MagicMock() + synapse.axon.auth_interceptor.parse_signature.return_value = (None, None, "5CtstubuSoVLJGCXkiWRNKrrGg2DVBZ9qMs2qYTLsZR4q1Wg", None) + + call = bittensor._synapse.text_prompting.synapse.SynapseForward( synapse, request, synapse.forward, context = context ) blacklist = synapse.blacklist( call ) assert blacklist == False def test_text_prompting_synapse_priority(): synapse = get_synapse() - request = bittensor.proto.ForwardTextPromptingRequest( - hotkey = "5DD26kC2kxajmwfbbZmVmxhrY9VeeyR1Gpzy9i8wxLUg6zxm" - ) - call = bittensor._synapse.text_prompting.synapse.SynapseForward( synapse, request, synapse.forward ) + request = bittensor.proto.ForwardTextPromptingRequest() + + # Mock the signature checking of the context. + context = MagicMock() + context.invocation_metadata.return_value = {} + synapse.axon = MagicMock() + synapse.axon.auth_interceptor = MagicMock() + synapse.axon.auth_interceptor.parse_signature.return_value = (None, None, "5CtstubuSoVLJGCXkiWRNKrrGg2DVBZ9qMs2qYTLsZR4q1Wg", None) + + call = bittensor._synapse.text_prompting.synapse.SynapseForward( synapse, request, synapse.forward, context = context ) priority = synapse.priority( call ) assert priority == 0.0 From 7f24987ece0de258158f741d5195e7450ce2bcdd Mon Sep 17 00:00:00 2001 From: ifrit98 Date: Sun, 18 Jun 2023 00:00:50 +0000 Subject: [PATCH 5/5] update bittensor version into to current (510) --- bittensor/_axon/__init__.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/bittensor/_axon/__init__.py b/bittensor/_axon/__init__.py index d73fb3941c..188268caeb 100644 --- a/bittensor/_axon/__init__.py +++ b/bittensor/_axon/__init__.py @@ -305,7 +305,7 @@ def parse_signature(self, metadata: Dict[str, str]) -> Tuple[int, str, str, str] version = metadata.get('bittensor-version') if signature is None: raise Exception("Request signature missing") - if int(version) < 370: + if int(version) < 510: raise Exception("Incorrect Version") parts = self.parse_signature_v2(signature) if parts is not None: