fix handling of invalid token

This commit is contained in:
Victor Shyba 2019-05-12 04:12:12 -03:00
parent f5cef14d8d
commit 32d1563885
2 changed files with 14 additions and 1 deletions

View file

@ -628,12 +628,15 @@ class KademliaProtocol(DatagramProtocol):
return True return True
async def store_to_peer(self, hash_value: bytes, peer: 'KademliaPeer') -> typing.Tuple[bytes, bool]: async def store_to_peer(self, hash_value: bytes, peer: 'KademliaPeer') -> typing.Tuple[bytes, bool]:
try: async def __store():
res = await self.get_rpc_peer(peer).store(hash_value) res = await self.get_rpc_peer(peer).store(hash_value)
if res != b"OK": if res != b"OK":
raise ValueError(res) raise ValueError(res)
log.debug("Stored %s to %s", binascii.hexlify(hash_value).decode()[:8], peer) log.debug("Stored %s to %s", binascii.hexlify(hash_value).decode()[:8], peer)
return peer.node_id, True return peer.node_id, True
try:
return await __store()
except asyncio.TimeoutError: except asyncio.TimeoutError:
log.debug("Timeout while storing blob_hash %s at %s", binascii.hexlify(hash_value).decode()[:8], peer) log.debug("Timeout while storing blob_hash %s at %s", binascii.hexlify(hash_value).decode()[:8], peer)
except ValueError as err: except ValueError as err:
@ -641,6 +644,10 @@ class KademliaProtocol(DatagramProtocol):
except Exception as err: except Exception as err:
if 'Invalid token' in str(err): if 'Invalid token' in str(err):
self.peer_manager.clear_token(peer.node_id) self.peer_manager.clear_token(peer.node_id)
try:
return await __store()
except:
return peer.node_id, False
else: else:
log.exception("Unexpected error while storing blob_hash") log.exception("Unexpected error while storing blob_hash")
return peer.node_id, False return peer.node_id, False

View file

@ -86,3 +86,9 @@ class DHTIntegrationTest(AsyncioTestCase):
blob_hash = hexlify(constants.generate_id(1337)).decode() blob_hash = hexlify(constants.generate_id(1337)).decode()
node_ids = await node1.announce_blob(blob_hash) node_ids = await node1.announce_blob(blob_hash)
self.assertIn(node2.protocol.node_id, node_ids) self.assertIn(node2.protocol.node_id, node_ids)
node2.protocol.node_rpc.refresh_token()
node_ids = await node1.announce_blob(blob_hash)
self.assertIn(node2.protocol.node_id, node_ids)
node2.protocol.node_rpc.refresh_token()
node_ids = await node1.announce_blob(blob_hash)
self.assertIn(node2.protocol.node_id, node_ids)