api.py 6.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176
  1. # -*- encoding: utf8 -*-
  2. #
  3. # The Qubes OS Project, http://www.qubes-os.org
  4. #
  5. # Copyright (C) 2017 Marek Marczykowski-Górecki
  6. # <marmarek@invisiblethingslab.com>
  7. #
  8. # This library is free software; you can redistribute it and/or
  9. # modify it under the terms of the GNU Lesser General Public
  10. # License as published by the Free Software Foundation; either
  11. # version 2.1 of the License, or (at your option) any later version.
  12. #
  13. # This library is distributed in the hope that it will be useful,
  14. # but WITHOUT ANY WARRANTY; without even the implied warranty of
  15. # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
  16. # Lesser General Public License for more details.
  17. #
  18. # You should have received a copy of the GNU Lesser General Public
  19. # License along with this library; if not, see <https://www.gnu.org/licenses/>.
  20. import asyncio
  21. import socket
  22. import unittest.mock
  23. import qubes.api
  24. import qubes.tests
  25. class TestMgmt(object):
  26. def __init__(self, app, src, method, dest, arg, send_event=None):
  27. self.app = app
  28. self.src = src
  29. self.method = method
  30. self.dest = dest
  31. self.arg = arg
  32. self.send_event = send_event
  33. try:
  34. self.function = {
  35. 'mgmt.success': self.success,
  36. 'mgmt.success_none': self.success_none,
  37. 'mgmt.qubesexception': self.qubesexception,
  38. 'mgmt.exception': self.exception,
  39. 'mgmt.event': self.event,
  40. }[self.method.decode()]
  41. except KeyError:
  42. raise qubes.api.ProtocolError('Invalid method')
  43. def execute(self, untrusted_payload):
  44. self.task = asyncio.Task(self.function(
  45. untrusted_payload=untrusted_payload))
  46. return self.task
  47. def cancel(self):
  48. self.task.cancel()
  49. @asyncio.coroutine
  50. def success(self, untrusted_payload):
  51. return 'src: {!r}, dest: {!r}, arg: {!r}, payload: {!r}'.format(
  52. self.src, self.dest, self.arg, untrusted_payload
  53. )
  54. @asyncio.coroutine
  55. def success_none(self, untrusted_payload):
  56. pass
  57. @asyncio.coroutine
  58. def qubesexception(self, untrusted_payload):
  59. raise qubes.exc.QubesException('qubes-exception')
  60. @asyncio.coroutine
  61. def exception(self, untrusted_payload):
  62. raise Exception('exception')
  63. @asyncio.coroutine
  64. def event(self, untrusted_payload):
  65. future = asyncio.get_event_loop().create_future()
  66. class Subject:
  67. name = 'subject'
  68. def __str__(self):
  69. return 'subject'
  70. self.send_event(Subject(), 'event', payload=untrusted_payload.decode())
  71. try:
  72. # give some time to close the other end
  73. yield from asyncio.sleep(0.1)
  74. # should be canceled
  75. self.send_event(Subject, 'event2',
  76. payload=untrusted_payload.decode())
  77. yield from future
  78. except asyncio.CancelledError:
  79. pass
  80. class TC_00_QubesDaemonProtocol(qubes.tests.QubesTestCase):
  81. def setUp(self):
  82. super(TC_00_QubesDaemonProtocol, self).setUp()
  83. self.app = unittest.mock.Mock()
  84. self.app.log = self.log
  85. self.sock_client, self.sock_server = socket.socketpair()
  86. self.reader, self.writer = self.loop.run_until_complete(
  87. asyncio.open_connection(sock=self.sock_client))
  88. connect_coro = self.loop.create_connection(
  89. lambda: qubes.api.QubesDaemonProtocol(
  90. TestMgmt, app=self.app),
  91. sock=self.sock_server)
  92. self.transport, self.protocol = self.loop.run_until_complete(
  93. connect_coro)
  94. def tearDown(self):
  95. self.writer.close()
  96. try:
  97. self.loop.run_until_complete(self.writer.wait_closed())
  98. except AttributeError: # old python in travis
  99. pass
  100. self.transport.close()
  101. super(TC_00_QubesDaemonProtocol, self).tearDown()
  102. def test_000_message_ok(self):
  103. self.writer.write(b'dom0\0mgmt.success\0dom0\0arg\0payload')
  104. self.writer.write_eof()
  105. with self.assertNotRaises(asyncio.TimeoutError):
  106. response = self.loop.run_until_complete(
  107. asyncio.wait_for(self.reader.read(), 1))
  108. self.assertEqual(response,
  109. b"0\0src: b'dom0', dest: b'dom0', arg: b'arg', payload: b'payload'")
  110. def test_001_message_ok_in_parts(self):
  111. self.writer.write(b'dom0\0mgmt.')
  112. self.loop.run_until_complete(self.writer.drain())
  113. self.writer.write(b'success\0dom0\0arg\0payload')
  114. self.writer.write_eof()
  115. with self.assertNotRaises(asyncio.TimeoutError):
  116. response = self.loop.run_until_complete(
  117. asyncio.wait_for(self.reader.read(), 1))
  118. self.assertEqual(response,
  119. b"0\0src: b'dom0', dest: b'dom0', arg: b'arg', payload: b'payload'")
  120. def test_002_message_ok_empty(self):
  121. self.writer.write(b'dom0\0mgmt.success_none\0dom0\0arg\0payload')
  122. self.writer.write_eof()
  123. with self.assertNotRaises(asyncio.TimeoutError):
  124. response = self.loop.run_until_complete(
  125. asyncio.wait_for(self.reader.read(), 1))
  126. self.assertEqual(response, b"0\0")
  127. def test_003_exception_qubes(self):
  128. self.writer.write(b'dom0\0mgmt.qubesexception\0dom0\0arg\0payload')
  129. self.writer.write_eof()
  130. with self.assertNotRaises(asyncio.TimeoutError):
  131. response = self.loop.run_until_complete(
  132. asyncio.wait_for(self.reader.read(), 1))
  133. self.assertEqual(response, b"2\0QubesException\0\0qubes-exception\0")
  134. def test_004_exception_generic(self):
  135. self.writer.write(b'dom0\0mgmt.exception\0dom0\0arg\0payload')
  136. self.writer.write_eof()
  137. with self.assertNotRaises(asyncio.TimeoutError):
  138. response = self.loop.run_until_complete(
  139. asyncio.wait_for(self.reader.read(), 1))
  140. self.assertEqual(response, b"")
  141. def test_005_event(self):
  142. self.writer.write(b'dom0\0mgmt.event\0dom0\0arg\0payload')
  143. self.writer.write_eof()
  144. with self.assertNotRaises(asyncio.TimeoutError):
  145. response = self.loop.run_until_complete(
  146. asyncio.wait_for(self.reader.readuntil(b'\0\0'), 1))
  147. self.assertEqual(response, b"1\0subject\0event\0payload\0payload\0\0")
  148. # this will trigger connection_lost, but only when next event is sent
  149. self.sock_client.shutdown(socket.SHUT_RD)
  150. # check if event-producing method is interrupted
  151. with self.assertNotRaises(asyncio.TimeoutError):
  152. self.loop.run_until_complete(
  153. asyncio.wait_for(self.protocol.mgmt.task, 1))