def test_rsa_signature_validator():
receiver_validator = RSASignatureValidator()
sender_validator = RSASignatureValidator(RSAPrivateKey())
mallory_validator = RSASignatureValidator(RSAPrivateKey())
plain_record = DHTRecord(key=b'key', subkey=b'subkey', value=b'value',
expiration_time=get_dht_time() + 10)
protected_records = [
dataclasses.replace(plain_record,
key=plain_record.key + sender_validator.local_public_key),
dataclasses.replace(plain_record,
subkey=plain_record.subkey + sender_validator.local_public_key),
]
# test 1: Non-protected record (no signature added)
assert sender_validator.sign_value(plain_record) == plain_record.value
assert receiver_validator.validate(plain_record)
# test 2: Correct signatures
signed_records = [dataclasses.replace(record, value=sender_validator.sign_value(record))
for record in protected_records]
for record in signed_records:
assert receiver_validator.validate(record)
assert receiver_validator.strip_value(record) == b'value'
# test 3: Invalid signatures
signed_records = protected_records # Without signature
signed_records += [dataclasses.replace(record,
value=record.value + b'[signature:INVALID_BYTES]')
for record in protected_records] # With invalid signature
signed_records += [dataclasses.replace(record, value=mallory_validator.sign_value(record))
for record in protected_records] # With someone else's signature
for record in signed_records:
assert not receiver_validator.validate(record)
async def test_auth_rpc_wrapper():
class Servicer:
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
class Client:
def __init__(self, servicer: Servicer):
self._servicer = servicer
async def rpc_increment(
self, request: dht_pb2.PingRequest) -> dht_pb2.PingResponse:
return await self._servicer.rpc_increment(request)
servicer = AuthRPCWrapper(Servicer(), AuthRole.SERVICER,
MockAuthorizer(RSAPrivateKey(), 'bob'))
client = AuthRPCWrapper(Client(servicer), AuthRole.CLIENT,
MockAuthorizer(RSAPrivateKey(), 'alice'))
request = dht_pb2.PingRequest()
request.peer.endpoint = '127.0.0.1:1111'
response = await client.rpc_increment(request)
assert response.sender_endpoint == '127.0.0.1:2222'
assert response.auth.service_access_token.username == 'bob'
async def test_dhtnode_signatures():
alice = await hivemind.DHTNode.create(record_validator=RSASignatureValidator())
bob = await hivemind.DHTNode.create(
record_validator=RSASignatureValidator(RSAPrivateKey()),
initial_peers=[f"{LOCALHOST}:{alice.port}"])
mallory = await hivemind.DHTNode.create(
record_validator=RSASignatureValidator(RSAPrivateKey()),
initial_peers=[f"{LOCALHOST}:{alice.port}"])
key = b'key'
subkey = b'protected_subkey' + bob.protocol.record_validator.local_public_key
assert await bob.store(key, b'true_value', hivemind.get_dht_time() + 10, subkey=subkey)
assert (await alice.get(key, latest=True)).value[subkey].value == b'true_value'
store_ok = await mallory.store(key, b'fake_value', hivemind.get_dht_time() + 10, subkey=subkey)
assert not store_ok
assert (await alice.get(key, latest=True)).value[subkey].value == b'true_value'
assert await bob.store(key, b'updated_true_value', hivemind.get_dht_time() + 10, subkey=subkey)
assert (await alice.get(key, latest=True)).value[subkey].value == b'updated_true_value'
await bob.shutdown() # Bob has shut down, now Mallory is the single peer of Alice
store_ok = await mallory.store(key, b'updated_fake_value',
hivemind.get_dht_time() + 10, subkey=subkey)
assert not store_ok
assert (await alice.get(key, latest=True)).value[subkey].value == b'updated_true_value'
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)
def test_cached_key():
first_validator = RSASignatureValidator()
second_validator = RSASignatureValidator()
assert first_validator.local_public_key == second_validator.local_public_key
third_validator = RSASignatureValidator(RSAPrivateKey())
assert first_validator.local_public_key != third_validator.local_public_key
def __init__(self, private_key: Optional[RSAPrivateKey] = None):
if private_key is None:
private_key = RSAPrivateKey.process_wide()
self._private_key = private_key
serialized_public_key = private_key.get_public_key().to_bytes()
self._local_public_key = self.PUBLIC_KEY_FORMAT.replace(
b'_key_', serialized_public_key)
def __init__(self, local_private_key: Optional[RSAPrivateKey]=None):
if local_private_key is None:
local_private_key = RSAPrivateKey.process_wide()
self._local_private_key = local_private_key
self._local_public_key = local_private_key.get_public_key()
self._local_access_token = None
self._refresh_lock = asyncio.Lock()
self._recent_nonces = TimedStorage()
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)
async def get_token(self) -> AccessToken:
if MockAuthorizer._authority_private_key is None:
MockAuthorizer._authority_private_key = RSAPrivateKey()
self._authority_public_key = MockAuthorizer._authority_private_key.get_public_key(
)
token = AccessToken(username=self._username,
public_key=self.local_public_key.to_bytes(),
expiration_time=str(datetime.utcnow() +
timedelta(minutes=1)))
token.signature = MockAuthorizer._authority_private_key.sign(
self._token_to_bytes(token))
return token