async def test_call_unary_handler(should_cancel,
replicate,
handle_name="handle"):
handler_cancelled = False
async def ping_handler(request, context):
try:
await asyncio.sleep(2)
except asyncio.CancelledError:
nonlocal handler_cancelled
handler_cancelled = True
return dht_pb2.PingResponse(peer=dht_pb2.NodeInfo(
node_id=context.id.encode(), rpc_port=context.port),
sender_endpoint=context.handle_name,
available=True)
server_primary = await P2P.create()
server = await replicate_if_needed(server_primary, replicate)
server_pid = server_primary._child.pid
await server.add_unary_handler(handle_name, ping_handler,
dht_pb2.PingRequest, dht_pb2.PingResponse)
assert is_process_running(server_pid)
nodes = bootstrap_from([server])
client_primary = await P2P.create(bootstrap=True, bootstrap_peers=nodes)
client = await replicate_if_needed(client_primary, replicate)
client_pid = client_primary._child.pid
assert is_process_running(client_pid)
ping_request = dht_pb2.PingRequest(peer=dht_pb2.NodeInfo(
node_id=client.id.encode(), rpc_port=client._host_port),
validate=True)
expected_response = dht_pb2.PingResponse(peer=dht_pb2.NodeInfo(
node_id=server.id.encode(), rpc_port=server._host_port),
sender_endpoint=handle_name,
available=True)
await client.wait_for_at_least_n_peers(1)
libp2p_server_id = PeerID.from_base58(server.id)
stream_info, reader, writer = await client._client.stream_open(
libp2p_server_id, (handle_name, ))
await P2P.send_protobuf(ping_request, dht_pb2.PingRequest, writer)
if should_cancel:
writer.close()
await asyncio.sleep(1)
assert handler_cancelled
else:
result, err = await P2P.receive_protobuf(dht_pb2.PingResponse, reader)
assert err is None
assert result == expected_response
assert not handler_cancelled
await server.stop_listening()
await server_primary.shutdown()
assert not is_process_running(server_pid)
await client_primary.shutdown()
assert not is_process_running(client_pid)
async def rpc_ping(self, request: dht_pb2.PingRequest,
context: grpc.ServicerContext):
""" Some node wants us to add it to our routing table. """
response = dht_pb2.PingResponse(peer=self.node_info,
sender_endpoint=context.peer(),
dht_time=get_dht_time(),
available=False)
if request.peer and request.peer.node_id and request.peer.rpc_port:
sender_id = DHTID.from_bytes(request.peer.node_id)
if request.peer.endpoint != dht_pb2.NodeInfo.endpoint.DESCRIPTOR.default_value:
sender_endpoint = request.peer.endpoint # if peer has preferred endpoint, use it
else:
sender_endpoint = replace_port(context.peer(),
new_port=request.peer.rpc_port)
response.sender_endpoint = sender_endpoint
if request.validate:
response.available = await self.call_ping(
response.sender_endpoint, validate=False) == sender_id
asyncio.create_task(
self.update_routing_table(sender_id,
sender_endpoint,
responded=response.available
or not request.validate))
return response
async def rpc_increment(
self, request: dht_pb2.PingRequest) -> dht_pb2.PingResponse:
assert request.peer.endpoint == '127.0.0.1:1111'
assert request.auth.client_access_token.username == 'alice'
response = dht_pb2.PingResponse()
response.sender_endpoint = '127.0.0.1:2222'
return response
async def ping_handler(request, context):
try:
await asyncio.sleep(2)
except asyncio.CancelledError:
nonlocal handler_cancelled
handler_cancelled = True
return dht_pb2.PingResponse(peer=dht_pb2.NodeInfo(
node_id=context.id.encode(), rpc_port=context.port),
sender_endpoint=context.handle_name,
available=True)
async def test_valid_request_and_response():
client_authorizer = MockAuthorizer(RSAPrivateKey())
service_authorizer = MockAuthorizer(RSAPrivateKey())
request = dht_pb2.PingRequest()
request.peer.endpoint = '127.0.0.1:7777'
await client_authorizer.sign_request(request,
service_authorizer.local_public_key)
assert await service_authorizer.validate_request(request)
response = dht_pb2.PingResponse()
response.sender_endpoint = '127.0.0.1:31337'
await service_authorizer.sign_response(response, request)
assert await client_authorizer.validate_response(response, request)
async def test_invalid_access_token():
client_authorizer = MockAuthorizer(RSAPrivateKey())
service_authorizer = MockAuthorizer(RSAPrivateKey())
request = dht_pb2.PingRequest()
request.peer.endpoint = '127.0.0.1:7777'
await client_authorizer.sign_request(request,
service_authorizer.local_public_key)
# Break the access token signature
request.auth.client_access_token.signature = b'broken'
assert not await service_authorizer.validate_request(request)
response = dht_pb2.PingResponse()
response.sender_endpoint = '127.0.0.1:31337'
await service_authorizer.sign_response(response, request)
# Break the access token signature
response.auth.service_access_token.signature = b'broken'
assert not await client_authorizer.validate_response(response, request)
async def test_invalid_signatures():
client_authorizer = MockAuthorizer(RSAPrivateKey())
service_authorizer = MockAuthorizer(RSAPrivateKey())
request = dht_pb2.PingRequest()
request.peer.endpoint = '127.0.0.1:7777'
await client_authorizer.sign_request(request,
service_authorizer.local_public_key)
# A man-in-the-middle attacker changes the request content
request.peer.endpoint = '127.0.0.2:7777'
assert not await service_authorizer.validate_request(request)
response = dht_pb2.PingResponse()
response.sender_endpoint = '127.0.0.1:31337'
await service_authorizer.sign_response(response, request)
# A man-in-the-middle attacker changes the response content
response.sender_endpoint = '127.0.0.2:31337'
assert not await client_authorizer.validate_response(response, request)