def testReceiveMessagesWithStatus(self):
"""Receiving a sequence of messages with a status."""
flow_obj = self.FlowSetup("FlowOrderTest")
session_id = flow_obj.session_id
messages = [
rdf_flows.GrrMessage(request_id=1,
response_id=i,
session_id=session_id,
payload=rdfvalue.RDFInteger(i),
task_id=15) for i in range(1, 10)
]
# Now add the status message
status = rdf_flows.GrrStatus(
status=rdf_flows.GrrStatus.ReturnedStatus.OK)
messages.append(
rdf_flows.GrrMessage(request_id=1,
response_id=len(messages) + 1,
task_id=15,
session_id=messages[0].session_id,
payload=status,
type=rdf_flows.GrrMessage.Type.STATUS))
self.server.ReceiveMessages(self.client_id, messages)
# Make sure the task is still on the client queue
manager = queue_manager.QueueManager(token=self.token)
tasks_on_client_queue = manager.Query(self.client_id, 100)
self.assertEqual(len(tasks_on_client_queue), 1)
stored_messages = data_store.DB.ReadResponsesForRequestId(
session_id, 1, token=self.token)
self.assertEqual(len(stored_messages), len(messages))
stored_messages.sort(key=lambda m: m.response_id)
# Check that messages were stored correctly
for stored_message, message in zip(stored_messages, messages):
self.assertRDFValuesEqual(stored_message, message)
def StartClients(cls, hunt_id, client_ids, token=None):
"""This method is called by the foreman for each client it discovers.
Note that this function is performance sensitive since it is called by the
foreman for every client which needs to be scheduled.
Args:
hunt_id: The hunt to schedule.
client_ids: List of clients that should be added to the hunt.
token: An optional access token to use.
"""
token = token or access_control.ACLToken(username="******",
reason="hunting")
with queue_manager.QueueManager(token=token) as flow_manager:
for client_id in client_ids:
# Now we construct a special response which will be sent to the hunt
# flow. Randomize the request_id so we do not overwrite other messages
# in the queue.
state = rdfvalue.RequestState(id=utils.PRNG.GetULong(),
session_id=hunt_id,
client_id=client_id,
next_state="AddClient")
# Queue the new request.
flow_manager.QueueRequest(hunt_id, state)
# Send a response.
msg = rdfvalue.GrrMessage(session_id=hunt_id,
request_id=state.id,
response_id=1,
auth_state=rdfvalue.GrrMessage.
AuthorizationState.AUTHENTICATED,
type=rdfvalue.GrrMessage.Type.STATUS,
payload=rdfvalue.GrrStatus())
flow_manager.QueueResponse(hunt_id, msg)
# And notify the worker about it.
flow_manager.QueueNotification(session_id=hunt_id)
def testReceiveMessagesWithStatus(self):
"""Receiving a sequence of messages with a status."""
flow_obj = self.FlowSetup("FlowOrderTest")
session_id = flow_obj.session_id
messages = [
rdf_flows.GrrMessage(request_id=1,
response_id=i,
session_id=session_id,
payload=rdfvalue.RDFInteger(i),
task_id=15) for i in range(1, 10)
]
# Now add the status message
status = rdf_flows.GrrStatus(
status=rdf_flows.GrrStatus.ReturnedStatus.OK)
messages.append(
rdf_flows.GrrMessage(request_id=1,
response_id=len(messages) + 1,
task_id=15,
session_id=messages[0].session_id,
payload=status,
type=rdf_flows.GrrMessage.Type.STATUS))
self.server.ReceiveMessages(self.client_id, messages)
# Make sure the task is still on the client queue
manager = queue_manager.QueueManager(token=self.token)
tasks_on_client_queue = manager.Query(self.client_id, 100)
self.assertEqual(len(tasks_on_client_queue), 1)
# Check that messages were stored correctly
for message in messages:
stored_message, _ = data_store.DB.Resolve(
session_id.Add("state/request:00000001"),
manager.FLOW_RESPONSE_TEMPLATE % (1, message.response_id),
token=self.token)
stored_message = rdf_flows.GrrMessage(stored_message)
self.assertRDFValueEqual(stored_message, message)
def BuildTable(self, start_row, end_row, request):
session_id = request.REQ.get("flow", "")
if not session_id:
return
manager = queue_manager.QueueManager(token=request.token)
for i, (request, responses) in enumerate(
manager.FetchRequestsAndResponses(rdfvalue.RDFURN(session_id))):
if request.id == 0:
continue
if i < start_row:
continue
if i > end_row:
break
# Tie up the request to each response to make it easier to render.
self.AddCell(i, "ID", manager.FLOW_REQUEST_TEMPLATE % request.id)
self.AddCell(i, "Request", request)
if responses:
self.AddCell(i, "Last Response", responses[-1])
def testNotificationRacesAreResolved(self):
# We need a random flow object for this test.
session_id = flow.GRRFlow.StartFlow(client_id=self.client_id,
flow_name="WorkerSendingTestFlow",
token=self.token)
worker_obj = worker.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)
manager.NotifyQueue(notification)
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 = config_lib.CONFIG["Worker.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 WakeStuckFlow(session_id):
"""Wake up stuck flows.
A stuck flow is one which is waiting for the client to do something, but the
client requests have been removed from the client queue. This can happen if
the system is too loaded and the client messages have TTLed out. In this case
we reschedule the client requests for this session.
Args:
session_id: The session for the flow to wake.
Returns:
The total number of client messages re-queued.
"""
session_id = rdfvalue.SessionID(session_id)
woken = 0
checked_pending = False
with queue_manager.QueueManager() as manager:
for request, responses in manager.FetchRequestsAndResponses(session_id):
# We need to check if there are client requests pending.
if not checked_pending:
task = manager.Query(request.client_id,
task_id="task:%s" % request.request.task_id)
if task:
# Client has tasks pending already.
return
checked_pending = True
if not responses or responses[-1].type != rdfvalue.GrrMessage.Type.STATUS:
manager.QueueClientMessage(request.request)
woken += 1
if responses and responses[-1].type == rdfvalue.GrrMessage.Type.STATUS:
manager.QueueNotification(session_id)
return woken
def CheckNotificationsDisappear(self, session_id):
worker_obj = worker.GRRWorker(worker.DEFAULT_WORKER_QUEUE,
token=self.token)
manager = queue_manager.QueueManager(token=self.token)
notification = rdfvalue.GrrNotification(session_id=session_id)
manager.NotifyQueue(notification)
notifications = manager.GetNotificationsByPriority(
worker.DEFAULT_WORKER_QUEUE).get(notification.priority, [])
# Check the notification is there.
self.assertEqual(len(notifications), 1)
self.assertEqual(notifications[0].session_id, session_id)
# Process all messages
worker_obj.RunOnce()
worker_obj.thread_pool.Join()
notifications = manager.GetNotificationsByPriority(
worker.DEFAULT_WORKER_QUEUE).get(notification.priority, [])
# Check the notification is now gone.
self.assertEqual(len(notifications), 0)
def testGetNotificationsForAllShards(self):
manager = queue_manager.QueueManager(token=self.token)
print "notification shards:" + str(manager.num_notification_shards)
manager.QueueNotification(session_id=rdfvalue.SessionID(
base="aff4:/hunts", queue=queues.HUNTS, flow_name="42"))
manager.Flush()
manager.QueueNotification(session_id=rdfvalue.SessionID(
base="aff4:/hunts", queue=queues.HUNTS, flow_name="43"))
manager.Flush()
live_shard_count = 0
for _ in range(manager.num_notification_shards):
shard_sessions = manager.GetNotifications(queues.HUNTS)
print "retrieved sessions:" + str(shard_sessions)
self.assertLess(len(shard_sessions), 2)
if len(shard_sessions) == 1:
live_shard_count += 1
self.assertEqual(live_shard_count, 2)
notifications = manager.GetNotificationsForAllShards(queues.HUNTS)
self.assertEqual(len(notifications), 2)
def _ScheduleResponseAndStatus(self, client_id, flow_id):
with queue_manager.QueueManager(token=self.token) as flow_manager:
# Schedule a response.
flow_manager.QueueResponse(
flow_id,
rdf_flows.GrrMessage(
source=client_id,
session_id=flow_id,
payload=rdf_protodict.DataBlob(string="Helllo"),
request_id=1,
response_id=1))
# And a STATUS message.
flow_manager.QueueResponse(
flow_id,
rdf_flows.GrrMessage(
source=client_id,
session_id=flow_id,
payload=rdf_flows.GrrStatus(
status=rdf_flows.GrrStatus.ReturnedStatus.OK),
request_id=1,
response_id=2,
type=rdf_flows.GrrMessage.Type.STATUS))
def testPriorityScheduling(self):
test_queue = rdfvalue.RDFURN("fooReschedule")
tasks = []
for i in range(10):
msg = rdf_flows.GrrMessage(session_id="Test%d" % i,
priority=i % 3,
queue=test_queue,
generate_task_id=True)
tasks.append(msg)
manager = queue_manager.QueueManager(token=self.token)
with data_store.DB.GetMutationPool(token=self.token) as pool:
manager.Schedule(tasks, pool)
tasks = manager.QueryAndOwn(test_queue, lease_seconds=100, limit=3)
self.assertEqual(len(tasks), 3)
for task in tasks:
self.assertEqual(task.priority, 2)
tasks = manager.QueryAndOwn(test_queue, lease_seconds=100, limit=3)
self.assertEqual(len(tasks), 3)
for task in tasks:
self.assertEqual(task.priority, 1)
tasks = manager.QueryAndOwn(test_queue, lease_seconds=100, limit=100)
self.assertEqual(len(tasks), 4)
for task in tasks:
self.assertEqual(task.priority, 0)
# Now for Query.
tasks = manager.Query(test_queue, limit=100)
self.assertEqual(len(tasks), 10)
self.assertEqual([task.priority for task in tasks],
[2, 2, 2, 1, 1, 1, 0, 0, 0, 0])
def testReSchedule(self):
"""Test the ability to re-schedule a task."""
test_queue = rdfvalue.RDFURN("fooReschedule")
task = rdf_flows.GrrMessage(queue=test_queue,
task_ttl=5,
session_id="aff4:/Test",
generate_task_id=True)
manager = queue_manager.QueueManager(token=self.token)
with data_store.DB.GetMutationPool(token=self.token) as pool:
manager.Schedule([task], pool)
# Get a lease on the task
tasks = manager.QueryAndOwn(test_queue, lease_seconds=100, limit=100)
self.assertEqual(len(tasks), 1)
# Record the task id
original_id = tasks[0].task_id
# If we try to get another lease on it we should fail
tasks_2 = manager.QueryAndOwn(test_queue, lease_seconds=100, limit=100)
self.assertEqual(len(tasks_2), 0)
# Now we reschedule it
with data_store.DB.GetMutationPool(token=self.token) as pool:
manager.Schedule([task], pool)
# The id should not change
self.assertEqual(tasks[0].task_id, original_id)
# If we try to get another lease on it we should not fail
tasks = manager.QueryAndOwn(test_queue, lease_seconds=100, limit=100)
self.assertEqual(len(tasks), 1)
# But the id should not change
self.assertEqual(tasks[0].task_id, original_id)
def SendResponse(self,
session_id,
data,
client_id=None,
well_known=False,
request_id=None):
if not isinstance(data, rdfvalue.RDFValue):
data = rdf_protodict.DataBlob(string=data)
if well_known:
request_id, response_id = 0, 12345
else:
request_id, response_id = request_id or 1, 1
with queue_manager.QueueManager(token=self.token) as flow_manager:
flow_manager.QueueResponse(
session_id,
rdf_flows.GrrMessage(source=client_id,
session_id=session_id,
payload=data,
request_id=request_id,
response_id=response_id))
if not well_known:
# For normal flows we have to send a status as well.
flow_manager.QueueResponse(
session_id,
rdf_flows.GrrMessage(
source=client_id,
session_id=session_id,
payload=rdf_flows.GrrStatus(
status=rdf_flows.GrrStatus.ReturnedStatus.OK),
request_id=request_id,
response_id=response_id + 1,
type=rdf_flows.GrrMessage.Type.STATUS))
flow_manager.QueueNotification(session_id=session_id,
last_status=request_id)
timestamp = flow_manager.frozen_timestamp
return timestamp
def testCallClientWellKnown(self):
"""Well known flows can also call the client."""
cls = flow.GRRFlow.classes["GetClientStatsAuto"]
flow_obj = cls(cls.well_known_session_id, mode="rw", token=self.token)
flow_obj.CallClient(self.client_id, admin.GetClientStats)
# Check that a message went out to the client
manager = queue_manager.QueueManager(token=self.token)
tasks = manager.Query(self.client_id, limit=100)
self.assertEqual(len(tasks), 1)
message = tasks[0]
# If we don't specify where to send the replies, they go to the devnull flow
devnull = flow.GRRFlow.classes["IgnoreResponses"]
self.assertEqual(message.session_id, devnull.well_known_session_id)
self.assertEqual(message.request_id, 0)
self.assertEqual(message.name, admin.GetClientStats.__name__)
messages = []
def StoreMessage(_, msg):
messages.append(msg)
with utils.Stubber(devnull, "ProcessMessage", StoreMessage):
client_mock = action_mocks.ActionMock(admin.GetClientStats)
for _ in test_lib.TestFlowHelper(
"ClientActionRunner",
client_mock,
client_id=self.client_id,
action="GetClientStats",
token=self.token):
pass
# Make sure the messages arrived.
self.assertEqual(len(messages), 1)
def testMultipleNotificationsForTheSameSessionId(self):
manager = queue_manager.QueueManager(token=self.token)
manager.QueueNotification(
session_id=rdfvalue.SessionID(
base="aff4:/hunts", queue=queues.HUNTS, flow_name="123456"),
timestamp=(self._current_mock_time + 10) * 1e6)
manager.QueueNotification(
session_id=rdfvalue.SessionID(
base="aff4:/hunts", queue=queues.HUNTS, flow_name="123456"),
timestamp=(self._current_mock_time + 20) * 1e6)
manager.QueueNotification(
session_id=rdfvalue.SessionID(
base="aff4:/hunts", queue=queues.HUNTS, flow_name="123456"),
timestamp=(self._current_mock_time + 30) * 1e6)
manager.Flush()
self.assertEqual(len(manager.GetNotificationsForAllShards(queues.HUNTS)), 0)
self._current_mock_time += 10
self.assertEqual(len(manager.GetNotificationsForAllShards(queues.HUNTS)), 1)
manager.DeleteNotification(
rdfvalue.SessionID(
base="aff4:/hunts", queue=queues.HUNTS, flow_name="123456"))
self._current_mock_time += 10
self.assertEqual(len(manager.GetNotificationsForAllShards(queues.HUNTS)), 1)
manager.DeleteNotification(
rdfvalue.SessionID(
base="aff4:/hunts", queue=queues.HUNTS, flow_name="123456"))
self._current_mock_time += 10
self.assertEqual(len(manager.GetNotificationsForAllShards(queues.HUNTS)), 1)
manager.DeleteNotification(
rdfvalue.SessionID(
base="aff4:/hunts", queue=queues.HUNTS, flow_name="123456"))
self._current_mock_time += 10
self.assertEqual(len(manager.GetNotificationsForAllShards(queues.HUNTS)), 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:
# 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])
if len(new_rules) < len(rules):
self.Set(self.Schema.RULES, new_rules)
self.Flush()
def testReceiveUnsolicitedClientMessage(self):
flow_obj = self.FlowSetup("FlowOrderTest")
session_id = flow_obj.session_id
status = rdf_flows.GrrStatus(
status=rdf_flows.GrrStatus.ReturnedStatus.OK)
messages = [
# This message has no task_id set...
rdf_flows.GrrMessage(request_id=1,
response_id=1,
session_id=session_id,
payload=rdfvalue.RDFInteger(1),
task_id=15),
rdf_flows.GrrMessage(request_id=1,
response_id=2,
session_id=session_id,
payload=status,
type=rdf_flows.GrrMessage.Type.STATUS)
]
self.server.ReceiveMessages(self.client_id, messages)
manager = queue_manager.QueueManager(token=self.token)
completed = list(manager.FetchCompletedRequests(session_id))
self.assertEqual(len(completed), 1)
def BuildTable(self, start_row, end_row, request):
client_id = rdf_client.ClientURN(request.REQ.get("client_id"))
now = rdfvalue.RDFDatetime().Now()
# Make a local QueueManager.
manager = queue_manager.QueueManager(token=request.token)
for i, task in enumerate(manager.Query(client_id, limit=end_row)):
if i < start_row:
continue
difference = now - task.eta
if difference > 0:
self.AddCell(i, "Status", dict(
icon="stock_yes", description="Available for Lease"))
else:
self.AddCell(i, "Status", dict(
icon="clock",
description="Leased for %s Seconds" % (difference / 1e6)))
self.AddCell(i, "ID", task.task_id)
self.AddCell(i, "Flow", task.session_id)
self.AddCell(i, "Due", rdfvalue.RDFDatetime(task.eta))
self.AddCell(i, "Client Action", task.name)
def testDelete(self):
"""Test that we can delete tasks."""
test_queue = rdfvalue.RDFURN("fooDelete")
task = rdf_flows.GrrMessage(queue=test_queue,
session_id="aff4:/Test",
generate_task_id=True)
with data_store.DB.GetMutationPool(token=self.token) as pool:
manager = queue_manager.QueueManager(token=self.token)
manager.Schedule([task], pool)
# Get a lease on the task
tasks = manager.QueryAndOwn(test_queue, lease_seconds=100, limit=100)
self.assertEqual(len(tasks), 1)
self.assertEqual(tasks[0].session_id, "aff4:/Test")
# Now delete the task
manager.Delete(test_queue, tasks)
# Should not exist in the table
value, ts = data_store.DB.Resolve(test_queue,
"task:%08d" % task.task_id,
token=self.token)
self.assertEqual(value, None)
self.assertEqual(ts, 0)
# If we try to get another lease on it we should fail - even after
# expiry time.
self._current_mock_time += 1000
tasks = manager.QueryAndOwn(test_queue, lease_seconds=100, limit=100)
self.assertEqual(len(tasks), 0)
def testNoNotificationRescheduling(self):
"""Test that no notifications are rescheduled when a flow raises."""
with test_lib.FakeTime(10000):
flow_obj = self.FlowSetup("RaisingTestFlow")
session_id = flow_obj.session_id
flow_obj.Close()
# Send the flow some messages.
self.SendResponse(session_id, "Hello1", request_id=1)
self.SendResponse(session_id, "Hello2", request_id=2)
self.SendResponse(session_id, "Hello3", request_id=3)
worker_obj = worker.GRRWorker(token=self.token)
# Process all messages.
worker_obj.RunOnce()
worker_obj.thread_pool.Join()
delay = flow_runner.FlowRunner.notification_retry_interval
with test_lib.FakeTime(10000 + 100 + delay):
manager = queue_manager.QueueManager(token=self.token)
self.assertFalse(
manager.GetNotificationsForAllShards(session_id.Queue()))
def Run(self):
with test_lib.FakeTime(42):
flow_urn = flow.GRRFlow.StartFlow(
flow_name=processes.ListProcesses.__name__,
client_id=self.client_id,
token=self.token)
mock = test_lib.MockClient(self.client_id, None, token=self.token)
while mock.Next():
pass
replace = {flow_urn.Basename(): "W:ABCDEF"}
manager = queue_manager.QueueManager(token=self.token)
requests_responses = manager.FetchRequestsAndResponses(flow_urn)
for request, responses in requests_responses:
replace[str(request.request.task_id)] = "42"
for response in responses:
replace[str(response.task_id)] = "42"
self.Check("GET",
"/api/clients/%s/flows/%s/requests" %
(self.client_id.Basename(), flow_urn.Basename()),
replace=replace)
def Handle(self, args, token=None):
flow_urn = args.flow_id.ResolveClientFlowURN(args.client_id,
token=token)
# Check if this flow really exists.
try:
aff4.FACTORY.Open(flow_urn,
aff4_type=flow.GRRFlow,
mode="r",
token=token)
except aff4.InstantiationError:
raise FlowNotFoundError()
result = ApiListFlowRequestsResult()
manager = queue_manager.QueueManager(token=token)
requests_responses = manager.FetchRequestsAndResponses(flow_urn)
stop = None
if args.count:
stop = args.offset + args.count
for request, responses in itertools.islice(requests_responses,
args.offset, stop):
if request.id == 0:
continue
api_request = ApiFlowRequest(
request_id=manager.FLOW_REQUEST_TEMPLATE % request.id,
request_state=request)
if responses:
api_request.responses = responses
result.items.append(api_request)
return result
def _ProcessCompletedRequests(self, notification):
"""Does the actual processing of the completed requests."""
# First ensure that client messages are all removed. NOTE: We make a new
# queue manager here because we want only the client messages to be removed
# ASAP. This must happen before we actually run the flow to ensure the
# client requests are removed from the client queues.
with queue_manager.QueueManager(token=self.token) as manager:
for request, _ in manager.FetchCompletedRequests(
self.session_id, timestamp=(0, notification.timestamp)):
# Requests which are not destined to clients have no embedded request
# message.
if request.HasField("request"):
manager.DeQueueClientRequest(request.client_id,
request.request.task_id)
# The flow is dead - remove all outstanding requests and responses.
if not self.IsRunning():
self.queue_manager.DestroyFlowStates(self.session_id)
return
processing = []
while True:
try:
# Here we only care about completed requests - i.e. those requests with
# responses followed by a status message.
for request, responses in self.queue_manager.FetchCompletedResponses(
self.session_id, timestamp=(0, notification.timestamp)):
if request.id == 0:
continue
if not responses:
break
# We are missing a needed request - maybe its not completed yet.
if request.id > self.context.next_processed_request:
stats.STATS.IncrementCounter("grr_response_out_of_order")
break
# Not the request we are looking for - we have seen it before
# already.
if request.id < self.context.next_processed_request:
self.queue_manager.DeleteFlowRequestStates(self.session_id, request)
continue
if not responses:
continue
# Do we have all the responses here? This can happen if some of the
# responses were lost.
if len(responses) != responses[-1].response_id:
# If we can retransmit do so. Note, this is different from the
# automatic retransmission facilitated by the task scheduler (the
# Task.task_ttl field) which would happen regardless of these.
if request.transmission_count < 5:
stats.STATS.IncrementCounter("grr_request_retransmission_count")
request.transmission_count += 1
self.ReQueueRequest(request)
break
# If we get here its all good - run the flow.
if self.IsRunning():
self.flow_obj.HeartBeat()
self.RunStateMethod(request.next_state, request, responses)
# Quit early if we are no longer alive.
else:
break
# At this point we have processed this request - we can remove it and
# its responses from the queue.
self.queue_manager.DeleteFlowRequestStates(self.session_id, request)
self.context.next_processed_request += 1
self.DecrementOutstandingRequests()
# Are there any more outstanding requests?
if not self.OutstandingRequests():
# Allow the flow to cleanup
if self.IsRunning() and self.context.current_state != "End":
self.RunStateMethod("End")
# Rechecking the OutstandingRequests allows the End state (which was
# called above) to issue further client requests - hence postpone
# termination.
if not self.OutstandingRequests():
# TODO(user): Deprecate in favor of 'flow_completions' metric.
stats.STATS.IncrementCounter("grr_flow_completed_count")
stats.STATS.IncrementCounter(
"flow_completions", fields=[self.flow_obj.Name()])
logging.debug("Destroying session %s(%s) for client %s",
self.session_id,
self.flow_obj.Name(), self.runner_args.client_id)
self.flow_obj.Terminate()
# We are done here.
return
except queue_manager.MoreDataException:
# Join any threads.
for event in processing:
event.wait()
# We did not read all the requests/responses in this run in order to
# keep a low memory footprint and have to make another pass.
self.FlushMessages()
self.flow_obj.Flush()
continue
finally:
# Join any threads.
for event in processing:
event.wait()
def __init__(self, flow_obj, parent_runner=None, runner_args=None,
token=None):
"""Constructor for the Flow Runner.
Args:
flow_obj: The flow object this runner will run states for.
parent_runner: The parent runner of this runner.
runner_args: A FlowRunnerArgs() instance containing initial values. If not
specified, we use the runner_args from the flow_obj.
token: An instance of access_control.ACLToken security token.
"""
self.token = token or flow_obj.token
self.parent_runner = parent_runner
# If we have a parent runner, we use its queue manager.
if parent_runner is not None:
self.queue_manager = parent_runner.queue_manager
else:
# Otherwise we use a new queue manager.
self.queue_manager = queue_manager.QueueManager(token=self.token)
self.queued_replies = []
self.outbound_lock = threading.Lock()
self.flow_obj = flow_obj
# Initialize from a new runner args proto.
if runner_args is not None:
self.runner_args = runner_args
self.session_id = self.GetNewSessionID()
self.flow_obj.urn = self.session_id
# Flow state does not have a valid context, we need to create one.
self.context = self.InitializeContext(runner_args)
self.flow_obj.context = self.context
self.context.session_id = self.session_id
else:
# Retrieve args from the flow object's context. The flow object is
# responsible for storing our context, although they do not generally
# access it directly.
self.context = self.flow_obj.context
self.runner_args = self.flow_obj.runner_args
# Populate the flow object's urn with the session id.
self.flow_obj.urn = self.session_id = self.context.session_id
# Sent replies are cached so that they can be processed by output plugins
# when the flow is saved.
self.sent_replies = []
# If we're running a child flow and send_replies=True, but
# write_intermediate_results=False, we don't want to create an output
# collection object. We also only want to create it if runner_args are
# passed as a parameter, so that the new context is initialized.
#
# We can't create the collection as part of InitializeContext, as flow's
# urn is not known when InitializeContext runs.
if runner_args is not None and self.IsWritingResults():
with data_store.DB.GetMutationPool(token=self.token) as mutation_pool:
self.CreateCollections(mutation_pool)
def ReceiveMessages(self, client_id, messages):
"""Receives and processes the messages from the source.
For each message we update the request object, and place the
response in that request's queue. If the request is complete, we
send a message to the worker.
Args:
client_id: The client which sent the messages.
messages: A list of GrrMessage RDFValues.
"""
now = time.time()
with queue_manager.QueueManager(token=self.token,
store=self.data_store) as manager:
for session_id, msgs in utils.GroupBy(
messages, operator.attrgetter("session_id")).iteritems():
# Remove and handle messages to WellKnownFlows
unprocessed_msgs = self.HandleWellKnownFlows(msgs)
if not unprocessed_msgs:
continue
for msg in unprocessed_msgs:
manager.QueueResponse(msg)
for msg in unprocessed_msgs:
# Messages for well known flows should notify even though they don't
# have a status.
if msg.request_id == 0:
manager.QueueNotification(session_id=msg.session_id,
priority=msg.priority)
# Those messages are all the same, one notification is enough.
break
elif msg.type == rdf_flows.GrrMessage.Type.STATUS:
# If we receive a status message from the client it means the client
# has finished processing this request. We therefore can de-queue it
# from the client queue. msg.task_id will raise if the task id is
# not set (message originated at the client, there was no request on
# the server), so we have to check .HasTaskID() first.
if msg.HasTaskID():
manager.DeQueueClientRequest(
client_id, msg.task_id)
manager.QueueNotification(session_id=msg.session_id,
priority=msg.priority,
last_status=msg.request_id)
stat = rdf_flows.GrrStatus(msg.payload)
if stat.status == rdf_flows.GrrStatus.ReturnedStatus.CLIENT_KILLED:
# A client crashed while performing an action, fire an event.
crash_details = rdf_client.ClientCrash(
client_id=client_id,
session_id=session_id,
backtrace=stat.backtrace,
crash_message=stat.error_message,
nanny_status=stat.nanny_status,
timestamp=rdfvalue.RDFDatetime.Now())
msg = rdf_flows.GrrMessage(
source=client_id,
payload=crash_details,
auth_state=(rdf_flows.GrrMessage.
AuthorizationState.AUTHENTICATED))
events.Events.PublishEvent("ClientCrash",
msg,
token=self.token)
logging.debug("Received %s messages from %s in %s sec", len(messages),
client_id,
time.time() - now)
def testQueueing(self):
"""Tests that queueing and fetching of requests and responses work."""
session_id = rdfvalue.SessionID(flow_name="test")
request = rdfvalue.RequestState(id=1,
client_id=self.client_id,
next_state="TestState",
session_id=session_id)
with queue_manager.QueueManager(token=self.token) as manager:
manager.QueueRequest(session_id, request)
# We only have one unanswered request on the queue.
all_requests = list(manager.FetchRequestsAndResponses(session_id))
self.assertEqual(len(all_requests), 1)
self.assertEqual(all_requests[0], (request, []))
# FetchCompletedRequests should return nothing now.
self.assertEqual(list(manager.FetchCompletedRequests(session_id)), [])
# Now queue more requests and responses:
with queue_manager.QueueManager(token=self.token) as manager:
# Start with request 2 - leave request 1 un-responded to.
for request_id in range(2, 5):
request = rdfvalue.RequestState(id=request_id,
client_id=self.client_id,
next_state="TestState",
session_id=session_id)
manager.QueueRequest(session_id, request)
response_id = None
for response_id in range(1, 10):
# Normal message.
manager.QueueResponse(
session_id,
rdfvalue.GrrMessage(request_id=request_id,
response_id=response_id))
# And a status message.
manager.QueueResponse(
session_id,
rdfvalue.GrrMessage(request_id=request_id,
response_id=response_id + 1,
type=rdfvalue.GrrMessage.Type.STATUS))
completed_requests = list(manager.FetchCompletedRequests(session_id))
self.assertEqual(len(completed_requests), 3)
# First completed message is request_id = 2 with 10 responses.
self.assertEqual(completed_requests[0][0].id, 2)
# Last message is the status message.
self.assertEqual(completed_requests[0][-1].type,
rdfvalue.GrrMessage.Type.STATUS)
self.assertEqual(completed_requests[0][-1].response_id, 10)
# Now fetch all the completed responses. Set the limit so we only fetch some
# of the responses.
completed_response = list(manager.FetchCompletedResponses(session_id))
self.assertEqual(len(completed_response), 3)
for i, (request, responses) in enumerate(completed_response, 2):
self.assertEqual(request.id, i)
self.assertEqual(len(responses), 10)
# Now check if the limit is enforced. The limit refers to the total number
# of responses to return. We ask for maximum 15 responses, so we should get
# a single request with 10 responses (since 2 requests will exceed the
# limit).
more_data = False
i = 0
try:
partial_response = manager.FetchCompletedResponses(session_id,
limit=15)
for i, (request, responses) in enumerate(partial_response, 2):
self.assertEqual(request.id, i)
self.assertEqual(len(responses), 10)
except queue_manager.MoreDataException:
more_data = True
# Returns the first request that is completed.
self.assertEqual(i, 3)
# Make sure the manager told us that more data is available.
self.assertTrue(more_data)
def SendMessage(self, message):
# Now messages are set in the data store
with queue_manager.QueueManager(token=self.token) as manager:
manager.QueueResponse(message.session_id, message)
def testHandleClientMessageRetransmission(self):
"""Check that requests get retransmitted but only if there is no status."""
# Make a new fake client
client_id = self.SetupClients(1)[0]
# Test the standard behavior.
base_time = 1000
msgs_recvd = []
default_ttl = rdf_flows.GrrMessage().task_ttl
with test_lib.FakeTime(base_time):
flow.GRRFlow.StartFlow(
client_id=client_id,
flow_name="SendingFlow",
message_count=1,
token=self.token)
for i in range(default_ttl):
with test_lib.FakeTime(base_time + i * (self.message_expiry_time + 1)):
tasks = self.server.DrainTaskSchedulerQueueForClient(client_id, 100000)
msgs_recvd.append(tasks)
# Should return a client message (ttl-1) times and nothing afterwards.
self.assertEqual(
map(bool, msgs_recvd),
[True] * (rdf_flows.GrrMessage().task_ttl - 1) + [False])
# Now we simulate that the workers are overloaded - the client messages
# arrive but do not get processed in time.
if default_ttl <= 3:
self.fail("TTL too low for this test.")
msgs_recvd = []
with test_lib.FakeTime(base_time):
flow_id = flow.GRRFlow.StartFlow(
client_id=client_id,
flow_name="SendingFlow",
message_count=1,
token=self.token)
for i in range(default_ttl):
if i == 2:
self._ScheduleResponseAndStatus(client_id, flow_id)
with test_lib.FakeTime(base_time + i * (self.message_expiry_time + 1)):
tasks = self.server.DrainTaskSchedulerQueueForClient(client_id, 100000)
msgs_recvd.append(tasks)
if not tasks:
# Even if the request has not been leased ttl times yet,
# it should be dequeued by now.
new_tasks = queue_manager.QueueManager(token=self.token).Query(
queue=rdf_client.ClientURN(client_id).Queue(), limit=1000)
self.assertEqual(len(new_tasks), 0)
# Should return a client message twice and nothing afterwards.
self.assertEqual(
map(bool, msgs_recvd),
[True] * 2 + [False] * (rdf_flows.GrrMessage().task_ttl - 2))
def RunOnce(self):
"""Processes one set of messages from Task Scheduler.
The worker processes new jobs from the task master. For each job
we retrieve the session from the Task Scheduler.
Returns:
Total number of messages processed by this call.
"""
start_time = time.time()
processed = 0
queue_manager = queue_manager_lib.QueueManager(token=self.token)
for queue in self.queues:
# Freezeing the timestamp used by queue manager to query/delete
# notifications to avoid possible race conditions.
queue_manager.FreezeTimestamp()
fetch_messages_start = time.time()
notifications_by_priority = queue_manager.GetNotificationsByPriority(
queue)
stats.STATS.RecordEvent("worker_time_to_retrieve_notifications",
time.time() - fetch_messages_start)
# Process stuck flows first
stuck_flows = notifications_by_priority.pop(
queue_manager.STUCK_PRIORITY, [])
if stuck_flows:
self.ProcessStuckFlows(stuck_flows, queue_manager)
notifications_available = []
for priority in sorted(notifications_by_priority, reverse=True):
for notification in notifications_by_priority[priority]:
# Filter out session ids we already tried to lock but failed.
if notification.session_id not in self.queued_flows:
notifications_available.append(notification)
try:
# If we spent too much time processing what we have so far, the
# active_sessions list might not be current. We therefore break here
# so we can re-fetch a more up to date version of the list, and try
# again later. The risk with running with an old active_sessions list
# is that another worker could have already processed this message,
# and when we try to process it, there is nothing to do - costing us a
# lot of processing time. This is a tradeoff between checking the data
# store for current information and processing out of date
# information.
processed += self.ProcessMessages(
notifications_available, queue_manager,
self.RUN_ONCE_MAX_SECONDS - (time.time() - start_time))
# We need to keep going no matter what.
except Exception as e: # pylint: disable=broad-except
logging.error("Error processing message %s. %s.", e,
traceback.format_exc())
stats.STATS.IncrementCounter("grr_worker_exceptions")
if flags.FLAGS.debug:
pdb.post_mortem()
queue_manager.UnfreezeTimestamp()
# If we have spent too much time, stop.
if (time.time() - start_time) > self.RUN_ONCE_MAX_SECONDS:
return processed
return processed
def testNoValidStatusRaceIsResolved(self):
# This tests for the regression of a long standing race condition we saw
# where notifications would trigger the reading of another request that
# arrives later but wasn't completely written to the database yet.
# Timestamp based notification handling should eliminate this bug.
# We need a random flow object for this test.
session_id = flow.GRRFlow.StartFlow(client_id=self.client_id,
flow_name="WorkerSendingTestFlow",
token=self.token)
worker_obj = worker.GRRWorker(token=self.token)
manager = queue_manager.QueueManager(token=self.token)
manager.DeleteNotification(session_id)
manager.Flush()
# We have a first request that is complete (request_id 1, response_id 1).
self.SendResponse(session_id, "Response 1")
# However, we also have request #2 already coming in. The race is that
# the queue manager might write the status notification to
# session_id/state as "status:00000002" but not the status response
# itself yet under session_id/state/request:00000002
request_id = 2
response_id = 1
flow_manager = queue_manager.QueueManager(token=self.token)
flow_manager.FreezeTimestamp()
flow_manager.QueueResponse(
session_id,
rdf_flows.GrrMessage(
source=self.client_id,
session_id=session_id,
payload=rdf_protodict.DataBlob(string="Response 2"),
request_id=request_id,
response_id=response_id))
status = rdf_flows.GrrMessage(
source=self.client_id,
session_id=session_id,
payload=rdf_flows.GrrStatus(
status=rdf_flows.GrrStatus.ReturnedStatus.OK),
request_id=request_id,
response_id=response_id + 1,
type=rdf_flows.GrrMessage.Type.STATUS)
# Now we write half the status information.
subject = session_id.Add("state")
queue = flow_manager.to_write.setdefault(subject, {})
queue.setdefault(flow_manager.FLOW_STATUS_TEMPLATE % request_id,
[]).append((status.SerializeToString(), None))
flow_manager.Flush()
# We make the race even a bit harder by saying the new notification gets
# written right before the old one gets deleted. If we are not careful here,
# we delete the new notification as well and the flow becomes stuck.
def WriteNotification(self, arg_session_id, start=None, end=None):
if arg_session_id == session_id:
flow_manager.QueueNotification(session_id=arg_session_id)
flow_manager.Flush()
self.DeleteNotification.old_target(self,
arg_session_id,
start=start,
end=end)
with utils.Stubber(queue_manager.QueueManager, "DeleteNotification",
WriteNotification):
# This should process request 1 but not touch request 2.
worker_obj.RunOnce()
worker_obj.thread_pool.Join()
flow_obj = aff4.FACTORY.Open(session_id, token=self.token)
self.assertFalse(flow_obj.context.backtrace)
self.assertNotEqual(flow_obj.context.state,
rdf_flows.FlowContext.State.ERROR)
request_data = data_store.DB.ResolvePrefix(session_id.Add("state"),
"flow:",
token=self.token)
subjects = [r[0] for r in request_data]
# Make sure the status field and the original request are still there.
self.assertIn("flow:request:00000002", subjects)
self.assertIn("flow:status:00000002", subjects)
# Everything from request 1 should have been deleted.
self.assertNotIn("flow:request:00000001", subjects)
self.assertNotIn("flow:status:00000001", subjects)
# The notification for request 2 should have survived.
with queue_manager.QueueManager(token=self.token) as manager:
notifications = manager.GetNotifications(queues.FLOWS)
self.assertEqual(len(notifications), 1)
notification = notifications[0]
self.assertEqual(notification.session_id, session_id)
self.assertEqual(notification.timestamp,
flow_manager.frozen_timestamp)
self.assertEqual(RESULTS, ["Response 1"])
# The last missing piece of request 2 is the actual status message.
flow_manager.QueueResponse(session_id, status)
flow_manager.Flush()
# Now make sure request 2 runs as expected.
worker_obj.RunOnce()
worker_obj.thread_pool.Join()
self.assertEqual(RESULTS, ["Response 1", "Response 2"])
def testProcessMessages(self):
"""Test processing of several inbound messages."""
# Create a couple of flows
flow_obj = self.FlowSetup("WorkerSendingTestFlow")
session_id_1 = flow_obj.session_id
flow_obj.Close()
flow_obj = self.FlowSetup("WorkerSendingTestFlow2")
session_id_2 = flow_obj.session_id
flow_obj.Close()
manager = queue_manager.QueueManager(token=self.token)
# Check that client queue has messages
tasks_on_client_queue = manager.Query(self.client_id.Queue(), 100)
# should have 10 requests from WorkerSendingTestFlow and 1 from
# SendingTestFlow2
self.assertEqual(len(tasks_on_client_queue), 11)
# Send each of the flows a repeated message
self.SendResponse(session_id_1, "Hello1")
self.SendResponse(session_id_2, "Hello2")
self.SendResponse(session_id_1, "Hello1")
self.SendResponse(session_id_2, "Hello2")
worker_obj = worker.GRRWorker(token=self.token)
# Process all messages
worker_obj.RunOnce()
worker_obj.thread_pool.Join()
# Ensure both requests ran exactly once
RESULTS.sort()
self.assertEqual(2, len(RESULTS))
self.assertEqual("Hello1", RESULTS[0])
self.assertEqual("Hello2", RESULTS[1])
# Check that client queue is cleared - should have 2 less messages (since
# two were completed).
tasks_on_client_queue = manager.Query(self.client_id.Queue(), 100)
self.assertEqual(len(tasks_on_client_queue), 9)
# Ensure that processed requests are removed from state subject
self.assertEqual(
(None, 0),
data_store.DB.Resolve(session_id_1.Add("state"),
manager.FLOW_REQUEST_TEMPLATE % 1,
token=self.token))
# This flow is still in state Incoming.
flow_obj = aff4.FACTORY.Open(session_id_1, token=self.token)
self.assertTrue(
flow_obj.context.state != rdf_flows.FlowContext.State.TERMINATED)
self.assertEqual(flow_obj.context.current_state, "Incoming")
# This flow should be done.
flow_obj = aff4.FACTORY.Open(session_id_2, token=self.token)
self.assertTrue(
flow_obj.context.state == rdf_flows.FlowContext.State.TERMINATED)
self.assertEqual(flow_obj.context.current_state, "End")
