def inject_messages(*args, **kwargs):
# A control plane message
writer.put(
Message("SESSION_CREATE_REQUEST", "plugin_fuz", None, None,
None))
low_body = DevicePluginMessage('low', PRIO_LOW)
normal_body = DevicePluginMessage('normal', PRIO_NORMAL)
high_body = DevicePluginMessage('high', PRIO_HIGH)
writer.put(Message("DATA", "plugin_foo", low_body, "foo", 0))
writer.put(Message("DATA", "plugin_bar", normal_body, "foo", 1))
writer.put(Message("DATA", "plugin_baz", high_body, "foo", 2))
def poll(self, plugin_name):
"""
For any plugins that don't have a session, try asking for one.
For any ongoing sessions, invoke the poll callback
"""
now = datetime.datetime.now()
try:
session = self._client.sessions.get(plugin_name)
except KeyError:
# Request to open a session
#
if plugin_name in self._client.sessions._requested_at:
next_request_at = (
self._client.sessions._requested_at[plugin_name] +
self._client.sessions._backoffs[plugin_name])
if now < next_request_at:
# We're still in our backoff period, skip requesting a session
daemon_log.debug("Delaying session request until %s" %
next_request_at)
return
else:
if (self._client.sessions._backoffs[plugin_name] <
MAX_SESSION_BACKOFF):
self._client.sessions._backoffs[plugin_name] *= 2
daemon_log.debug("Requesting session for plugin %s" % plugin_name)
self._client.sessions._requested_at[plugin_name] = now
self.put(Message("SESSION_CREATE_REQUEST", plugin_name))
else:
try:
data = session.poll()
except Exception:
backtrace = "\n".join(
traceback.format_exception(*(sys.exc_info())))
daemon_log.error("Error in plugin %s: %s" %
(plugin_name, backtrace))
self._client.sessions.terminate(plugin_name)
self.put(Message("SESSION_CREATE_REQUEST", plugin_name))
else:
if data is not None:
if isinstance(data, DevicePluginMessageCollection):
for message in data:
session.send_message(
DevicePluginMessage(message,
priority=data.priority))
elif isinstance(data, DevicePluginMessage):
session.send_message(data)
else:
session.send_message(DevicePluginMessage(data))
def test_oversized_messages(self):
"""
Test that oversized messages are dropped and the session is terminated
"""
# Monkey-patch this setting to a lower limit to make testing easier
MAX_BYTES_PER_POST = 1024
from chroma_agent.agent_client import MAX_BYTES_PER_POST as LARGE_MAX_BYTES_PER_POST
def set_post_limit(size):
import chroma_agent.agent_client
chroma_agent.agent_client.MAX_BYTES_PER_POST = size
self.addCleanup(set_post_limit, LARGE_MAX_BYTES_PER_POST)
set_post_limit(MAX_BYTES_PER_POST)
client = mock.Mock()
client._fqdn = "test_server"
client.boot_time = IMLDateTime.utcnow()
client.start_time = IMLDateTime.utcnow()
writer = HttpWriter(client)
def fake_post(envelope):
if len(json.dumps(envelope)) > MAX_BYTES_PER_POST:
daemon_log.info("fake_post(): rejecting oversized message")
raise HttpError()
client.post = mock.Mock(side_effect=fake_post)
TestPlugin = mock.Mock()
mock_plugin_instance = mock.Mock()
mock_plugin_instance.start_session = mock.Mock(
return_value={'foo': 'bar'})
client.device_plugins.get = mock.Mock(
return_value=lambda (plugin_name): mock_plugin_instance)
client.device_plugins.get_plugins = mock.Mock(
return_value={'test_plugin': TestPlugin})
client.sessions = SessionTable(client)
daemon_log.setLevel(logging.DEBUG)
import string
from random import choice
oversized_string = "".join(
choice(string.printable) for i in range(MAX_BYTES_PER_POST))
# There should be one message to set up the session
writer.poll('test_plugin')
self.assertTrue(writer.send())
self.assertEqual(client.post.call_count, 1)
messages = client.post.call_args[0][0]['messages']
self.assertEqual(len(messages), 1)
self.assertEqual(messages[0]['type'], "SESSION_CREATE_REQUEST")
# Pretend we got a SESSION_CREATE_RESPONSE
client.sessions.create('test_plugin', 'id_foo')
self.assertEqual(len(client.sessions._sessions), 1)
# Inject a normal and an oversized message
normal_body = DevicePluginMessage('normal', PRIO_NORMAL)
oversized_body = DevicePluginMessage(oversized_string, PRIO_NORMAL)
writer.put(Message("DATA", "test_plugin", normal_body, "id_foo", 0))
writer.put(Message("DATA", "test_plugin", oversized_body, "id_foo", 1))
# Only the normal message should get through
self.assertTrue(writer.send())
self.assertEqual(client.post.call_count, 2)
messages = client.post.call_args[0][0]['messages']
self.assertEqual(len(messages), 1)
self.assertEqual(messages[0]['type'], "DATA")
# The oversized message should be dropped and the session
# terminated
self.assertFalse(writer.send())
self.assertEqual(client.post.call_count, 3)
self.assertEqual(len(client.sessions._sessions), 0)
# However, we should eventually get a new session for the
# offending plugin
writer.poll('test_plugin')
self.assertTrue(writer.send())
self.assertEqual(client.post.call_count, 4)
messages = client.post.call_args[0][0]['messages']
self.assertEqual(len(messages), 1)
self.assertEqual(messages[0]['type'], "SESSION_CREATE_REQUEST")