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=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 _SendMessages(self, grr_msgs, background=False):
"""Sends a block of messages through Fleetspeak."""
message_list = rdf_flows.PackedMessageList()
communicator.Communicator.EncodeMessageList(
rdf_flows.MessageList(job=grr_msgs), message_list)
fs_msg = fs_common_pb2.Message(
message_type="MessageList",
destination=fs_common_pb2.Address(service_name="GRR"),
background=background)
fs_msg.data.Pack(message_list.AsPrimitiveProto())
try:
sent_bytes = self._fs.Send(fs_msg)
except (IOError, struct.error) as e:
logging.critical("Broken local Fleetspeak connection (write end).")
raise e
stats.STATS.IncrementCounter("grr_client_sent_bytes", sent_bytes)
def testReceiveMessageListFleetspeak(self):
service_name = "GRR"
fake_service_client = _FakeGRPCServiceClient(service_name)
fleetspeak_connector.Reset()
fleetspeak_connector.Init(service_client=fake_service_client)
fsd = fs_frontend_tool.GRRFSServer()
grr_client_nr = 0xab
grr_client_id_urn = self.SetupClient(grr_client_nr)
flow_obj = self.FlowSetup(flow_test_lib.FlowOrderTest.__name__,
grr_client_id_urn)
num_msgs = 9
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 xrange(1, num_msgs + 1)
]
fs_client_id = "\x10\x00\x00\x00\x00\x00\x00\xab"
# fs_client_id should be equivalent to grr_client_id_urn
self.assertEqual(
fs_client_id,
fleetspeak_utils.GRRIDToFleetspeakID(grr_client_id_urn.Basename()))
message_list = rdf_flows.PackedMessageList()
communicator.Communicator.EncodeMessageList(
rdf_flows.MessageList(job=messages), message_list)
fs_message = fs_common_pb2.Message(
message_type="MessageList",
source=fs_common_pb2.Address(
client_id=fs_client_id, service_name=service_name))
fs_message.data.Pack(message_list.AsPrimitiveProto())
fsd.Process(fs_message, None)
# Make sure the task is still on the client queue
manager = queue_manager.QueueManager(token=self.token)
tasks_on_client_queue = manager.Query(grr_client_id_urn, 100)
self.assertEqual(len(tasks_on_client_queue), 1)
want_messages = [message.Copy() for message in messages]
for want_message in want_messages:
# This is filled in by the frontend as soon as it gets the message.
want_message.auth_state = (
rdf_flows.GrrMessage.AuthorizationState.AUTHENTICATED)
want_message.source = grr_client_id_urn
stored_messages = data_store.DB.ReadResponsesForRequestId(session_id, 1)
self.assertEqual(len(stored_messages), len(want_messages))
stored_messages.sort(key=lambda m: m.response_id)
# Check that messages were stored correctly
for stored_message, want_message in itertools.izip(stored_messages,
want_messages):
stored_message.timestamp = None
self.assertRDFValuesEqual(stored_message, want_message)