Skip to content

Commit

Permalink
push_notifications: Upgrade bouncer to handle encryption.
Browse files Browse the repository at this point in the history
  • Loading branch information
@hashirsarwar
hashirsarwar committed Jul 3, 2020
1 parent 44958dc commit bc499a6
Show file tree
Hide file tree
Showing 3 changed files with 114 additions and 13 deletions.
37 changes: 30 additions & 7 deletions zerver/lib/push_notifications.py
View file
Expand Up @@ -353,15 +353,20 @@ def send_android_push_notification(devices: List[DeviceToken], data: Dict[str, A
def uses_notification_bouncer() -> bool:
return settings.PUSH_NOTIFICATION_BOUNCER_URL is not None

EncryptedPayload = List[Tuple[str, Dict[str, Any]]]
BouncerPayload = Union[Dict[str, Any], EncryptedPayload]

def send_notifications_to_bouncer(user_profile_id: int,
apns_payload: Dict[str, Any],
gcm_payload: Dict[str, Any],
gcm_options: Dict[str, Any]) -> None:
apns_payload: BouncerPayload,
gcm_payload: BouncerPayload,
gcm_options: Dict[str, Any],
encrypted: bool=False) -> None:
post_data = {
'user_id': user_profile_id,
'apns_payload': apns_payload,
'gcm_payload': gcm_payload,
'gcm_options': gcm_options,
'encrypted': encrypted,
}
# Calls zilencer.views.remote_server_notify_push
send_json_to_push_bouncer('POST', 'push/notify', post_data)
Expand Down Expand Up @@ -418,6 +423,7 @@ def add_push_device_token(user_profile: UserProfile,
'user_id': user_profile.id,
'token': token_str,
'token_kind': kind,
'encrypted': encrypt_notifications,
}

if kind == PushDeviceToken.APNS:
Expand Down Expand Up @@ -801,6 +807,13 @@ def separate_devices(devices: DeviceList) -> DevicesByKind:

return encrypted, non_encrypted

def prepare_encrypted_payload(data: EncryptedData) -> EncryptedPayload:
payload = []
for device, content in data:
payload.append((device.token, content))

return payload

@statsd_increment("push_notifications")
def handle_push_notification(user_profile_id: int, missed_message: Dict[str, Any]) -> None:
"""
Expand Down Expand Up @@ -877,10 +890,20 @@ def handle_push_notification(user_profile_id: int, missed_message: Dict[str, Any
encrypt_apns_payload)

if uses_notification_bouncer():
send_notifications_to_bouncer(user_profile_id,
apns_payload,
gcm_payload,
gcm_options)
if android_devices or apple_devices:
send_notifications_to_bouncer(user_profile_id,
apns_payload,
gcm_payload,
gcm_options)

if encrypted_android_devices or encrypted_apple_devices:
encrypted_apns_payload = prepare_encrypted_payload(encrypted_apns_data)
encrypted_gcm_payload = prepare_encrypted_payload(encrypted_gcm_data)
send_notifications_to_bouncer(user_profile_id,
encrypted_apns_payload,
encrypted_gcm_payload,
gcm_options,
encrypted=True)
return

if apple_devices:
Expand Down
59 changes: 53 additions & 6 deletions zerver/tests/test_push_notifications.py
View file
Expand Up @@ -321,6 +321,24 @@ def test_push_bouncer_api(self, mock_request: Any) -> None:
subdomain="zulip")
self.assert_json_error(result, 'Token does not exist')

# Test when token exists in bouncer but not in local server
local_server_token_count = PushDeviceToken.objects.all().count()
RemotePushDeviceToken.objects.create(
kind=kind,
token=token,
user_id=user.id,
server=RemoteZulipServer.objects.get(uuid=self.server_uuid),
)
self.assertEqual(RemotePushDeviceToken.objects.all().count(), 1)

result = self.client_delete(endpoint, {'token': token,
'token_kind': kind,
'user_id': user.id},
subdomain="zulip")
self.assert_json_success(result)
self.assertEqual(PushDeviceToken.objects.all().count(), local_server_token_count)
self.assertEqual(RemotePushDeviceToken.objects.all().count(), 0)

with mock.patch('zerver.lib.remote_server.requests.request',
side_effect=requests.ConnectionError):
result = self.client_post(endpoint, {'token': token},
Expand Down Expand Up @@ -865,23 +883,51 @@ def test_send_notifications_to_bouncer(self) -> None:
message=message,
)

token = hex_to_b64(u'aaaa')
PushDeviceToken.objects.create(
kind=PushDeviceToken.GCM,
token=token,
user=user_profile,
)

token = hex_to_b64(u'eeee')
PushDeviceToken.objects.create(
kind=PushDeviceToken.GCM,
token=token,
user=user_profile,
notification_encryption_key=generate_encryption_key()
)

missed_message = {
'message_id': message.id,
'trigger': 'private_message',
}
with self.settings(PUSH_NOTIFICATION_BOUNCER_URL=True), \
with self.settings(PUSH_NOTIFICATION_BOUNCER_URL=True,
PUSH_NOTIFICATION_ENCRYPTION=True), \
mock.patch('zerver.lib.push_notifications.get_message_payload_apns',
return_value={'apns': True}), \
mock.patch('zerver.lib.push_notifications.get_message_payload_gcm',
return_value=({'gcm': True}, {})), \
mock.patch('zerver.lib.push_notifications.encrypt_apns_payload',
return_value={'apns_encrypted': True}), \
mock.patch('zerver.lib.push_notifications.encrypt_gcm_payload',
return_value={'gcm_encrypted': True}), \
mock.patch('zerver.lib.push_notifications'
'.send_notifications_to_bouncer') as mock_send:
handle_push_notification(user_profile.id, missed_message)
mock_send.assert_called_with(user_profile.id,
{'apns': True},
{'gcm': True},
{},
)
mock_send.assert_any_call(user_profile.id,
{'apns': True},
{'gcm': True},
{},
)

mock_send.assert_any_call(user_profile.id,
[],
[(token, {'gcm_encrypted': True})],
{},
encrypted=True,
)
self.assertEqual(mock_send.call_count, 2)

def test_non_bouncer_push(self) -> None:
self.setup_apns_tokens()
Expand Down Expand Up @@ -1652,6 +1698,7 @@ def test_send_notifications_to_bouncer(self, mock_send: mock.MagicMock) -> None:
'apns_payload': {'apns': True},
'gcm_payload': {'gcm': True},
'gcm_options': {},
'encrypted': False,
}
mock_send.assert_called_with('POST',
'push/notify',
Expand Down
31 changes: 31 additions & 0 deletions zilencer/views.py
View file
Expand Up @@ -15,7 +15,11 @@
from zerver.decorator import InvalidZulipServerKeyError, require_post
from zerver.lib.exceptions import JsonableError
from zerver.lib.push_notifications import (
EncryptedData,
EncryptedPayload,
send_android_encrypted_push_notification,
send_android_push_notification,
send_apple_encrypted_push_notification,
send_apple_push_notification,
)
from zerver.lib.request import REQ, has_request_variables
Expand Down Expand Up @@ -55,6 +59,16 @@ def validate_bouncer_token_request(entity: Union[UserProfile, RemoteZulipServer]
validate_token(token, kind)
return server

def prepare_encrypted_data(devices: List[RemotePushDeviceToken],
payload: EncryptedPayload) -> EncryptedData:
device_map = {d.token: d for d in devices}
encrypted_data = []
for token, data in payload:
if token in device_map:
encrypted_data.append((device_map[token], data))

return encrypted_data

@csrf_exempt
@require_post
@has_request_variables
Expand Down Expand Up @@ -102,6 +116,7 @@ def register_remote_server(
def register_remote_push_device(request: HttpRequest, entity: Union[UserProfile, RemoteZulipServer],
user_id: int=REQ(validator=check_int), token: str=REQ(),
token_kind: int=REQ(validator=check_int),
encrypted: bool=REQ(default=False, validator=check_bool),
ios_app_id: Optional[str]=None) -> HttpResponse:
server = validate_bouncer_token_request(entity, token, token_kind)

Expand All @@ -112,6 +127,7 @@ def register_remote_push_device(request: HttpRequest, entity: Union[UserProfile,
server=server,
kind=token_kind,
token=token,
encrypted=encrypted,
ios_app_id=ios_app_id,
# last_updated is to be renamed to date_created.
last_updated=timezone.now())
Expand Down Expand Up @@ -153,6 +169,7 @@ def remote_server_notify_push(request: HttpRequest, entity: Union[UserProfile, R
gcm_payload = payload['gcm_payload']
apns_payload = payload['apns_payload']
gcm_options = payload.get('gcm_options', {})
encrypted = payload['encrypted']

android_devices = list(RemotePushDeviceToken.objects.filter(
user_id=user_id,
Expand All @@ -166,10 +183,24 @@ def remote_server_notify_push(request: HttpRequest, entity: Union[UserProfile, R
server=server,
))

if encrypted:
android_data = prepare_encrypted_data(android_devices, gcm_payload)
apple_data = prepare_encrypted_data(apple_devices, apns_payload)

if android_data:
send_android_encrypted_push_notification(android_data, gcm_options)

if apple_data:
send_apple_encrypted_push_notification(user_id, apple_data)

return json_success()

if android_devices:
android_devices = [d for d in android_devices if not d.encrypted]
send_android_push_notification(android_devices, gcm_payload, gcm_options, remote=True)

if apple_devices:
apple_devices = [d for d in apple_devices if not d.encrypted]
send_apple_push_notification(user_id, apple_devices, apns_payload, remote=True)

return json_success()
Expand Down

0 comments on commit bc499a6

Please sign in to comment.