def __init__(self, server_address, handler, frontend=None, *args, **kwargs):
stats.STATS.SetGaugeValue("frontend_max_active_count",
self.request_queue_size)
if frontend:
self.frontend = frontend
else:
self.frontend = flow.FrontEndServer(
certificate=config_lib.CONFIG["Frontend.certificate"],
private_key=config_lib.CONFIG["PrivateKeys.server_key"],
max_queue_size=config_lib.CONFIG["Frontend.max_queue_size"],
message_expiry_time=config_lib.CONFIG["Frontend.message_expiry_time"],
max_retransmission_time=config_lib.CONFIG[
"Frontend.max_retransmission_time"])
self.server_cert = config_lib.CONFIG["Frontend.certificate"]
(address, _) = server_address
version = ipaddr.IPAddress(address).version
if version == 4:
self.address_family = socket.AF_INET
elif version == 6:
self.address_family = socket.AF_INET6
logging.info("Will attempt to listen on %s", server_address)
BaseHTTPServer.HTTPServer.__init__(self, server_address, handler, *args,
**kwargs)
def __init__(self):
startup.Init()
self.server_pem = config_lib.CONFIG["Frontend.certificate"]
self.front_end = flow.FrontEndServer(
certificate=config_lib.CONFIG["Frontend.certificate"],
private_key=config_lib.CONFIG["PrivateKeys.server_key"],
max_queue_size=config_lib.CONFIG["Frontend.max_queue_size"],
message_expiry_time=config_lib.CONFIG["Frontend.message_expiry_time"],
max_retransmission_time=config_lib.CONFIG[
"Frontend.max_retransmission_time"])
def setUp(self):
"""Setup the server."""
super(GRRFEServerTest, self).setUp()
# Whitelist test flow.
config_lib.CONFIG.Set("Frontend.well_known_flows", [utils.SmartStr(
test_lib.WellKnownSessionTest.well_known_session_id)])
# For tests, small pools are ok.
config_lib.CONFIG.Set("Threadpool.size", 10)
prefix = "pool-%s" % self._testMethodName
self.server = flow.FrontEndServer(
certificate=config_lib.CONFIG["Frontend.certificate"],
private_key=config_lib.CONFIG["PrivateKeys.server_key"],
message_expiry_time=self.message_expiry_time,
threadpool_prefix=prefix)
def testEqualTimestampNotifications(self):
frontend_server = flow.FrontEndServer(
certificate=config_lib.CONFIG["Frontend.certificate"],
private_key=config_lib.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))
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)
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)
def setUp(self):
"""Setup the server."""
super(GRRFEServerTest, self).setUp()
self.config_overrider = test_lib.ConfigOverrider({
# Whitelist test flow.
"Frontend.well_known_flows": [
utils.SmartStr(test_lib.WellKnownSessionTest.
well_known_session_id.FlowName())
],
# For tests, small pools are ok.
"Threadpool.size":
10
})
self.config_overrider.Start()
prefix = "pool-%s" % self._testMethodName
self.server = flow.FrontEndServer(
certificate=config_lib.CONFIG["Frontend.certificate"],
private_key=config_lib.CONFIG["PrivateKeys.server_key"],
message_expiry_time=self.message_expiry_time,
threadpool_prefix=prefix)
