def testUsesFrozenTimestampWhenDeletingAndFetchingNotifications(self):
# When used in "with" statement QueueManager uses the frozen timestamp
# when fetching and deleting data. Test that if we have 2 managers
# created at different times, they will behave correctly when dealing
# with notifications for the same session ids. I.e. older queue_manager
# will only "see" it's own notification and younger queue_manager will
# "see" both.
with queue_manager.QueueManager(token=self.token) as manager1:
manager1.QueueNotification(session_id=rdfvalue.SessionID(
base="aff4:/hunts", queue=queues.HUNTS, flow_name="123456"))
manager1.Flush()
self._current_mock_time += 10
with queue_manager.QueueManager(token=self.token) as manager2:
manager2.QueueNotification(session_id=rdfvalue.SessionID(
base="aff4:/hunts", queue=queues.HUNTS, flow_name="123456"))
manager2.Flush()
self.assertEqual(
len(manager1.GetNotificationsForAllShards(queues.HUNTS)), 1)
self.assertEqual(
len(manager2.GetNotificationsForAllShards(queues.HUNTS)), 1)
manager1.DeleteNotification(
rdfvalue.SessionID(
base="aff4:/hunts", queue=queues.HUNTS, flow_name="123456"))
self.assertEqual(
len(manager1.GetNotificationsForAllShards(queues.HUNTS)), 0)
self.assertEqual(
len(manager2.GetNotificationsForAllShards(queues.HUNTS)), 1)
def testDeleteRequest(self):
"""Check that we can efficiently destroy a single flow request."""
session_id = rdfvalue.SessionID(flow_name="test3")
request = rdf_flows.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(request)
manager.QueueResponse(
rdf_flows.GrrMessage(session_id=session_id,
request_id=1,
response_id=1))
# Check the request and responses are there.
all_requests = list(manager.FetchRequestsAndResponses(session_id))
self.assertEqual(len(all_requests), 1)
self.assertEqual(all_requests[0][0], request)
with queue_manager.QueueManager(token=self.token) as manager:
manager.DeleteRequest(request)
all_requests = list(manager.FetchRequestsAndResponses(session_id))
self.assertEqual(len(all_requests), 0)
def Stop(self, reason=None):
super(GenericHunt, self).Stop(reason=reason)
started_flows = grr_collections.RDFUrnCollection(
self.started_flows_collection_urn)
num_terminated_flows = 0
self.Log("Hunt stop. Terminating all the started flows.")
# Delete hunt flows states.
for flows_batch in utils.Grouper(started_flows,
self.__class__.STOP_BATCH_SIZE):
with queue_manager.QueueManager(token=self.token) as manager:
manager.MultiDestroyFlowStates(flows_batch)
with data_store.DB.GetMutationPool() as mutation_pool:
for f in flows_batch:
flow.GRRFlow.MarkForTermination(
f,
reason="Parent hunt stopped.",
mutation_pool=mutation_pool)
num_terminated_flows += len(flows_batch)
# Delete hunt's requests and responses to ensure no more
# processing is going to occur.
with queue_manager.QueueManager(token=self.token) as manager:
manager.DestroyFlowStates(self.session_id)
self.Log("%d flows terminated.", num_terminated_flows)
def testNotificationRequeueing(self):
with test_lib.ConfigOverrider({"Worker.queue_shards": 1}):
session_id = rdfvalue.SessionID(base="aff4:/testflows",
queue=queues.HUNTS,
flow_name="123")
with test_lib.FakeTime(1000):
# Schedule a notification.
with queue_manager.QueueManager(token=self.token) as manager:
manager.QueueNotification(session_id=session_id)
with test_lib.FakeTime(1100):
with queue_manager.QueueManager(token=self.token) as manager:
notifications = manager.GetNotifications(queues.HUNTS)
self.assertEqual(len(notifications), 1)
# This notification was first queued and last queued at time 1000.
notification = notifications[0]
self.assertEqual(
notification.timestamp.AsSecondsFromEpoch(), 1000)
self.assertEqual(
notification.first_queued.AsSecondsFromEpoch(), 1000)
# Now requeue the same notification.
manager.DeleteNotification(session_id)
manager.QueueNotification(notification)
with test_lib.FakeTime(1200):
with queue_manager.QueueManager(token=self.token) as manager:
notifications = manager.GetNotifications(queues.HUNTS)
self.assertEqual(len(notifications), 1)
notification = notifications[0]
# Now the last queue time is 1100, the first queue time is still 1000.
self.assertEqual(
notification.timestamp.AsSecondsFromEpoch(), 1100)
self.assertEqual(
notification.first_queued.AsSecondsFromEpoch(), 1000)
# Again requeue the same notification.
manager.DeleteNotification(session_id)
manager.QueueNotification(notification)
expired = 1000 + queue_manager.QueueManager.notification_expiry_time
with test_lib.FakeTime(expired):
with queue_manager.QueueManager(token=self.token) as manager:
notifications = manager.GetNotifications(queues.HUNTS)
self.assertEqual(len(notifications), 1)
# Again requeue the notification, this time it should be dropped.
manager.DeleteNotification(session_id)
manager.QueueNotification(notifications[0])
with queue_manager.QueueManager(token=self.token) as manager:
notifications = manager.GetNotifications(queues.HUNTS)
self.assertEqual(len(notifications), 0)
def testDestroyFlowStates(self):
"""Check that we can efficiently destroy the flow's request queues."""
session_id = rdfvalue.SessionID(flow_name="test2")
request = rdf_flows.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(request)
manager.QueueResponse(
rdf_flows.GrrMessage(request_id=1,
response_id=1,
session_id=session_id))
# Check the request and responses are there.
all_requests = list(manager.FetchRequestsAndResponses(session_id))
self.assertEqual(len(all_requests), 1)
self.assertEqual(all_requests[0][0], request)
# Read the response directly.
responses = data_store.DB.ReadResponsesForRequestId(session_id, 1)
self.assertEqual(len(responses), 1)
response = responses[0]
self.assertEqual(response.request_id, 1)
self.assertEqual(response.response_id, 1)
self.assertEqual(response.session_id, session_id)
with queue_manager.QueueManager(token=self.token) as manager:
manager.DestroyFlowStates(session_id)
all_requests = list(manager.FetchRequestsAndResponses(session_id))
self.assertEqual(len(all_requests), 0)
# Check that the response is gone.
responses = data_store.DB.ReadResponsesForRequestId(session_id, 1)
self.assertEqual(len(responses), 0)
# Ensure the rows are gone from the data store. Some data stores
# don't store the queues in that way but there is no harm in
# checking.
self.assertEqual(
data_store.DB.ResolveRow(session_id.Add("state/request:00000001"),
token=self.token), [])
self.assertEqual(
data_store.DB.ResolveRow(session_id.Add("state"),
token=self.token), [])
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(token=self.token) 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 testReceiveUnsolicitedClientMessage(self):
client_id = test_lib.TEST_CLIENT_ID
flow_obj = self.FlowSetup(
flow_test_lib.FlowOrderTest.__name__, client_id=client_id)
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(client_id, messages)
manager = queue_manager.QueueManager(token=self.token)
completed = list(manager.FetchCompletedRequests(session_id))
self.assertEqual(len(completed), 1)
def FinalizeProcessCompletedRequests(self, notification):
# Delete kill notification as the flow got processed and is not
# stuck.
with queue_manager.QueueManager(token=self.token) as manager:
manager.DeleteNotification(self.session_id,
start=self.context.kill_timestamp,
end=self.context.kill_timestamp)
self.context.kill_timestamp = None
# If a flow raises in one state, the remaining states will not
# be processed. This is indistinguishable from an incomplete
# state due to missing responses / status so we need to check
# here if the flow is still running before rescheduling.
if (self.IsRunning() and notification.last_status
and (self.context.next_processed_request <=
notification.last_status)):
logging.debug("Had to reschedule a notification: %s",
notification)
# We have received a notification for a specific request but
# could not process that request. This might be a race
# condition in the data store so we reschedule the
# notification in the future.
delay = self.notification_retry_interval
notification.ttl -= 1
if notification.ttl:
manager.QueueNotification(
notification, timestamp=notification.timestamp + delay)
def testReceiveMessages(self):
"""Test Receiving messages with no status."""
client_id = test_lib.TEST_CLIENT_ID
flow_obj = self.FlowSetup(
flow_test_lib.FlowOrderTest.__name__, client_id=client_id)
session_id = flow_obj.session_id
messages = [
rdf_flows.GrrMessage(
request_id=1,
response_id=i,
session_id=session_id,
payload=rdfvalue.RDFInteger(i)) for i in range(1, 10)
]
self.server.ReceiveMessages(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(client_id, 100)
self.assertEqual(len(tasks_on_client_queue), 1)
stored_messages = data_store.DB.ReadResponsesForRequestId(session_id, 1)
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):
# We don't care about the last queueing time.
stored_message.timestamp = None
self.assertRDFValuesEqual(stored_message, message)
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
# TODO(user): The request_id field should be an int.
api_request = ApiFlowRequest(
request_id=str(request.id), request_state=request)
if responses:
api_request.responses = responses
result.items.append(api_request)
return result
def testKillNotificationsScheduledForFlows(self):
worker_obj = worker.GRRWorker(token=self.token)
initial_time = rdfvalue.RDFDatetime().FromSecondsFromEpoch(100)
try:
with test_lib.FakeTime(initial_time.AsSecondsFromEpoch()):
flow.GRRFlow.StartFlow(
flow_name=WorkerStuckableTestFlow.__name__,
client_id=self.client_id,
token=self.token,
sync=False)
# Process all messages
worker_obj.RunOnce()
# Wait until worker thread starts processing the flow.
WorkerStuckableTestFlow.WaitUntilWorkerStartsProcessing()
# Assert that there are no stuck notifications in the worker's
# queue.
with queue_manager.QueueManager(token=self.token) as manager:
for queue in worker_obj.queues:
notifications = manager.GetNotificationsByPriority(queue)
self.assertFalse(manager.STUCK_PRIORITY in notifications)
finally:
# Release the semaphore so that worker thread unblocks and finishes
# processing the flow.
WorkerStuckableTestFlow.LetWorkerFinishProcessing()
worker_obj.thread_pool.Join()
def testNotificationsAreDeletedFromAllShards(self):
manager = queue_manager.QueueManager(token=self.token)
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()
# There should be two notifications in two different shards.
shards_with_data = 0
for _ in range(manager.num_notification_shards):
shard_sessions = manager.GetNotifications(queues.HUNTS)
if shard_sessions:
shards_with_data += 1
self.assertEqual(len(shard_sessions), 1)
self.assertEqual(shards_with_data, 2)
# This should still work, as we delete notifications from all shards.
manager.DeleteNotification(
rdfvalue.SessionID(base="aff4:/hunts",
queue=queues.HUNTS,
flow_name="43"))
manager.DeleteNotification(
rdfvalue.SessionID(base="aff4:/hunts",
queue=queues.HUNTS,
flow_name="42"))
for _ in range(manager.num_notification_shards):
shard_sessions = manager.GetNotifications(queues.HUNTS)
self.assertFalse(shard_sessions)
def Run(self):
client_ids = self.SetupClients(1)
client_id = client_ids[0]
replace = {}
with test_lib.FakeTime(42):
flow_urn = flow.GRRFlow.StartFlow(
client_id=client_id,
flow_name=processes.ListProcesses.__name__,
token=self.token)
replace[flow_urn.Basename()] = "F:123456"
# Here we emulate a mock client with no actions (None) that should produce
# an error.
mock = flow_test_lib.MockClient(client_id, None, token=self.token)
while mock.Next():
pass
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)] = "43"
self.Check("ListClientActionRequests",
args=client_plugin.ApiListClientActionRequestsArgs(
client_id=client_id.Basename()),
replace=replace)
self.Check("ListClientActionRequests",
args=client_plugin.ApiListClientActionRequestsArgs(
client_id=client_id.Basename(), fetch_responses=True),
replace=replace)
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() 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
with data_store.DB.GetMutationPool() as pool:
manager.Delete(test_queue, tasks, mutation_pool=pool)
# Task is now deleted.
tasks = manager.QueryAndOwn(test_queue, lease_seconds=100, limit=100)
self.assertEqual(len(tasks), 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 testCountsActualNumberOfCompletedResponsesWhenApplyingTheLimit(self):
session_id = rdfvalue.SessionID(flow_name="test")
# Now queue more requests and responses:
with queue_manager.QueueManager(token=self.token) as manager:
# Start with request 1 - leave request 1 un-responded to.
for request_id in range(5):
request = rdf_flows.RequestState(id=request_id,
client_id=self.client_id,
next_state="TestState",
session_id=session_id)
manager.QueueRequest(request)
# Don't queue any actual responses, just a status message with a
# fake response_id.
manager.QueueResponse(
rdf_flows.GrrMessage(
session_id=session_id,
request_id=request_id,
response_id=1000,
type=rdf_flows.GrrMessage.Type.STATUS))
# Check that even though status message for every request indicates 1000
# responses, only the actual response count is used to apply the limit
# when FetchCompletedResponses is called.
completed_response = list(
manager.FetchCompletedResponses(session_id, limit=5))
self.assertEqual(len(completed_response), 5)
for i, (request, responses) in enumerate(completed_response):
self.assertEqual(request.id, i)
# Responses contain just the status message.
self.assertEqual(len(responses), 1)
def testTaskRetransmissionsAreCorrectlyAccounted(self):
test_queue = rdfvalue.RDFURN("fooSchedule")
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() 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)
self.assertEqual(tasks[0].task_ttl, 4)
self.assertEqual(
stats.STATS.GetMetricValue("grr_task_retransmission_count"),
self.retransmission_metric_value)
# Get a lease on the task 100 seconds later
self._current_mock_time += 110
tasks = manager.QueryAndOwn(test_queue, lease_seconds=100, limit=100)
self.assertEqual(len(tasks), 1)
self.assertEqual(tasks[0].task_ttl, 3)
self.assertEqual(
stats.STATS.GetMetricValue("grr_task_retransmission_count"),
self.retransmission_metric_value + 1)
def testSchedule(self):
"""Test the ability to schedule a task."""
test_queue = rdfvalue.RDFURN("fooSchedule")
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() as pool:
manager.Schedule([task], pool)
self.assertGreater(task.task_id, 0)
self.assertGreater(task.task_id & 0xffffffff, 0)
self.assertEqual(
(long(self._current_mock_time * 1000) & 0xffffffff) << 32,
task.task_id
& 0x1fffffff00000000)
self.assertEqual(task.task_ttl, 5)
stored_tasks = data_store.DB.QueueQueryTasks(test_queue, limit=100000)
self.assertEqual(len(stored_tasks), 1)
stored_task = stored_tasks[0]
self.assertGreater(stored_task.eta, 0)
stored_task.eta = None
self.assertRDFValuesEqual(stored_task, task)
# 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].task_ttl, 4)
self.assertEqual(tasks[0].session_id, "aff4:/Test")
# If we try to get another lease on it we should fail
self._current_mock_time += 10
tasks = manager.QueryAndOwn(test_queue, lease_seconds=100, limit=100)
self.assertEqual(len(tasks), 0)
# However after 100 seconds this should work again
self._current_mock_time += 110
tasks = manager.QueryAndOwn(test_queue, lease_seconds=100, limit=100)
self.assertEqual(len(tasks), 1)
self.assertEqual(tasks[0].task_ttl, 3)
# Check now that after a few retransmits we drop the message
for i in range(2, 0, -1):
self._current_mock_time += 110
tasks = manager.QueryAndOwn(test_queue, lease_seconds=100)
self.assertEqual(len(tasks), 1)
self.assertEqual(tasks[0].task_ttl, i)
# The task is now gone
self._current_mock_time += 110
tasks = manager.QueryAndOwn(test_queue, lease_seconds=100)
self.assertEqual(len(tasks), 0)
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)
test_process = client_test_lib.MockWindowsProcess(name="test_process")
with utils.Stubber(psutil, "Process", lambda: test_process):
mock = flow_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(
"ListFlowRequests",
args=flow_plugin.ApiListFlowRequestsArgs(
client_id=self.client_id.Basename(), flow_id=flow_urn.Basename()),
replace=replace)
def testNoKillNotificationsScheduledForHunts(self):
worker_obj = worker.GRRWorker(token=self.token)
initial_time = rdfvalue.RDFDatetime().FromSecondsFromEpoch(100)
try:
with test_lib.FakeTime(initial_time.AsSecondsFromEpoch()):
with implementation.GRRHunt.StartHunt(
hunt_name=WorkerStuckableHunt.__name__,
client_rate=0,
token=self.token) as hunt:
hunt.GetRunner().Start()
implementation.GRRHunt.StartClients(hunt.session_id, [self.client_id])
# Process all messages
while worker_obj.RunOnce():
pass
# Wait until worker thread starts processing the flow.
WorkerStuckableHunt.WaitUntilWorkerStartsProcessing()
# Assert that there are no stuck notifications in the worker's queue.
with queue_manager.QueueManager(token=self.token) as manager:
for queue in worker_obj.queues:
notifications = manager.GetNotificationsByPriority(queue)
self.assertFalse(manager.STUCK_PRIORITY in notifications)
finally:
# Release the semaphore so that worker thread unblocks and finishes
# processing the flow.
WorkerStuckableHunt.LetWorkerFinishProcessing()
worker_obj.thread_pool.Join()
def CheckNotificationsDisappear(self, session_id):
worker_obj = worker.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 testDestroyFlowStates(self):
"""Check that we can efficiently destroy the flow's request queues."""
session_id = rdfvalue.SessionID(flow_name="test2")
request = rdf_flows.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(request)
manager.QueueResponse(
rdf_flows.GrrMessage(request_id=1,
response_id=1,
session_id=session_id))
# Check the request and responses are there.
all_requests = list(manager.FetchRequestsAndResponses(session_id))
self.assertEqual(len(all_requests), 1)
self.assertEqual(all_requests[0][0], request)
# Ensure the rows are in the data store:
self.assertEqual(
data_store.DB.ResolveRow(session_id.Add("state"),
token=self.token)[0][0],
"flow:request:00000001")
self.assertEqual(
data_store.DB.ResolveRow(session_id.Add("state/request:00000001"),
token=self.token)[0][0],
"flow:response:00000001:00000001")
with queue_manager.QueueManager(token=self.token) as manager:
manager.DestroyFlowStates(session_id)
all_requests = list(manager.FetchRequestsAndResponses(session_id))
self.assertEqual(len(all_requests), 0)
# Ensure the rows are gone from the data store.
self.assertEqual(
data_store.DB.ResolveRow(session_id.Add("state/request:00000001"),
token=self.token), [])
self.assertEqual(
data_store.DB.ResolveRow(session_id.Add("state"),
token=self.token), [])
def CreateLeasedClientRequest(
client_id=rdf_client.ClientURN("C.0000000000000001"), token=None):
flow.GRRFlow.StartFlow(
client_id=client_id,
flow_name=processes.ListProcesses.__name__,
token=token)
with queue_manager.QueueManager(token=token) as manager:
manager.QueryAndOwn(client_id.Queue(), limit=1, lease_seconds=10000)
def testFirstShardNameIsEqualToTheQueue(self):
manager = queue_manager.QueueManager(token=self.token)
while True:
shard = manager.GetNotificationShard(queues.HUNTS)
if (manager.notification_shard_counters[str(queues.HUNTS)] %
manager.num_notification_shards) == 0:
break
self.assertEqual(shard, queues.HUNTS)
def testNotFirstShardNameHasIndexSuffix(self):
manager = queue_manager.QueueManager(token=self.token)
while True:
shard = manager.GetNotificationShard(queues.HUNTS)
if (manager.notification_shard_counters[str(queues.HUNTS)] %
manager.num_notification_shards) == 1:
break
self.assertEqual(shard, queues.HUNTS.Add("1"))
def HandleMessageBundles(self, request_comms, response_comms):
"""Processes a queue of messages as passed from the client.
We basically dispatch all the GrrMessages in the queue to the task scheduler
for backend processing. We then retrieve from the TS the messages destined
for this client.
Args:
request_comms: A ClientCommunication rdfvalue with messages sent by the
client. source should be set to the client CN.
response_comms: A ClientCommunication rdfvalue of jobs destined to this
client.
Returns:
tuple of (source, message_count) where message_count is the number of
messages received from the client with common name source.
"""
messages, source, timestamp = self._communicator.DecodeMessages(
request_comms)
now = time.time()
if messages:
# Receive messages in line.
self.ReceiveMessages(source, messages)
# We send the client a maximum of self.max_queue_size messages
required_count = max(0, self.max_queue_size - request_comms.queue_size)
tasks = []
message_list = rdf_flows.MessageList()
# Only give the client messages if we are able to receive them in a
# reasonable time.
if time.time() - now < 10:
tasks = self.DrainTaskSchedulerQueueForClient(
source, required_count)
message_list.job = tasks
# Encode the message_list in the response_comms using the same API version
# the client used.
try:
self._communicator.EncodeMessages(
message_list,
response_comms,
destination=str(source),
timestamp=timestamp,
api_version=request_comms.api_version)
except communicator.UnknownClientCert:
# We can not encode messages to the client yet because we do not have the
# client certificate - return them to the queue so we can try again later.
with data_store.DB.GetMutationPool() as pool:
queue_manager.QueueManager(token=self.token).Schedule(
tasks, pool)
raise
return source, len(messages)
def DeleteJobFlows(self, age=None):
"""Deletes flows initiated by the job that are older than specified."""
if age is None:
raise ValueError("age can't be None")
child_flows = list(self.ListChildren(age=age))
with queue_manager.QueueManager(token=self.token) as queuemanager:
queuemanager.MultiDestroyFlowStates(child_flows)
aff4.FACTORY.MultiDelete(child_flows, token=self.token)
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.queue_manager.FreezeTimestamp()
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 = []
def testHandleMessageBundle(self):
"""Check that HandleMessageBundles() requeues messages if it failed.
This test makes sure that when messages are pending for a client, and which
we have no certificate for, the messages are requeued when sending fails.
"""
# Make a new fake client
client_id = self.SetupClient(0)
class MockCommunicator(object):
"""A fake that simulates an unenrolled client."""
def DecodeMessages(self, *unused_args):
"""For simplicity client sends an empty request."""
return ([], client_id, 100)
def EncodeMessages(self, *unused_args, **unused_kw):
"""Raise because the server has no certificates for this client."""
raise communicator.UnknownClientCert()
# Install the mock.
self.server._communicator = MockCommunicator()
# First request, the server will raise UnknownClientCert.
request_comms = rdf_flows.ClientCommunication()
self.assertRaises(communicator.UnknownClientCert,
self.server.HandleMessageBundles, request_comms, 2)
# We can still schedule a flow for it
flow.GRRFlow.StartFlow(
client_id=client_id,
flow_name=flow_test_lib.SendingFlow.__name__,
message_count=1,
token=self.token)
manager = queue_manager.QueueManager(token=self.token)
tasks = manager.Query(client_id, limit=100)
self.assertRaises(communicator.UnknownClientCert,
self.server.HandleMessageBundles, request_comms, 2)
new_tasks = manager.Query(client_id, limit=100)
# The different in eta times reflect the lease that the server took on the
# client messages.
lease_time = (new_tasks[0].eta - tasks[0].eta) / 1e6
# This lease time must be small, as the HandleMessageBundles() call failed,
# the pending client messages must be put back on the queue.
self.assertLess(lease_time, 1)
# Since the server tried to send it, the ttl must be decremented
self.assertEqual(tasks[0].task_ttl - new_tasks[0].task_ttl, 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 testEqualTimestampNotifications(self):
frontend_server = front_end.FrontEndServer(
certificate=config.CONFIG["Frontend.certificate"],
private_key=config.CONFIG["PrivateKeys.server_key"],
message_expiry_time=100,
threadpool_prefix="notification-test")
# This schedules 10 requests.
session_id = flow.GRRFlow.StartFlow(
client_id=self.client_id,
flow_name="WorkerSendingTestFlow",
token=self.token)
# We pretend that the client processed all the 10 requests at once and
# sends the replies in a single http poll.
messages = [
rdf_flows.GrrMessage(
request_id=i,
response_id=1,
session_id=session_id,
payload=rdf_protodict.DataBlob(string="test%s" % i),
generate_task_id=True) for i in range(1, 11)
]
status = rdf_flows.GrrStatus(status=rdf_flows.GrrStatus.ReturnedStatus.OK)
statuses = [
rdf_flows.GrrMessage(
request_id=i,
response_id=2,
session_id=session_id,
payload=status,
type=rdf_flows.GrrMessage.Type.STATUS,
generate_task_id=True) for i in range(1, 11)
]
frontend_server.ReceiveMessages(self.client_id, messages + statuses)
with queue_manager.QueueManager(token=self.token) as q:
all_notifications = q.GetNotificationsByPriorityForAllShards(
rdfvalue.RDFURN("aff4:/F"))
medium_priority = rdf_flows.GrrNotification.Priority.MEDIUM_PRIORITY
medium_notifications = all_notifications[medium_priority]
my_notifications = [
n for n in medium_notifications if n.session_id == session_id
]
# There must not be more than one notification.
self.assertEqual(len(my_notifications), 1)
notification = my_notifications[0]
self.assertEqual(notification.first_queued, notification.timestamp)
self.assertEqual(notification.last_status, 10)
