diff --git a/bittensor/axon.py b/bittensor/axon.py index d91039fb91..c7df06363e 100644 --- a/bittensor/axon.py +++ b/bittensor/axon.py @@ -24,6 +24,7 @@ import copy import json import time +import base64 import asyncio import inspect import uvicorn @@ -509,39 +510,36 @@ async def some_endpoint(body_dict: dict = Depends(verify_body_integrity)): # The function only executes if the body integrity verification is successful. ... """ - # Properly extract keys and values that have 'hash' - hash_headers = { - "_".join(k.split("_")[k.split("_").index("hash") + 1 :]): v - for k, v in request.headers.items() - if "_hash_" in k - } - # Await and load the request body so we can inspect it body = await request.body() request_body = body.decode() if isinstance(body, bytes) else body + # Gather the required field names from the headers of the request + required_hash_fields = json.loads( + base64.b64decode(request.headers.get("hash_fields", "").encode()).decode( + "utf-8" + ) + ) + # Load the body dict and check if all required field hashes match body_dict = json.loads(request_body) - for required_field in list(hash_headers): + field_hashes = [] + for required_field in required_hash_fields: # Hash the field in the body to compare against the header hashes body_value = body_dict.get(required_field, None) if body_value == None: - return JSONResponse( - content={"error": f"Missing required field {required_field}"}, - status_code=400, - ) + raise ValueError(f"Missing required field {required_field}") field_hash = bittensor.utils.hash(str(body_value)) + field_hashes.append(field_hash) - # If any hashes fail to match up, return a 400 error as the body is invalid - if field_hash != hash_headers[required_field]: - return JSONResponse( - content={ - "error": f"Hash mismatch with {field_hash} and {getattr(synapse, required_field + '_hash')}" - }, - status_code=400, - ) + parsed_body_hash = bittensor.utils.hash("".join(field_hashes)) + body_hash = request.headers.get("computed_body_hash", "") + if parsed_body_hash != body_hash: + raise ValueError( + f"Hash mismatch between header body hash {body_hash} and parsed body hash {parsed_body_hash}" + ) - # If body is good, return the parsed body so that it can be injected into the route function + # If body is good, return the parsed body so that it can be passed onto the route function return body_dict @classmethod @@ -653,13 +651,8 @@ def default_verify(self, synapse: bittensor.Synapse): # Build the keypair from the dendrite_hotkey keypair = Keypair(ss58_address=synapse.dendrite.hotkey) - # Pull body hashes from synapse recieved with request. - body_hashes = [ - getattr(synapse, field + "_hash") for field in synapse.required_hash_fields - ] - # Build the signature messages. - message = f"{synapse.dendrite.nonce}.{synapse.dendrite.hotkey}.{self.wallet.hotkey.ss58_address}.{synapse.dendrite.uuid}.{body_hashes}" + message = f"{synapse.dendrite.nonce}.{synapse.dendrite.hotkey}.{self.wallet.hotkey.ss58_address}.{synapse.dendrite.uuid}.{synapse.computed_body_hash}" # Build the unique endpoint key. endpoint_key = f"{synapse.dendrite.hotkey}:{synapse.dendrite.uuid}" diff --git a/bittensor/commands/network.py b/bittensor/commands/network.py index d67d3d5db7..f1a9a6f428 100644 --- a/bittensor/commands/network.py +++ b/bittensor/commands/network.py @@ -166,7 +166,7 @@ def add_args(parser: argparse.ArgumentParser): "immunity_period": "sudo_set_immunity_period", "min_allowed_weights": "sudo_set_min_allowed_weights", "activity_cutoff": "sudo_set_activity_cutoff", - "max_validators": "sudo_set_max_allowed_validators" + "max_validators": "sudo_set_max_allowed_validators", } diff --git a/bittensor/commands/root.py b/bittensor/commands/root.py index 099c79cd3c..55bea86b50 100644 --- a/bittensor/commands/root.py +++ b/bittensor/commands/root.py @@ -121,8 +121,10 @@ def run(cli): if neuron_data.hotkey in delegate_info else "", neuron_data.hotkey, - "{:.5f}".format(float(subtensor.get_total_stake_for_hotkey(neuron_data.hotkey))), - "Yes" if neuron_data.hotkey in senate_members else "No" + "{:.5f}".format( + float(subtensor.get_total_stake_for_hotkey(neuron_data.hotkey)) + ), + "Yes" if neuron_data.hotkey in senate_members else "No", ) table.box = None diff --git a/bittensor/dendrite.py b/bittensor/dendrite.py index eb9c020824..e831fa0d02 100644 --- a/bittensor/dendrite.py +++ b/bittensor/dendrite.py @@ -347,41 +347,12 @@ def preprocess_synapse_for_request( } ) - # Sign the request using the dendrite, axon info, and the synapse body hashes - body_hashes = self.hash_synapse_body(synapse) - - message = f"{synapse.dendrite.nonce}.{synapse.dendrite.hotkey}.{synapse.axon.hotkey}.{synapse.dendrite.uuid}.{body_hashes}" + # Sign the request using the dendrite, axon info, and the synapse body hash + message = f"{synapse.dendrite.nonce}.{synapse.dendrite.hotkey}.{synapse.axon.hotkey}.{synapse.dendrite.uuid}.{synapse.body_hash}" synapse.dendrite.signature = f"0x{self.keypair.sign(message).hex()}" return synapse - @staticmethod - def hash_synapse_body(synapse: bittensor.Synapse) -> str: - """ - Compute a SHA-256 hash of the serialized body of the Synapse instance. - - The body of the Synapse instance comprises its serialized and encoded - non-optional fields. This property retrieves these fields using the - `get_body` method, then concatenates their string representations, and - finally computes a SHA-256 hash of the resulting string. - - Returns: - str: The hexadecimal representation of the SHA-256 hash of the instance's body. - """ - # Hash the body for verification - hashes = [] - - # Getting the fields of the instance - instance_fields = synapse.__dict__ - - # Iterating over the fields of the instance - for field, value in instance_fields.items(): - # If the field is required in the subclass schema, add it. - if field in synapse.required_hash_fields: - hashes.append(bittensor.utils.hash(str(value))) - - return hashes - def process_server_response( self, server_response: Response, diff --git a/bittensor/synapse.py b/bittensor/synapse.py index 20a9095e2e..9e43b5cb58 100644 --- a/bittensor/synapse.py +++ b/bittensor/synapse.py @@ -297,6 +297,24 @@ def set_name_type(cls, values) -> dict: repr=False, ) + computed_body_hash: Optional[str] = pydantic.Field( + title="computed_body_hash", + description="The computed body hash of the request.", + examples="0x0813029319030129u4120u10841824y0182u091u230912u", + default="", + allow_mutation=False, + repr=False, + ) + + hash_fields: Optional[List[str]] = pydantic.Field( + title="hash_fields", + description="The list of required fields to compute the body hash.", + examples=["roles", "messages"], + default=[], + allow_mutation=False, + repr=False, + ) + def __setattr__(self, name: str, value: Any): """ Override the __setattr__ method to make the `required_hash_fields` property read-only. @@ -305,6 +323,10 @@ def __setattr__(self, name: str, value: Any): raise AttributeError( "required_hash_fields property is read-only and cannot be overridden." ) + if name == "body_hash": + raise AttributeError( + "body_hash property is read-only and cannot be overridden." + ) super().__setattr__(name, value) @property @@ -505,15 +527,12 @@ def to_headers(self) -> dict: elif required and field in required: try: - # Create an empty (dummy) instance of type(value) to pass pydantic validation on the axon side + # create an empty (dummy) instance of type(value) to pass pydantic validation on the axon side serialized_value = json.dumps(value.__class__.__call__()) - # Create a hash of the original data so we can verify on the axon side - hash_value = bittensor.utils.hash(str(value)) encoded_value = base64.b64encode(serialized_value.encode()).decode( "utf-8" ) headers[f"bt_header_input_obj_{field}"] = encoded_value - headers[f"bt_header_input_hash_{field}"] = hash_value except TypeError as e: raise ValueError( f"Error serializing {field} with value {value}. Objects must be json serializable." @@ -522,9 +541,40 @@ def to_headers(self) -> dict: # Adding the size of the headers and the total size to the headers headers["header_size"] = str(sys.getsizeof(headers)) headers["total_size"] = str(self.get_total_size()) - + headers["computed_body_hash"] = self.body_hash + headers["hash_fields"] = base64.b64encode( + json.dumps(self.required_hash_fields).encode() + ).decode("utf-8") return headers + @property + def body_hash(self) -> str: + """ + Compute a SHA-256 hash of the serialized body of the Synapse instance. + + The body of the Synapse instance comprises its serialized and encoded + non-optional fields. This property retrieves these fields using the + `get_body` method, then concatenates their string representations, and + finally computes a SHA-256 hash of the resulting string. + + Returns: + str: The hexadecimal representation of the SHA-256 hash of the instance's body. + """ + # Hash the body for verification + hashes = [] + + # Getting the fields of the instance + instance_fields = self.__dict__ + + # Iterating over the fields of the instance + for field, value in instance_fields.items(): + # If the field is required in the subclass schema, add it. + if field in self.required_hash_fields: + hashes.append(bittensor.utils.hash(str(value))) + + # Hash and return the hashes that have been concatenated + return bittensor.utils.hash("".join(hashes)) + @classmethod def parse_headers_to_inputs(cls, headers: dict) -> dict: """ @@ -630,27 +680,20 @@ def parse_headers_to_inputs(cls, headers: dict) -> dict: f"Error while parsing 'input_obj' header {key}: {e}" ) continue - elif "bt_header_input_hash" in key: - try: - new_key = key.split("bt_header_input_hash_")[1] + "_hash" - # Skip if the key already exists in the dictionary - if new_key in inputs_dict: - continue - # Decode and load the serialized object - inputs_dict[new_key] = value - except Exception as e: - bittensor.logging.error( - f"Error while parsing 'input_hash' header {key}: {e}" - ) - continue else: - pass # log unexpected keys + pass # TODO: log unexpected keys # Assign the remaining known headers directly inputs_dict["timeout"] = headers.get("timeout", None) inputs_dict["name"] = headers.get("name", None) inputs_dict["header_size"] = headers.get("header_size", None) inputs_dict["total_size"] = headers.get("total_size", None) + inputs_dict["computed_body_hash"] = headers.get("computed_body_hash", None) + inputs_dict["hash_fields"] = json.loads( + base64.b64decode(headers.get("hash_fields", "W10=").encode()).decode( + "utf-8" + ) + ) return inputs_dict diff --git a/tests/unit_tests/test_synapse.py b/tests/unit_tests/test_synapse.py index 27a81412fe..a8b5c57f7f 100644 --- a/tests/unit_tests/test_synapse.py +++ b/tests/unit_tests/test_synapse.py @@ -39,11 +39,16 @@ class Test(bittensor.Synapse): "name": "Test", "header_size": "111", "total_size": "111", + "computed_body_hash": "0xabcdef", + "hash_fields": base64.b64encode( + json.dumps(["key1", "key2"]).encode("utf-8") + ).decode("utf-8"), } + print(headers) # Run the function to test inputs_dict = Test.parse_headers_to_inputs(headers) - + print(inputs_dict) # Check the resulting dictionary assert inputs_dict == { "axon": {"nonce": "111"}, @@ -54,6 +59,8 @@ class Test(bittensor.Synapse): "name": "Test", "header_size": "111", "total_size": "111", + "computed_body_hash": "0xabcdef", + "hash_fields": ["key1", "key2"], } @@ -74,6 +81,10 @@ class Test(bittensor.Synapse): "name": "Test", "header_size": "111", "total_size": "111", + "computed_body_hash": "0xabcdef", + "hash_fields": base64.b64encode( + json.dumps(["key1", "key2"]).encode("utf-8") + ).decode("utf-8"), } # Run the function to test @@ -241,6 +252,18 @@ def test_body_hash_override(): # Create a Synapse instance synapse_instance = bittensor.Synapse() + # Try to set the body_hash property and expect an AttributeError + with pytest.raises( + AttributeError, + match="body_hash property is read-only and cannot be overridden.", + ): + synapse_instance.body_hash = [] + + +def test_required_fields_override(): + # Create a Synapse instance + synapse_instance = bittensor.Synapse() + # Try to set the body_hash property and expect an AttributeError with pytest.raises( AttributeError,