From b3e9fc2ffd3528cb876dbbf79cbaff339ca9c04f Mon Sep 17 00:00:00 2001 From: ifrit98 Date: Mon, 28 Aug 2023 13:43:05 +0000 Subject: [PATCH 1/5] Add body hash to signature for verification, stop overrides --- bittensor/axon.py | 3 +- bittensor/dendrite.py | 3 +- bittensor/synapse.py | 122 +++++++++++++++++++++++++++++++++++++++++- 3 files changed, 124 insertions(+), 4 deletions(-) diff --git a/bittensor/axon.py b/bittensor/axon.py index 8183723f4c..801d1dbad8 100644 --- a/bittensor/axon.py +++ b/bittensor/axon.py @@ -24,6 +24,7 @@ import copy import time import asyncio +import hashlib import inspect import uvicorn import argparse @@ -583,7 +584,7 @@ def default_verify(self, synapse: bittensor.Synapse) -> Request: keypair = Keypair(ss58_address=synapse.dendrite.hotkey) # Build the signature messages. - message = f"{synapse.dendrite.nonce}.{synapse.dendrite.hotkey}.{self.wallet.hotkey.ss58_address}.{synapse.dendrite.uuid}" + message = f"{synapse.dendrite.nonce}.{synapse.dendrite.hotkey}.{self.wallet.hotkey.ss58_address}.{synapse.dendrite.uuid}{synapse.body_hash}" # Build the unique endpoint key. endpoint_key = f"{synapse.dendrite.hotkey}:{synapse.dendrite.uuid}" diff --git a/bittensor/dendrite.py b/bittensor/dendrite.py index 153e410820..7eab76ebbb 100644 --- a/bittensor/dendrite.py +++ b/bittensor/dendrite.py @@ -22,6 +22,7 @@ import time import torch import httpx +import hashlib import bittensor as bt from typing import Union, Optional, List @@ -314,7 +315,7 @@ def preprocess_synapse_for_request( ) # Sign the request using the dendrite and axon information - message = f"{synapse.dendrite.nonce}.{synapse.dendrite.hotkey}.{synapse.axon.hotkey}.{synapse.dendrite.uuid}" + 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 diff --git a/bittensor/synapse.py b/bittensor/synapse.py index a75324cdd1..620ff4b89d 100644 --- a/bittensor/synapse.py +++ b/bittensor/synapse.py @@ -22,13 +22,14 @@ import pickle import base64 import typing +import hashlib import pydantic from pydantic.schema import schema import bittensor from abc import abstractmethod from fastapi.responses import Response from fastapi import Request -from typing import Dict, Optional, Tuple, Union, List, Callable +from typing import Dict, Optional, Tuple, Union, List, Callable, Any def get_size(obj, seen=None): @@ -204,7 +205,37 @@ class Config: ) -class Synapse(pydantic.BaseModel): +class ProtectOverride(type): + """ + Metaclass to prevent subclasses from overriding specified methods or attributes. + + When a subclass attempts to override a protected attribute or method, a `TypeError` is raised. + The current implementation specifically checks for overriding the 'body_hash' attribute. + + Overriding `protected_method` in a subclass of `MyClass` will raise a TypeError. + """ + + def __new__(cls, name, bases, class_dict): + # Check if the derived class tries to override the 'body_hash' method or attribute. + if ( + any(base for base in bases if hasattr(base, "body_hash")) + and "body_hash" in class_dict + ): + raise TypeError("You can't override the body_hash attribute!") + return super(ProtectOverride, cls).__new__(cls, name, bases, class_dict) + + +class CombinedMeta(ProtectOverride, type(pydantic.BaseModel)): + """ + Metaclass combining functionality of ProtectOverride and BaseModel's metaclass. + + Inherits the attributes and methods from both parent metaclasses to provide combined behavior. + """ + + pass + + +class Synapse(pydantic.BaseModel, metaclass=CombinedMeta): class Config: validate_assignment = True @@ -285,6 +316,16 @@ def set_name_type(cls, values): repr=False, ) + def __setattr__(self, name, value): + """ + Override the __setattr__ method to make the body_hash property read-only. + """ + if name == "body_hash": + raise AttributeError( + "body_hash property is read-only and cannot be overridden." + ) + super().__setattr__(name, value) + def get_total_size(self) -> int: """ Get the total size of the current object. @@ -298,6 +339,83 @@ def get_total_size(self) -> int: self.total_size = get_size(self) return self.total_size + def get_body(self) -> List[Any]: + """ + Retrieve the serialized and encoded non-optional fields of the Synapse instance. + + This method filters through the fields of the Synapse instance and identifies + non-optional attributes that have non-null values, excluding specific attributes + such as `name`, `timeout`, `total_size`, `header_size`, `dendrite`, and `axon`. + It returns a list containing these selected field values. + + Returns: + List[Any]: A list of values from the non-optional fields of the Synapse instance. + + Note: + The determination of whether a field is optional or not is based on the + schema definition for the Synapse class. + """ + fields = [] + + # 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 object is not optional and non-null, add to the list of returned body fields + required = schema([self.__class__])["definitions"][self.name].get( + "required" + ) + if ( + required + and value != None + and field + not in [ + "name", + "timeout", + "total_size", + "header_size", + "dendrite", + "axon", + ] + ): + fields.append(value) + + return fields + + @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. + + Note: + This property is intended to be read-only. Any attempts to override + or set its value will raise an AttributeError due to the protections + set in the __setattr__ method. + + Returns: + str: The hexadecimal representation of the SHA-256 hash of the instance's body. + """ + # Hash the body for verification + body = self.get_body() + + # Convert elements to string and concatenate + concat = "".join(map(str, body)) + + # Create a SHA-256 hash object + sha256 = hashlib.sha256() + + # Update the hash object with the concatenated string + sha256.update(concat.encode("utf-8")) + + # Produce the hash + return sha256.hexdigest() + @property def is_success(self) -> bool: """ From ccc16ca25e6bb0da7c46c57e14c952dcbfab1323 Mon Sep 17 00:00:00 2001 From: ifrit98 Date: Mon, 28 Aug 2023 13:58:14 +0000 Subject: [PATCH 2/5] add tests --- tests/unit_tests/test_synapse.py | 24 +++++++++++++++++++++--- 1 file changed, 21 insertions(+), 3 deletions(-) diff --git a/tests/unit_tests/test_synapse.py b/tests/unit_tests/test_synapse.py index 764736fe49..a44c0d92c6 100644 --- a/tests/unit_tests/test_synapse.py +++ b/tests/unit_tests/test_synapse.py @@ -20,9 +20,6 @@ import typing import pytest import bittensor -import unittest -import unittest.mock as mock -from unittest.mock import MagicMock def test_parse_headers_to_inputs(): @@ -238,3 +235,24 @@ class Test(bittensor.Synapse): assert next_synapse.a["cat"].shape == [10] assert next_synapse.a["dog"].dtype == "torch.float32" assert next_synapse.a["dog"].shape == [11] + + +def test_override_protection(): + with pytest.raises(TypeError, match="You can't override the body_hash attribute!"): + + class DerivedModel(bittensor.Synapse): + @property + def body_hash(self): + return "new_value" + + +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 = "some_value" From 3f181e0a3738ddb3c2ed04a36c013eadf5e358aa Mon Sep 17 00:00:00 2001 From: ifrit98 Date: Mon, 28 Aug 2023 14:00:10 +0000 Subject: [PATCH 3/5] remove unused imports --- bittensor/synapse.py | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/bittensor/synapse.py b/bittensor/synapse.py index 620ff4b89d..f5abbe785d 100644 --- a/bittensor/synapse.py +++ b/bittensor/synapse.py @@ -18,7 +18,6 @@ import ast import sys -import torch import pickle import base64 import typing @@ -26,10 +25,7 @@ import pydantic from pydantic.schema import schema import bittensor -from abc import abstractmethod -from fastapi.responses import Response -from fastapi import Request -from typing import Dict, Optional, Tuple, Union, List, Callable, Any +from typing import Optional, List, Any def get_size(obj, seen=None): From e6d52db177bcfe02b9899ec9f08b191bb2f2e70e Mon Sep 17 00:00:00 2001 From: ifrit98 Date: Mon, 28 Aug 2023 14:02:00 +0000 Subject: [PATCH 4/5] remove hashlib import --- bittensor/axon.py | 1 - bittensor/dendrite.py | 1 - 2 files changed, 2 deletions(-) diff --git a/bittensor/axon.py b/bittensor/axon.py index 801d1dbad8..06c9779c77 100644 --- a/bittensor/axon.py +++ b/bittensor/axon.py @@ -24,7 +24,6 @@ import copy import time import asyncio -import hashlib import inspect import uvicorn import argparse diff --git a/bittensor/dendrite.py b/bittensor/dendrite.py index 7eab76ebbb..82c0b2b1bc 100644 --- a/bittensor/dendrite.py +++ b/bittensor/dendrite.py @@ -22,7 +22,6 @@ import time import torch import httpx -import hashlib import bittensor as bt from typing import Union, Optional, List From a98573f598d9ce1840229a592bfb13303d87d29f Mon Sep 17 00:00:00 2001 From: ifrit98 Date: Mon, 28 Aug 2023 17:57:38 +0000 Subject: [PATCH 5/5] missing . --- bittensor/axon.py | 2 +- bittensor/dendrite.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/bittensor/axon.py b/bittensor/axon.py index 06c9779c77..7a59902df3 100644 --- a/bittensor/axon.py +++ b/bittensor/axon.py @@ -583,7 +583,7 @@ def default_verify(self, synapse: bittensor.Synapse) -> Request: keypair = Keypair(ss58_address=synapse.dendrite.hotkey) # Build the signature messages. - message = f"{synapse.dendrite.nonce}.{synapse.dendrite.hotkey}.{self.wallet.hotkey.ss58_address}.{synapse.dendrite.uuid}{synapse.body_hash}" + message = f"{synapse.dendrite.nonce}.{synapse.dendrite.hotkey}.{self.wallet.hotkey.ss58_address}.{synapse.dendrite.uuid}.{synapse.body_hash}" # Build the unique endpoint key. endpoint_key = f"{synapse.dendrite.hotkey}:{synapse.dendrite.uuid}" diff --git a/bittensor/dendrite.py b/bittensor/dendrite.py index 82c0b2b1bc..5b72f2f6ae 100644 --- a/bittensor/dendrite.py +++ b/bittensor/dendrite.py @@ -314,7 +314,7 @@ def preprocess_synapse_for_request( ) # Sign the request using the dendrite and axon information - message = f"{synapse.dendrite.nonce}.{synapse.dendrite.hotkey}.{synapse.axon.hotkey}.{synapse.dendrite.uuid}{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