qemu

protocol.py (18678B)

---

 import asyncio
 from contextlib import contextmanager
 import os
 import socket
 from tempfile import TemporaryDirectory
 
 import avocado
 
 from qemu.qmp import ConnectError, Runstate
 from qemu.qmp.protocol import AsyncProtocol, StateError
 from qemu.qmp.util import asyncio_run, create_task
 
 
 class NullProtocol(AsyncProtocol[None]):
     """
     NullProtocol is a test mockup of an AsyncProtocol implementation.
 
     It adds a fake_session instance variable that enables a code path
     that bypasses the actual connection logic, but still allows the
     reader/writers to start.
 
     Because the message type is defined as None, an asyncio.Event named
     'trigger_input' is created that prohibits the reader from
     incessantly being able to yield None; this event can be poked to
     simulate an incoming message.
 
     For testing symmetry with do_recv, an interface is added to "send" a
     Null message.
 
     For testing purposes, a "simulate_disconnection" method is also
     added which allows us to trigger a bottom half disconnect without
     injecting any real errors into the reader/writer loops; in essence
     it performs exactly half of what disconnect() normally does.
     """
     def __init__(self, name=None):
         self.fake_session = False
         self.trigger_input: asyncio.Event
         super().__init__(name)
 
     async def _establish_session(self):
         self.trigger_input = asyncio.Event()
         await super()._establish_session()
 
     async def _do_start_server(self, address, ssl=None):
         if self.fake_session:
             self._accepted = asyncio.Event()
             self._set_state(Runstate.CONNECTING)
             await asyncio.sleep(0)
         else:
             await super()._do_start_server(address, ssl)
 
     async def _do_accept(self):
         if self.fake_session:
             self._accepted = None
         else:
             await super()._do_accept()
 
     async def _do_connect(self, address, ssl=None):
         if self.fake_session:
             self._set_state(Runstate.CONNECTING)
             await asyncio.sleep(0)
         else:
             await super()._do_connect(address, ssl)
 
     async def _do_recv(self) -> None:
         await self.trigger_input.wait()
         self.trigger_input.clear()
 
     def _do_send(self, msg: None) -> None:
         pass
 
     async def send_msg(self) -> None:
         await self._outgoing.put(None)
 
     async def simulate_disconnect(self) -> None:
         """
         Simulates a bottom-half disconnect.
 
         This method schedules a disconnection but does not wait for it
         to complete. This is used to put the loop into the DISCONNECTING
         state without fully quiescing it back to IDLE. This is normally
         something you cannot coax AsyncProtocol to do on purpose, but it
         will be similar to what happens with an unhandled Exception in
         the reader/writer.
 
         Under normal circumstances, the library design requires you to
         await on disconnect(), which awaits the disconnect task and
         returns bottom half errors as a pre-condition to allowing the
         loop to return back to IDLE.
         """
         self._schedule_disconnect()
 
 
 class LineProtocol(AsyncProtocol[str]):
     def __init__(self, name=None):
         super().__init__(name)
         self.rx_history = []
 
     async def _do_recv(self) -> str:
         raw = await self._readline()
         msg = raw.decode()
         self.rx_history.append(msg)
         return msg
 
     def _do_send(self, msg: str) -> None:
         assert self._writer is not None
         self._writer.write(msg.encode() + b'\n')
 
     async def send_msg(self, msg: str) -> None:
         await self._outgoing.put(msg)
 
 
 def run_as_task(coro, allow_cancellation=False):
     """
     Run a given coroutine as a task.
 
     Optionally, wrap it in a try..except block that allows this
     coroutine to be canceled gracefully.
     """
     async def _runner():
         try:
             await coro
         except asyncio.CancelledError:
             if allow_cancellation:
                 return
             raise
     return create_task(_runner())
 
 
 @contextmanager
 def jammed_socket():
     """
     Opens up a random unused TCP port on localhost, then jams it.
     """
     socks = []
 
     try:
         sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
         sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
         sock.bind(('127.0.0.1', 0))
         sock.listen(1)
         address = sock.getsockname()
 
         socks.append(sock)
 
         # I don't *fully* understand why, but it takes *two* un-accepted
         # connections to start jamming the socket.
         for _ in range(2):
             sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
             sock.connect(address)
             socks.append(sock)
 
         yield address
 
     finally:
         for sock in socks:
             sock.close()
 
 
 class Smoke(avocado.Test):
 
     def setUp(self):
         self.proto = NullProtocol()
 
     def test__repr__(self):
         self.assertEqual(
             repr(self.proto),
             "<NullProtocol runstate=IDLE>"
         )
 
     def testRunstate(self):
         self.assertEqual(
             self.proto.runstate,
             Runstate.IDLE
         )
 
     def testDefaultName(self):
         self.assertEqual(
             self.proto.name,
             None
         )
 
     def testLogger(self):
         self.assertEqual(
             self.proto.logger.name,
             'qemu.qmp.protocol'
         )
 
     def testName(self):
         self.proto = NullProtocol('Steve')
 
         self.assertEqual(
             self.proto.name,
             'Steve'
         )
 
         self.assertEqual(
             self.proto.logger.name,
             'qemu.qmp.protocol.Steve'
         )
 
         self.assertEqual(
             repr(self.proto),
             "<NullProtocol name='Steve' runstate=IDLE>"
         )
 
 
 class TestBase(avocado.Test):
 
     def setUp(self):
         self.proto = NullProtocol(type(self).__name__)
         self.assertEqual(self.proto.runstate, Runstate.IDLE)
         self.runstate_watcher = None
 
     def tearDown(self):
         self.assertEqual(self.proto.runstate, Runstate.IDLE)
 
     async def _asyncSetUp(self):
         pass
 
     async def _asyncTearDown(self):
         if self.runstate_watcher:
             await self.runstate_watcher
 
     @staticmethod
     def async_test(async_test_method):
         """
         Decorator; adds SetUp and TearDown to async tests.
         """
         async def _wrapper(self, *args, **kwargs):
             loop = asyncio.get_event_loop()
             loop.set_debug(True)
 
             await self._asyncSetUp()
             await async_test_method(self, *args, **kwargs)
             await self._asyncTearDown()
 
         return _wrapper
 
     # Definitions
 
     # The states we expect a "bad" connect/accept attempt to transition through
     BAD_CONNECTION_STATES = (
         Runstate.CONNECTING,
         Runstate.DISCONNECTING,
         Runstate.IDLE,
     )
 
     # The states we expect a "good" session to transition through
     GOOD_CONNECTION_STATES = (
         Runstate.CONNECTING,
         Runstate.RUNNING,
         Runstate.DISCONNECTING,
         Runstate.IDLE,
     )
 
     # Helpers
 
     async def _watch_runstates(self, *states):
         """
         This launches a task alongside (most) tests below to confirm that
         the sequence of runstate changes that occur is exactly as
         anticipated.
         """
         async def _watcher():
             for state in states:
                 new_state = await self.proto.runstate_changed()
                 self.assertEqual(
                     new_state,
                     state,
                     msg=f"Expected state '{state.name}'",
                 )
 
         self.runstate_watcher = create_task(_watcher())
         # Kick the loop and force the task to block on the event.
         await asyncio.sleep(0)
 
 
 class State(TestBase):
 
     @TestBase.async_test
     async def testSuperfluousDisconnect(self):
         """
         Test calling disconnect() while already disconnected.
         """
         await self._watch_runstates(
             Runstate.DISCONNECTING,
             Runstate.IDLE,
         )
         await self.proto.disconnect()
 
 
 class Connect(TestBase):
     """
     Tests primarily related to calling Connect().
     """
     async def _bad_connection(self, family: str):
         assert family in ('INET', 'UNIX')
 
         if family == 'INET':
             await self.proto.connect(('127.0.0.1', 0))
         elif family == 'UNIX':
             await self.proto.connect('/dev/null')
 
     async def _hanging_connection(self):
         with jammed_socket() as addr:
             await self.proto.connect(addr)
 
     async def _bad_connection_test(self, family: str):
         await self._watch_runstates(*self.BAD_CONNECTION_STATES)
 
         with self.assertRaises(ConnectError) as context:
             await self._bad_connection(family)
 
         self.assertIsInstance(context.exception.exc, OSError)
         self.assertEqual(
             context.exception.error_message,
             "Failed to establish connection"
         )
 
     @TestBase.async_test
     async def testBadINET(self):
         """
         Test an immediately rejected call to an IP target.
         """
         await self._bad_connection_test('INET')
 
     @TestBase.async_test
     async def testBadUNIX(self):
         """
         Test an immediately rejected call to a UNIX socket target.
         """
         await self._bad_connection_test('UNIX')
 
     @TestBase.async_test
     async def testCancellation(self):
         """
         Test what happens when a connection attempt is aborted.
         """
         # Note that accept() cannot be cancelled outright, as it isn't a task.
         # However, we can wrap it in a task and cancel *that*.
         await self._watch_runstates(*self.BAD_CONNECTION_STATES)
         task = run_as_task(self._hanging_connection(), allow_cancellation=True)
 
         state = await self.proto.runstate_changed()
         self.assertEqual(state, Runstate.CONNECTING)
 
         # This is insider baseball, but the connection attempt has
         # yielded *just* before the actual connection attempt, so kick
         # the loop to make sure it's truly wedged.
         await asyncio.sleep(0)
 
         task.cancel()
         await task
 
     @TestBase.async_test
     async def testTimeout(self):
         """
         Test what happens when a connection attempt times out.
         """
         await self._watch_runstates(*self.BAD_CONNECTION_STATES)
         task = run_as_task(self._hanging_connection())
 
         # More insider baseball: to improve the speed of this test while
         # guaranteeing that the connection even gets a chance to start,
         # verify that the connection hangs *first*, then await the
         # result of the task with a nearly-zero timeout.
 
         state = await self.proto.runstate_changed()
         self.assertEqual(state, Runstate.CONNECTING)
         await asyncio.sleep(0)
 
         with self.assertRaises(asyncio.TimeoutError):
             await asyncio.wait_for(task, timeout=0)
 
     @TestBase.async_test
     async def testRequire(self):
         """
         Test what happens when a connection attempt is made while CONNECTING.
         """
         await self._watch_runstates(*self.BAD_CONNECTION_STATES)
         task = run_as_task(self._hanging_connection(), allow_cancellation=True)
 
         state = await self.proto.runstate_changed()
         self.assertEqual(state, Runstate.CONNECTING)
 
         with self.assertRaises(StateError) as context:
             await self._bad_connection('UNIX')
 
         self.assertEqual(
             context.exception.error_message,
             "NullProtocol is currently connecting."
         )
         self.assertEqual(context.exception.state, Runstate.CONNECTING)
         self.assertEqual(context.exception.required, Runstate.IDLE)
 
         task.cancel()
         await task
 
     @TestBase.async_test
     async def testImplicitRunstateInit(self):
         """
         Test what happens if we do not wait on the runstate event until
         AFTER a connection is made, i.e., connect()/accept() themselves
         initialize the runstate event. All of the above tests force the
         initialization by waiting on the runstate *first*.
         """
         task = run_as_task(self._hanging_connection(), allow_cancellation=True)
 
         # Kick the loop to coerce the state change
         await asyncio.sleep(0)
         assert self.proto.runstate == Runstate.CONNECTING
 
         # We already missed the transition to CONNECTING
         await self._watch_runstates(Runstate.DISCONNECTING, Runstate.IDLE)
 
         task.cancel()
         await task
 
 
 class Accept(Connect):
     """
     All of the same tests as Connect, but using the accept() interface.
     """
     async def _bad_connection(self, family: str):
         assert family in ('INET', 'UNIX')
 
         if family == 'INET':
             await self.proto.start_server_and_accept(('example.com', 1))
         elif family == 'UNIX':
             await self.proto.start_server_and_accept('/dev/null')
 
     async def _hanging_connection(self):
         with TemporaryDirectory(suffix='.qmp') as tmpdir:
             sock = os.path.join(tmpdir, type(self.proto).__name__ + ".sock")
             await self.proto.start_server_and_accept(sock)
 
 
 class FakeSession(TestBase):
 
     def setUp(self):
         super().setUp()
         self.proto.fake_session = True
 
     async def _asyncSetUp(self):
         await super()._asyncSetUp()
         await self._watch_runstates(*self.GOOD_CONNECTION_STATES)
 
     async def _asyncTearDown(self):
         await self.proto.disconnect()
         await super()._asyncTearDown()
 
     ####
 
     @TestBase.async_test
     async def testFakeConnect(self):
 
         """Test the full state lifecycle (via connect) with a no-op session."""
         await self.proto.connect('/not/a/real/path')
         self.assertEqual(self.proto.runstate, Runstate.RUNNING)
 
     @TestBase.async_test
     async def testFakeAccept(self):
         """Test the full state lifecycle (via accept) with a no-op session."""
         await self.proto.start_server_and_accept('/not/a/real/path')
         self.assertEqual(self.proto.runstate, Runstate.RUNNING)
 
     @TestBase.async_test
     async def testFakeRecv(self):
         """Test receiving a fake/null message."""
         await self.proto.start_server_and_accept('/not/a/real/path')
 
         logname = self.proto.logger.name
         with self.assertLogs(logname, level='DEBUG') as context:
             self.proto.trigger_input.set()
             self.proto.trigger_input.clear()
             await asyncio.sleep(0)  # Kick reader.
 
         self.assertEqual(
             context.output,
             [f"DEBUG:{logname}:<-- None"],
         )
 
     @TestBase.async_test
     async def testFakeSend(self):
         """Test sending a fake/null message."""
         await self.proto.start_server_and_accept('/not/a/real/path')
 
         logname = self.proto.logger.name
         with self.assertLogs(logname, level='DEBUG') as context:
             # Cheat: Send a Null message to nobody.
             await self.proto.send_msg()
             # Kick writer; awaiting on a queue.put isn't sufficient to yield.
             await asyncio.sleep(0)
 
         self.assertEqual(
             context.output,
             [f"DEBUG:{logname}:--> None"],
         )
 
     async def _prod_session_api(
             self,
             current_state: Runstate,
             error_message: str,
             accept: bool = True
     ):
         with self.assertRaises(StateError) as context:
             if accept:
                 await self.proto.start_server_and_accept('/not/a/real/path')
             else:
                 await self.proto.connect('/not/a/real/path')
 
         self.assertEqual(context.exception.error_message, error_message)
         self.assertEqual(context.exception.state, current_state)
         self.assertEqual(context.exception.required, Runstate.IDLE)
 
     @TestBase.async_test
     async def testAcceptRequireRunning(self):
         """Test that accept() cannot be called when Runstate=RUNNING"""
         await self.proto.start_server_and_accept('/not/a/real/path')
 
         await self._prod_session_api(
             Runstate.RUNNING,
             "NullProtocol is already connected and running.",
             accept=True,
         )
 
     @TestBase.async_test
     async def testConnectRequireRunning(self):
         """Test that connect() cannot be called when Runstate=RUNNING"""
         await self.proto.start_server_and_accept('/not/a/real/path')
 
         await self._prod_session_api(
             Runstate.RUNNING,
             "NullProtocol is already connected and running.",
             accept=False,
         )
 
     @TestBase.async_test
     async def testAcceptRequireDisconnecting(self):
         """Test that accept() cannot be called when Runstate=DISCONNECTING"""
         await self.proto.start_server_and_accept('/not/a/real/path')
 
         # Cheat: force a disconnect.
         await self.proto.simulate_disconnect()
 
         await self._prod_session_api(
             Runstate.DISCONNECTING,
             ("NullProtocol is disconnecting."
              " Call disconnect() to return to IDLE state."),
             accept=True,
         )
 
     @TestBase.async_test
     async def testConnectRequireDisconnecting(self):
         """Test that connect() cannot be called when Runstate=DISCONNECTING"""
         await self.proto.start_server_and_accept('/not/a/real/path')
 
         # Cheat: force a disconnect.
         await self.proto.simulate_disconnect()
 
         await self._prod_session_api(
             Runstate.DISCONNECTING,
             ("NullProtocol is disconnecting."
              " Call disconnect() to return to IDLE state."),
             accept=False,
         )
 
 
 class SimpleSession(TestBase):
 
     def setUp(self):
         super().setUp()
         self.server = LineProtocol(type(self).__name__ + '-server')
 
     async def _asyncSetUp(self):
         await super()._asyncSetUp()
         await self._watch_runstates(*self.GOOD_CONNECTION_STATES)
 
     async def _asyncTearDown(self):
         await self.proto.disconnect()
         try:
             await self.server.disconnect()
         except EOFError:
             pass
         await super()._asyncTearDown()
 
     @TestBase.async_test
     async def testSmoke(self):
         with TemporaryDirectory(suffix='.qmp') as tmpdir:
             sock = os.path.join(tmpdir, type(self.proto).__name__ + ".sock")
             server_task = create_task(self.server.start_server_and_accept(sock))
 
             # give the server a chance to start listening [...]
             await asyncio.sleep(0)
             await self.proto.connect(sock)