Skip to content

Commit

Permalink
Blacklist fixes + depreciation of old signatures (#1240)
Browse files Browse the repository at this point in the history
* fixes blacklist error message + remove

* remove checks for parse signature

* remove sign v1 tests

* fix for the syanpse checks

* fix tests and remove duplicate sign
  • Loading branch information
Eugene-hu authored Apr 12, 2023
1 parent 1cb9bed commit 4741f62
Show file tree
Hide file tree
Showing 7 changed files with 31 additions and 116 deletions.
43 changes: 11 additions & 32 deletions bittensor/_axon/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -374,22 +374,6 @@ def __init__(
self.blacklist = blacklist
self.receiver_hotkey = receiver_hotkey

def parse_legacy_signature(
self, signature: str
) -> Union[Tuple[int, str, str, str, int], None]:
r"""Attempts to parse a signature using the legacy format, using `bitxx` as a separator"""
parts = signature.split("bitxx")
if len(parts) < 4:
return None
try:
nonce = int(parts[0])
parts = parts[1:]
except ValueError:
return None
receptor_uuid, parts = parts[-1], parts[:-1]
signature, parts = parts[-1], parts[:-1]
sender_hotkey = "".join(parts)
return (nonce, sender_hotkey, signature, receptor_uuid, 1)

def parse_signature_v2(
self, signature: str
Expand All @@ -405,7 +389,7 @@ def parse_signature_v2(
sender_hotkey = parts[1]
signature = parts[2]
receptor_uuid = parts[3]
return (nonce, sender_hotkey, signature, receptor_uuid, 2)
return (nonce, sender_hotkey, signature, receptor_uuid)

def parse_signature(
self, metadata: Dict[str, str]
Expand All @@ -418,10 +402,9 @@ def parse_signature(
if int(version) < 370:
raise Exception("Incorrect Version")

for parser in [self.parse_signature_v2, self.parse_legacy_signature]:
parts = parser(signature)
if parts is not None:
return parts
parts = self.parse_signature_v2(signature)
if parts is not None:
return parts
raise Exception("Unknown signature format")

def check_signature(
Expand All @@ -430,17 +413,12 @@ def check_signature(
sender_hotkey: str,
signature: str,
receptor_uuid: str,
format: int,
):
r"""verification of signature in metadata. Uses the pubkey and nonce"""
keypair = Keypair(ss58_address=sender_hotkey)
# Build the expected message which was used to build the signature.
if format == 2:
message = f"{nonce}.{sender_hotkey}.{self.receiver_hotkey}.{receptor_uuid}"
elif format == 1:
message = f"{nonce}{sender_hotkey}{receptor_uuid}"
else:
raise Exception("Invalid signature version")
message = f"{nonce}.{sender_hotkey}.{self.receiver_hotkey}.{receptor_uuid}"

# Build the key which uniquely identifies the endpoint that has signed
# the message.
endpoint_key = f"{sender_hotkey}:{receptor_uuid}"
Expand All @@ -467,8 +445,10 @@ def black_list_checking(self, hotkey: str, method: str):
if request_type is None:
raise Exception("Unknown request type")

if self.blacklist(hotkey, request_type):
raise Exception("Request type is blacklisted")
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"""
Expand All @@ -481,12 +461,11 @@ def intercept_service(self, continuation, handler_call_details):
sender_hotkey,
signature,
receptor_uuid,
signature_format,
) = self.parse_signature(metadata)

# signature checking
self.check_signature(
nonce, sender_hotkey, signature, receptor_uuid, signature_format
nonce, sender_hotkey, signature, receptor_uuid
)

# blacklist checking
Expand Down
40 changes: 1 addition & 39 deletions bittensor/_metagraph/naka_metagraph_impl.py
Original file line number Diff line number Diff line change
Expand Up @@ -381,45 +381,7 @@ def retrieve_cached_neurons( self, block: int = None ):
"""
Retrieves cached metagraph syncs from IPFS.
"""
ipfs = bittensor.Ipfs()
ipns_hash = ipfs.latest_neurons_ipns
ipfs_hash = ipfs.cat

if block != None:
ipns_hash = ipfs.historical_neurons_ipns
ipfs_hash = ipfs.node_get

try:
# Ping IPNS for latest IPFS hash
ipns_resolve = ipfs.retrieve_directory(ipfs.ipns_resolve, (('arg', ipns_hash),))

# Extract IPFS hash from IPNS response
ipfs_path = ast.literal_eval(ipns_resolve.text)
except Exception as e:
logger.error("Error detected in metagraph sync: {} with sample text {}".format(e,ipns_resolve.text))

# Try Again
# Ping IPNS for latest IPFS hash
ipns_resolve = ipfs.retrieve_directory(ipfs.ipns_resolve, (('arg', ipns_hash),))

# Extract IPFS hash from IPNS response
ipfs_path = ast.literal_eval(ipns_resolve.text)

ipfs_resolved_hash = ipfs_path['Path'].split("ipfs/")[1]
ipfs_response = ipfs.retrieve_directory(ipfs_hash, (('arg', ipfs_resolved_hash),))

# Extract all neuron sync hashes
if block != None:
historical_neurons = json.loads(ipfs_response.content)['Links']
# Find the one that corresponds to our block
sync_data = next(item for item in historical_neurons if item["Name"] == "nakamoto-{}.pkl".format(block))
# Retrieve Neuron contents
ipfs_response = ipfs.retrieve_directory(ipfs.cat, (('arg', sync_data['Hash']),))

# Unpickle the response
neurons = pickle.loads(ipfs_response.content)

return neurons
raise Exception('Nakamoto is deprecated. Please use finney instead')

def sync ( self, block: int = None, cached: bool = True, subtensor = None, netuid: int = 1 ) -> 'Metagraph':
r""" Synchronizes this metagraph with the chain state.
Expand Down
14 changes: 10 additions & 4 deletions bittensor/_neuron/text/core_server/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -425,7 +425,13 @@ def synapse_check(self, synapse, hotkey, inputs_x=None):
"""
## Uid that sent the request
incoming_uid = self.metagraph.hotkeys.index(hotkey)
try:
incoming_uid = self.metagraph.hotkeys.index(hotkey)
except Exception as e:
if self.config.neuron.blacklist_allow_non_registered:
return False
return True

batch_size, sequence_len = inputs_x[0].size()
if synapse.synapse_type == bittensor.proto.Synapse.SynapseType.TEXT_LAST_HIDDEN_STATE:
if self.metagraph.S[incoming_uid] < self.config.neuron.lasthidden_stake \
Expand Down Expand Up @@ -643,10 +649,10 @@ def hotkey_check():
time_check()
stake_check()
hotkey_check()
return False
except Exception as e:
return False, None
except Exception as error:
self.prometheus_counters.labels("blacklisted").inc()
return True
return True, error

def get_neuron(self):
if self.subtensor.network == 'nakamoto':
Expand Down
15 changes: 1 addition & 14 deletions bittensor/_receptor/receptor_impl.py
Original file line number Diff line number Diff line change
Expand Up @@ -123,27 +123,14 @@ def __del__ ( self ):
def __exit__ ( self ):
self.__del__()

def sign_v1( self ):
r""" Uses the wallet pubkey to sign a message containing the pubkey and the time
"""
nonce = self.nonce()
message = str(nonce) + str(self.wallet.hotkey.ss58_address) + str(self.receptor_uid)
spliter = 'bitxx'
signature = spliter.join([ str(nonce), str(self.wallet.hotkey.ss58_address), "0x" + self.wallet.hotkey.sign(message).hex(), str(self.receptor_uid) ])
return signature

def sign_v2(self):
def sign(self):
nonce = f"{self.nonce()}"
sender_hotkey = self.wallet.hotkey.ss58_address
receiver_hotkey = self.endpoint.hotkey
message = f"{nonce}.{sender_hotkey}.{receiver_hotkey}.{self.receptor_uid}"
signature = f"0x{self.wallet.hotkey.sign(message).hex()}"
return ".".join([nonce, sender_hotkey, signature, self.receptor_uid])

def sign(self):
if self.endpoint.version >= bittensor.__new_signature_version__:
return self.sign_v2()
return self.sign_v1()

def nonce ( self ):
r"""creates a string representation of the time
Expand Down
19 changes: 4 additions & 15 deletions tests/unit_tests/bittensor_tests/test_axon.py
Original file line number Diff line number Diff line change
Expand Up @@ -39,12 +39,6 @@
def gen_nonce():
return f"{time.monotonic_ns()}"

def sign_v1(wallet):
nonce, receptor_uid = gen_nonce(), str(uuid.uuid1())
message = "{}{}{}".format(nonce, str(wallet.hotkey.ss58_address), receptor_uid)
spliter = 'bitxx'
signature = spliter.join([ nonce, str(wallet.hotkey.ss58_address), "0x" + wallet.hotkey.sign(message).hex(), receptor_uid])
return signature

def sign_v2(sender_wallet, receiver_wallet):
nonce, receptor_uid = gen_nonce(), str(uuid.uuid1())
Expand All @@ -55,13 +49,8 @@ def sign_v2(sender_wallet, receiver_wallet):
return ".".join([nonce, sender_hotkey, signature, receptor_uid])

def sign(sender_wallet, receiver_wallet, receiver_version):
if receiver_version >= bittensor.__new_signature_version__:
return sign_v2(sender_wallet, receiver_wallet)
return sign_v1(sender_wallet)

def test_sign_v1():
sign_v1(wallet)
sign_v1(axon.wallet)

return sign_v2(sender_wallet, receiver_wallet)

def test_sign_v2():
sign_v2(sender_wallet, wallet)
Expand Down Expand Up @@ -951,7 +940,7 @@ def forward( inputs_x:torch.FloatTensor, synapse , model_output = None):
axon.stop()

def test_grpc_forward_works():
for receiver_version in [341, bittensor.__new_signature_version__, bittensor.__version_as_int__]:
for receiver_version in [bittensor.__new_signature_version__, bittensor.__version_as_int__]:
run_test_grpc_forward_works(receiver_version)

def run_test_grpc_backward_works(receiver_version):
Expand Down Expand Up @@ -994,7 +983,7 @@ def backward( inputs_x:torch.FloatTensor, grads_dy:torch.FloatTensor, synapses):
axon.stop()

def test_grpc_backward_works():
for receiver_version in [341, bittensor.__new_signature_version__, bittensor.__version_as_int__]:
for receiver_version in [bittensor.__new_signature_version__, bittensor.__version_as_int__]:
run_test_grpc_backward_works(receiver_version)

def test_grpc_forward_fails():
Expand Down
8 changes: 4 additions & 4 deletions tests/unit_tests/bittensor_tests/test_neuron.py
Original file line number Diff line number Diff line change
Expand Up @@ -346,7 +346,6 @@ def test_stake_blacklist(self):
mock_hotkey,
mock_hotkey_1,
],
total_stake= mock_total_stake,
S=torch.tensor(mock_total_stake),
)

Expand Down Expand Up @@ -382,12 +381,13 @@ def test_stake_blacklist(self):
# args, kwargs
_, kwargs = mock_new_axon.call_args
blacklist = kwargs['blacklist']

# Check that the blacklist rejects below min stake
assert blacklist(mock_hotkey, bittensor.proto.RequestType.FORWARD) == True
check, error = blacklist(mock_hotkey, bittensor.proto.RequestType.FORWARD)
assert check == True

# Check that the blacklist accepts above min stake
assert blacklist(mock_hotkey_1, bittensor.proto.RequestType.FORWARD) == False
check, error = blacklist(mock_hotkey_1, bittensor.proto.RequestType.FORWARD)
assert check == False


if __name__ == '__main__':
Expand Down
8 changes: 0 additions & 8 deletions tests/unit_tests/bittensor_tests/test_receptor.py
Original file line number Diff line number Diff line change
Expand Up @@ -433,13 +433,6 @@ def backward_break():
assert ops == [bittensor.proto.ReturnCode.UnknownException] * len(synapses)

def test_receptor_signature_output():
def verify_v1(signature: str):
(nonce, sender_address, signature, receptor_uuid) = signature.split("bitxx")
assert nonce == "123"
assert sender_address == "5Ey8t8pBJSYqLYCzeC3HiPJu5DxzXy2Dzheaj29wRHvhjoai"
assert receptor_uuid == "6d8b8788-6b6a-11ed-916f-0242c0a85003"
message = f"{nonce}{sender_address}{receptor_uuid}"
assert wallet.hotkey.verify(message, signature)

def verify_v2(signature: str):
(nonce, sender_address, signature, receptor_uuid) = signature.split(".")
Expand All @@ -450,7 +443,6 @@ def verify_v2(signature: str):
assert wallet.hotkey.verify(message, signature)

matrix = {
bittensor.__new_signature_version__ - 1: verify_v1,
bittensor.__new_signature_version__: verify_v2,
}

Expand Down

0 comments on commit 4741f62

Please sign in to comment.