From b997d6b0cc490f863bc03612ccfa688526125326 Mon Sep 17 00:00:00 2001 From: jrconlin Date: Tue, 15 Mar 2016 11:45:03 -0700 Subject: [PATCH] bug: limit valid months to acceptable range A user that tries to connect from a period longer than we currently allow for could cause a "KeyError" on the server. Instead, we should require that the user use a new UAID, which shoud cause the client to re-register older connections. Closes #350 --- autopush/db.py | 19 ++++- autopush/settings.py | 35 ++++++-- autopush/tests/test_integration.py | 2 +- autopush/tests/test_websocket.py | 130 ++++++++++++++++++++++++++++- autopush/websocket.py | 9 +- 5 files changed, 179 insertions(+), 16 deletions(-) diff --git a/autopush/db.py b/autopush/db.py index 6c45f650..a229ce0a 100644 --- a/autopush/db.py +++ b/autopush/db.py @@ -64,10 +64,11 @@ def normalize_id(id): return '-'.join((raw[:8], raw[8:12], raw[12:16], raw[16:20], raw[20:])) -def make_rotating_tablename(prefix, delta=0): +def make_rotating_tablename(prefix, delta=0, date=None): """Creates a tablename for table rotation based on a prefix with a given month delta.""" - date = get_month(delta=delta) + if not date: + date = get_month(delta=delta) return "{}_{}_{}".format(prefix, date.year, date.month) @@ -83,11 +84,11 @@ def create_rotating_message_table(prefix="message", read_throughput=5, ) -def get_rotating_message_table(prefix="message", delta=0): +def get_rotating_message_table(prefix="message", delta=0, date=None): """Gets the message table for the current month.""" db = DynamoDBConnection() dblist = db.list_tables()["TableNames"] - tablename = make_rotating_tablename(prefix, delta) + tablename = make_rotating_tablename(prefix, delta, date) if tablename not in dblist: return create_rotating_message_table(prefix=prefix, delta=delta) else: @@ -513,6 +514,16 @@ def fetch_messages(self, uaid, limit=10): consistent=True, limit=limit)) +class MessageSink(Message): + """DefaultDict base class for invalid message dates""" + + def __init__(self): + pass + + def fetch_messages(self, uaid, limit=0): + return list() + + class Router(object): """Create a Router table abstraction on top of a DynamoDB Table object""" def __init__(self, table, metrics): diff --git a/autopush/settings.py b/autopush/settings.py index 1a5a8bba..c5cd8c1b 100644 --- a/autopush/settings.py +++ b/autopush/settings.py @@ -19,11 +19,10 @@ get_router_table, get_storage_table, get_rotating_message_table, - make_rotating_tablename, preflight_check, Storage, Router, - Message + Message, ) from autopush.exceptions import InvalidTokenException from autopush.metrics import ( @@ -162,9 +161,12 @@ def __init__(self, self.router = Router(self.router_table, self.metrics) # Used to determine whether a connection is out of date with current - # db objects - self.current_msg_month = make_rotating_tablename(self._message_prefix) - self.current_month = datetime.date.today().month + # db objects. There are three noteworty cases: + # 1 "Last Month" the table requires a rollover. + # 2 "This Month" the most common case. + # 3 "Tomorrow" in the case where the system is actively rolling over + # and one node is not yet using the latest clock, no rollover + # required. self.create_initial_message_tables() # Run preflight check @@ -203,18 +205,35 @@ def message(self, value): """Setter to set the current message table""" self.message_tables[self.current_msg_month] = value + def add_tomorrow(self, today, this_month_table): + """check to see if tomorrow is a new month, in case we're a node that + has a slow clock.""" + tomorrow = get_rotating_message_table( + self._message_prefix, 0, today + datetime.timedelta(days=1)) + if tomorrow.table_name != this_month_table: + self.message_tables[tomorrow.table_name] = Message( + tomorrow, self.metrics) + def create_initial_message_tables(self): """Initializes a dict of the initial rotating messages tables. - An entry for last months table, and an entry for this months table. + An entry for last months table, an entry for this months table, + and an entry for tomorrow, if it starts a new month. """ + today = datetime.date.today() last_month = get_rotating_message_table(self._message_prefix, -1) this_month = get_rotating_message_table(self._message_prefix) + self.current_month = today.month + self.current_msg_month = this_month.table_name self.message_tables = { last_month.table_name: Message(last_month, self.metrics), - this_month.table_name: Message(this_month, self.metrics), + this_month.table_name: Message(this_month, self.metrics) } + # Due to issues around mocking "datetime" and desire to + # reduce excessive hotpath function calls, we have to skip this + # line + self.add_tomorrow(today, this_month.table_name) @inlineCallbacks def update_rotating_tables(self): @@ -226,6 +245,7 @@ def update_rotating_tables(self): """ today = datetime.date.today() + self.add_tomorrow(today, self.current_msg_month) if today.month == self.current_month: # No change in month, we're fine. returnValue(False) @@ -240,7 +260,6 @@ def update_rotating_tables(self): self.current_msg_month = message_table.table_name self.message_tables[self.current_msg_month] = \ Message(message_table, self.metrics) - returnValue(True) def update(self, **kwargs): diff --git a/autopush/tests/test_integration.py b/autopush/tests/test_integration.py index 41925cda..8feea84c 100644 --- a/autopush/tests/test_integration.py +++ b/autopush/tests/test_integration.py @@ -291,7 +291,7 @@ def disconnect(self): self.ws.close() def sleep(self, duration): - time.sleep(duration) + time.sleep(duration) # pragma: nocover def wait_for(self, func): """Waits several seconds for a function to return True""" diff --git a/autopush/tests/test_websocket.py b/autopush/tests/test_websocket.py index a451f7bb..b38daebd 100644 --- a/autopush/tests/test_websocket.py +++ b/autopush/tests/test_websocket.py @@ -1,4 +1,5 @@ import json +import datetime import time import uuid from hashlib import sha256 @@ -18,7 +19,9 @@ from twisted.trial import unittest import autopush.db as db -from autopush.db import create_rotating_message_table +from autopush.db import ( + create_rotating_message_table, +) from autopush.settings import AutopushSettings from autopush.websocket import ( PushState, @@ -405,6 +408,131 @@ def wait_for_agent_call(): # pragma: nocover reactor.callLater(0.1, wait_for_agent_call) return d + def test_hello_old(self): + orig_uaid = "deadbeef12345678decafbad12345678" + # router.register_user returns (registered, previous + target_day = datetime.date(2016, 2, 29) + msg_day = datetime.date(2015, 12, 15) + msg_date = "{}_{}_{}".format( + self.proto.ap_settings._message_prefix, + msg_day.year, + msg_day.month) + msg_data = { + "router_type": "webpush", + "node_id": "http://localhost", + "last_connect": int(msg_day.strftime("%s")), + "current_month": msg_date, + } + + def fake_msg(data): + return (True, msg_data, data) + + mock_msg = Mock(wraps=db.Message) + mock_msg.fetch_messages.return_value = [] + self.proto.ap_settings.router.register_user = fake_msg + # massage message_tables to include our fake range + mt = self.proto.ps.settings.message_tables + for k in mt.keys(): + del(mt[k]) + mt['message_2016_1'] = mock_msg + mt['message_2016_2'] = mock_msg + mt['message_2016_3'] = mock_msg + with patch.object(datetime, 'date', + Mock(wraps=datetime.date)) as patched: + patched.today.return_value = target_day + self._connect() + self._send_message(dict(messageType="hello", + uaid=orig_uaid, + channelIDs=[], + use_webpush=True)) + + def check_result(msg): + eq_(self.proto.ps.rotate_message_table, False) + # it's fine you've not connected in a while, but + # you should recycle your endpoints since they're probably + # invalid by now anyway. + eq_(msg["status"], 200) + ok_(msg["uaid"] != orig_uaid) + + return self._check_response(check_result) + + def test_hello_tomorrow(self): + orig_uaid = "deadbeef12345678decafbad12345678" + # router.register_user returns (registered, previous + target_day = datetime.date(2016, 2, 29) + msg_day = datetime.date(2016, 3, 1) + msg_date = "{}_{}_{}".format( + self.proto.ap_settings._message_prefix, + msg_day.year, + msg_day.month) + msg_data = { + "router_type": "webpush", + "node_id": "http://localhost", + "last_connect": int(msg_day.strftime("%s")), + "current_month": msg_date, + } + + def fake_msg(data): + return (True, msg_data, data) + + mock_msg = Mock(wraps=db.Message) + mock_msg.fetch_messages.return_value = [] + self.proto.ap_settings.router.register_user = fake_msg + # massage message_tables to include our fake range + mt = self.proto.ps.settings.message_tables + for k in mt.keys(): + del(mt[k]) + mt['message_2016_1'] = mock_msg + mt['message_2016_2'] = mock_msg + mt['message_2016_3'] = mock_msg + with patch.object(datetime, 'date', + Mock(wraps=datetime.date)) as patched: + patched.today.return_value = target_day + self._connect() + self._send_message(dict(messageType="hello", + uaid=orig_uaid, + channelIDs=[], + use_webpush=True)) + + def check_result(msg): + eq_(self.proto.ps.rotate_message_table, False) + # it's fine you've not connected in a while, but + # you should recycle your endpoints since they're probably + # invalid by now anyway. + eq_(msg["status"], 200) + eq_(msg["uaid"], orig_uaid) + + return self._check_response(check_result) + + def test_add_tomorrow(self): + today = datetime.date(2016, 2, 29) + yester = datetime.date(2016, 1, 1) + tomorrow = datetime.date(2016, 3, 1) + today_table = "{}_{}_{}".format( + self.proto.ap_settings._message_prefix, + today.year, + today.month) + yester_table = "{}_{}_{}".format( + self.proto.ap_settings._message_prefix, + yester.year, + yester.month) + tomorrow_table = "{}_{}_{}".format( + self.proto.ap_settings._message_prefix, + tomorrow.year, + tomorrow.month) + + mock_msg = Mock(wraps=db.Message) + mock_msg.fetch_messages.return_value = [] + mt = self.proto.ps.settings.message_tables + for k in mt.keys(): + del(mt[k]) + mt[yester_table] = mock_msg + mt[today_table] = mock_msg + + self._connect() + self.proto.ps.settings.add_tomorrow(today, today_table) + ok_(tomorrow_table in self.proto.ps.settings.message_tables) + def test_hello(self): self._connect() self._send_message(dict(messageType="hello", channelIDs=[])) diff --git a/autopush/websocket.py b/autopush/websocket.py index 00fde0b1..bc0a2854 100644 --- a/autopush/websocket.py +++ b/autopush/websocket.py @@ -773,9 +773,14 @@ def _check_message_table_rotation(self, previous): self.transport.pauseProducing() # Check for table rotation cur_month = previous.get("current_month") + # Previous month user or new user, flag for message rotation and + # set the message_month to the router month if cur_month != self.ps.message_month: - # Previous month user or new user, flag for message rotation and - # set the message_month to the router month + if cur_month not in self.ps.settings.message_tables.keys(): + # This UAID has expired. Force client to reregister. + self.ps.uaid = uuid.uuid4().hex + self._finish_webpush_hello() + return self.ps.message_month = cur_month self.ps.rotate_message_table = True