def QueueNotification(self, notification=None, timestamp=None, **kw):
"""Queues a notification for a flow."""
if notification is None:
notification = rdf_flows.GrrNotification(**kw)
session_id = notification.session_id
if session_id:
if timestamp is None:
timestamp = self.frozen_timestamp
notification.timestamp = timestamp
# We must not store more than one notification per session id and
# timestamp or there will be race conditions. We therefore only keep
# the one with the highest request number (indicated by last_status).
# Note that timestamp might be None. In that case we also only want
# to keep the latest.
if timestamp is None:
ts_str = "None"
else:
ts_str = int(timestamp)
key = "%s!%s" % (session_id, ts_str)
existing = self.notifications.get(key)
if existing is not None:
if existing.last_status < notification.last_status:
self.notifications[key] = notification
else:
self.notifications[key] = notification
def CheckNotificationsDisappear(self, session_id):
worker_obj = worker_lib.GRRWorker(token=self.token)
manager = queue_manager.QueueManager(token=self.token)
notification = rdf_flows.GrrNotification(session_id=session_id)
with data_store.DB.GetMutationPool() as pool:
manager.NotifyQueue(notification, mutation_pool=pool)
notifications = manager.GetNotificationsByPriority(queues.FLOWS).get(
notification.priority, [])
# Check the notification is there. With multiple worker queue shards we can
# get other notifications such as for audit event listeners, so we need to
# filter out ours.
notifications = [
x for x in notifications if x.session_id == session_id
]
self.assertEqual(len(notifications), 1)
# Process all messages
worker_obj.RunOnce()
worker_obj.thread_pool.Join()
notifications = manager.GetNotificationsByPriority(queues.FLOWS).get(
notification.priority, [])
notifications = [
x for x in notifications if x.session_id == session_id
]
# Check the notification is now gone.
self.assertEqual(len(notifications), 0)
def ExpireRules(self):
"""Removes any rules with an expiration date in the past."""
rules = self.Get(self.Schema.RULES)
new_rules = self.Schema.RULES()
now = time.time() * 1e6
expired_session_ids = set()
for rule in rules:
if rule.expires > now:
new_rules.Append(rule)
else:
for action in rule.actions:
if action.hunt_id:
expired_session_ids.add(action.hunt_id)
if expired_session_ids:
with data_store.DB.GetMutationPool() as pool:
# Notify the worker to mark this hunt as terminated.
manager = queue_manager.QueueManager(token=self.token)
manager.MultiNotifyQueue(
[
rdf_flows.GrrNotification(session_id=session_id)
for session_id in expired_session_ids
],
mutation_pool=pool)
if len(new_rules) < len(rules):
self.Set(self.Schema.RULES, new_rules)
self.Flush()
def testNotificationReschedulingTTL(self):
"""Test that notifications are not rescheduled forever."""
with test_lib.FakeTime(10000):
worker_obj = worker_lib.GRRWorker(token=self.token)
flow_obj = self.FlowSetup("RaisingTestFlow")
session_id = flow_obj.session_id
flow_obj.Close()
with queue_manager.QueueManager(token=self.token) as manager:
notification = rdf_flows.GrrNotification(session_id=session_id,
timestamp=time.time(),
last_status=1)
with data_store.DB.GetMutationPool() as pool:
manager.NotifyQueue(notification, mutation_pool=pool)
notifications = manager.GetNotifications(queues.FLOWS)
# Check the notification is there.
notifications = [
n for n in notifications if n.session_id == session_id
]
self.assertEqual(len(notifications), 1)
delay = flow_runner.FlowRunner.notification_retry_interval
ttl = notification.ttl
for i in xrange(ttl - 1):
with test_lib.FakeTime(10000 + 100 + delay * (i + 1)):
# Process all messages.
worker_obj.RunOnce()
worker_obj.thread_pool.Join()
notifications = manager.GetNotifications(queues.FLOWS)
# Check the notification is for the correct session_id.
notifications = [
n for n in notifications if n.session_id == session_id
]
self.assertEqual(len(notifications), 1)
with test_lib.FakeTime(10000 + 100 + delay * ttl):
# Process all messages.
worker_obj.RunOnce()
worker_obj.thread_pool.Join()
notifications = manager.GetNotifications(queues.FLOWS)
self.assertEqual(len(notifications), 0)
def testReordering(self):
"""Check that out of order client messages are reordered."""
flow_obj = self.FlowSetup(flow_test_lib.FlowOrderTest.__name__)
# Simulate processing messages arriving in random order
message_ids = [2, 1, 4, 3, 5]
self.SendMessages(message_ids, flow_obj.session_id)
# Send the status message
self.SendOKStatus(6, flow_obj.session_id)
runner = flow_obj.GetRunner()
notification = rdf_flows.GrrNotification(
timestamp=rdfvalue.RDFDatetime.Now())
runner.ProcessCompletedRequests(notification)
# Check that the messages were processed in order
self.assertEqual(flow_obj.messages, [1, 2, 3, 4, 5])
def testNotificationRacesAreResolved(self):
# We need a random flow object for this test.
session_id = flow.StartFlow(client_id=self.client_id,
flow_name="WorkerSendingTestFlow",
token=self.token)
worker_obj = worker_lib.GRRWorker(token=self.token)
manager = queue_manager.QueueManager(token=self.token)
manager.DeleteNotification(session_id)
manager.Flush()
# We simulate a race condition here - the notification for request #1 is
# there but the actual request #1 is not. The worker should pick up the
# notification, notice that the request #1 is not there yet and reschedule
# the notification.
notification = rdf_flows.GrrNotification(session_id=session_id,
last_status=1)
with data_store.DB.GetMutationPool() as pool:
manager.NotifyQueue(notification, mutation_pool=pool)
notifications = manager.GetNotifications(queues.FLOWS)
# Check the notification is there.
notifications = [
n for n in notifications if n.session_id == session_id
]
self.assertEqual(len(notifications), 1)
# Process all messages
worker_obj.RunOnce()
worker_obj.thread_pool.Join()
delay = flow_runner.FlowRunner.notification_retry_interval
with test_lib.FakeTime(time.time() + 10 + delay):
requeued_notifications = manager.GetNotifications(queues.FLOWS)
# Check that there is a new notification.
notifications = [
n for n in notifications if n.session_id == session_id
]
self.assertEqual(len(requeued_notifications), 1)
self.assertEqual(requeued_notifications[0].first_queued,
notifications[0].first_queued)
self.assertNotEqual(requeued_notifications[0].timestamp,
notifications[0].timestamp)
def testAuthentication1(self):
"""Test that flows refuse to processes unauthenticated messages."""
flow_obj = self.FlowSetup(flow_test_lib.FlowOrderTest.__name__)
# Simulate processing messages arriving in random order
message_ids = [2, 1, 4, 3, 5]
self.SendMessages(message_ids,
flow_obj.session_id,
authenticated=False)
# Send the status message
self.SendOKStatus(6, flow_obj.session_id)
runner = flow_obj.GetRunner()
notification = rdf_flows.GrrNotification(
timestamp=rdfvalue.RDFDatetime.Now())
runner.ProcessCompletedRequests(notification)
# Now messages should actually be processed
self.assertEqual(flow_obj.messages, [])
def testAuthentication2(self):
"""Test that flows refuse to processes unauthenticated messages.
Here we try to simulate an attacker injecting unauthenticated
messages midstream.
The current implementation actually fails to process the entire
flow since the injected messages displace the real ones if they
arrive earlier. This can be an effective DoS against legitimate
clients but would require attackers to guess session ids.
"""
flow_obj = self.FlowSetup(flow_test_lib.FlowOrderTest.__name__)
# Simulate processing messages arriving in random order
message_ids = [1, 2]
self.SendMessages(message_ids, flow_obj.session_id, authenticated=True)
# Now suppose some of the messages are spoofed
message_ids = [3, 4, 5]
self.SendMessages(message_ids,
flow_obj.session_id,
authenticated=False)
# And now our real messages arrive
message_ids = [5, 6]
self.SendMessages(message_ids, flow_obj.session_id, authenticated=True)
# Send the status message
self.SendOKStatus(7, flow_obj.session_id)
runner = flow_obj.GetRunner()
notification = rdf_flows.GrrNotification(
timestamp=rdfvalue.RDFDatetime.Now())
runner.ProcessCompletedRequests(notification)
# Some messages should actually be processed
self.assertEqual(flow_obj.messages, [1, 2, 5, 6])