From aeeea3abf4cd697e6400bbe8dd61c5eb2b3131f9 Mon Sep 17 00:00:00 2001
From: Ben Bangert <ben@groovie.org>
Date: Sat, 16 Jul 2016 21:03:50 -0700
Subject: [PATCH] fix: ensure router_type is present in all records

It was possible previously for a corrupted record to be missing a router_type
which caused AWS ValidationError's. This change ensures all records have
a router_type and connected_at which are required values for the register
user expression condition that AWS checks.

Closes #526
---
 autopush/tests/test_websocket.py | 38 ++++++++++++++++++++++++++++++++
 autopush/websocket.py            |  6 +++++
 2 files changed, 44 insertions(+)

diff --git a/autopush/tests/test_websocket.py b/autopush/tests/test_websocket.py
index 505408f4..2230b933 100644
--- a/autopush/tests/test_websocket.py
+++ b/autopush/tests/test_websocket.py
@@ -484,6 +484,7 @@ def test_hello_tomorrow(self):
             uaid=orig_uaid,
             connected_at=ms_time(),
             current_month="message_2016_3",
+            router_type="simplepush",
         ))
 
         # router.register_user returns (registered, previous
@@ -565,6 +566,39 @@ def check_result(msg):
         eq_(self.proto.base_tags, ['use_webpush:True'])
         return self._check_response(check_result)
 
+    def test_hello_with_missing_router_type(self):
+        self._connect()
+        uaid = uuid.uuid4().hex
+        router = self.proto.ap_settings.router
+        router.register_user(dict(
+            uaid=uaid,
+            connected_at=ms_time(),
+        ))
+        self._send_message(dict(messageType="hello", channelIDs=[],
+                                uaid=uaid))
+
+        def check_result(msg):
+            eq_(msg["status"], 200)
+            ok_(msg["uaid"] != uaid)
+        return self._check_response(check_result)
+
+    def test_hello_with_missing_current_month(self):
+        self._connect()
+        uaid = uuid.uuid4().hex
+        router = self.proto.ap_settings.router
+        router.register_user(dict(
+            uaid=uaid,
+            connected_at=ms_time(),
+            router_type="webpush",
+        ))
+        self._send_message(dict(messageType="hello", channelIDs=[],
+                                uaid=uaid, use_webpush=True))
+
+        def check_result(msg):
+            eq_(msg["status"], 200)
+            ok_(msg["uaid"] != uaid)
+        return self._check_response(check_result)
+
     def test_hello_with_uaid(self):
         self._connect()
         uaid = uuid.uuid4().hex
@@ -572,6 +606,7 @@ def test_hello_with_uaid(self):
         router.register_user(dict(
             uaid=uaid,
             connected_at=ms_time(),
+            router_type="simplepush",
         ))
         self._send_message(dict(messageType="hello", channelIDs=[],
                                 uaid=uaid))
@@ -1657,6 +1692,7 @@ def test_notification_results(self):
         router.register_user(dict(
             uaid=uaid,
             connected_at=ms_time(),
+            router_type="simplepush",
         ))
 
         storage = self.proto.ap_settings.storage
@@ -1711,6 +1747,7 @@ def test_notification_dont_deliver_after_ack(self):
         router.register_user(dict(
             uaid=uaid,
             connected_at=ms_time(),
+            router_type="simplepush",
         ))
 
         storage = self.proto.ap_settings.storage
@@ -1772,6 +1809,7 @@ def test_notification_dont_deliver(self):
         router.register_user(dict(
             uaid=uaid,
             connected_at=ms_time(),
+            router_type="simplepush",
         ))
 
         storage = self.proto.ap_settings.storage
diff --git a/autopush/websocket.py b/autopush/websocket.py
index 781362c7..adfec15c 100644
--- a/autopush/websocket.py
+++ b/autopush/websocket.py
@@ -686,6 +686,12 @@ def _verify_user_record(self):
         except ItemNotFound:
             return None
 
+        # All records must have a router_type and connected_at, in some odd
+        # cases a record exists for some users that doesn't
+        if "router_type" not in record or "connected_at" not in record:
+            self.force_retry(self.ap_settings.router.drop_user, self.ps.uaid)
+            return None
+
         # Validate webpush records
         if self.ps.use_webpush:
             # Current month must exist and be a valid prior month