Skip to content

Commit

Permalink
simplify hashing for request verification (#1521)
Browse files Browse the repository at this point in the history 
* simplify hashing for request verification

* add override protection for  property

* run black

* update synapse tests

* update default json b64 value for hash_fields

* remove print statement

* run black
  • Loading branch information
@ifrit98
ifrit98 authored Sep 29, 2023
1 parent d34b735 commit 697e611
Show file tree
Hide file tree
Showing 4 changed files with 108 additions and 78 deletions.
47 changes: 20 additions & 27 deletions bittensor/axon.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,7 @@
import copy
import json
import time
import base64
import asyncio
import inspect
import uvicorn
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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}"
Expand Down
33 changes: 2 additions & 31 deletions bittensor/dendrite.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
81 changes: 62 additions & 19 deletions bittensor/synapse.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand All @@ -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
Expand Down Expand Up @@ -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."
Expand All @@ -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:
"""
Expand Down Expand Up @@ -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

Expand Down
25 changes: 24 additions & 1 deletion tests/unit_tests/test_synapse.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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"},
Expand All @@ -54,6 +59,8 @@ class Test(bittensor.Synapse):
"name": "Test",
"header_size": "111",
"total_size": "111",
"computed_body_hash": "0xabcdef",
"hash_fields": ["key1", "key2"],
}


Expand All @@ -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
Expand Down Expand Up @@ -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,
Expand Down

0 comments on commit 697e611

Please sign in to comment.