Merge pull request #23 from tomato42/response-handling

NBD response handling

fsresck/nbd/response.py

@@ -0,0 +1,93 @@
+# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+#
+#   Description: File system resilience testing application
+#   Author: Hubert Kario <hubert@kario.pl>
+#
+# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+#
+#   Copyright (c) 2017 Hubert Kario. All rights reserved.
+#
+#   This copyrighted material is made available to anyone wishing
+#   to use, modify, copy, or redistribute it subject to the terms
+#   and conditions of the GNU General Public License version 2.
+#
+#   This program is distributed in the hope that it will be
+#   useful, but WITHOUT ANY WARRANTY; without even the implied
+#   warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR
+#   PURPOSE. See the GNU General Public License for more details.
+#
+#   You should have received a copy of the GNU General Public
+#   License along with this program; if not, write to the Free
+#   Software Foundation, Inc., 51 Franklin Street, Fifth Floor,
+#   Boston, MA 02110-1301, USA.
+#
+# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+"""Handling of NBD replies."""
+
+import struct
+from .request import recvexactly
+from ..compat import compat_str
+from .constants import Magic
+from ..errors import FSError
+
+class NBDResponse(object):
+    """Representation of a single NBD protocol response."""
+
+    def __init__(self, error, handle):
+        """Initialise NBD response."""
+        self.error = error
+        self.handle = handle
+        self.data = None
+
+
+class NBDResponseSocket(object):
+    """Handle NBD responses on NBD socket."""
+
+    response_fmt = ">IIQ"
+    response_length = struct.calcsize(response_fmt)
+
+    def __init__(self, sock_in, sock_out, len_dict):
+        """
+        Initialise the object
+
+        :param sock_in: socket to read messages from
+        :param sock_out: socket to write messages to
+        :param len_dict: shared dictionary for keeping the expected lengths
+            of responses
+        """
+        self.sock_in = sock_in
+        self.sock_out = sock_out
+        self.len_dict = len_dict
+
+    def recv(self):
+        """Read a single response from socket."""
+        header = recvexactly(self.sock_in, self.response_length)
+        assert len(header) == self.response_length
+        header = compat_str(header)
+        result_tuple = struct.unpack(self.response_fmt, header)
+        magic, error, handle = result_tuple
+
+        if magic != Magic.NBD_REPLY_MAGIC:
+            raise FSError("Response magic invalid: {0}".format(magic))
+
+        if error != 0:
+            # TODO
+            raise FSError("Received error response, don't know how to handle")
+
+        data_len = self.len_dict.pop(handle)
+        data = recvexactly(self.sock_in, data_len)
+
+        resp = NBDResponse(error, handle)
+        resp.data = data
+
+        return resp
+
+    def send(self, resp):
+        data = struct.pack(self.response_fmt,
+                           Magic.NBD_REPLY_MAGIC,
+                           resp.error,
+                           resp.handle)
+        if resp.data:
+            data += resp.data
+        self.sock_out.sendall(data)

tests/nbd/test_response.py

@@ -0,0 +1,124 @@
+# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+#
+#   Description: File system resilience testing application
+#   Author: Hubert Kario <hubert@kario.pl>
+#
+# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+#
+#   Copyright (c) 2017 Hubert Kario. All rights reserved.
+#
+#   This copyrighted material is made available to anyone wishing
+#   to use, modify, copy, or redistribute it subject to the terms
+#   and conditions of the GNU General Public License version 2.
+#
+#   This program is distributed in the hope that it will be
+#   useful, but WITHOUT ANY WARRANTY; without even the implied
+#   warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR
+#   PURPOSE. See the GNU General Public License for more details.
+#
+#   You should have received a copy of the GNU General Public
+#   License along with this program; if not, write to the Free
+#   Software Foundation, Inc., 51 Franklin Street, Fifth Floor,
+#   Boston, MA 02110-1301, USA.
+#
+# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+# compatibility with Python 2.6, for that we need unittest2 package,
+# which is not available on 3.3 or 3.4
+try:
+    import unittest2 as unittest
+except ImportError:
+    import unittest
+
+try:
+    import mock
+    from mock import call
+except ImportError:
+    import unittest.mock as mock
+    from unittest.mock import call
+
+
+from fsresck.nbd.response import NBDResponseSocket, NBDResponse
+from fsresck.errors import FSError
+
+class TestNBDResponseSocket(unittest.TestCase):
+    def test___init__(self):
+        obj = NBDResponseSocket(None, None, None)
+
+        self.assertIsInstance(obj, NBDResponseSocket)
+
+    @mock.patch('fsresck.nbd.response.recvexactly')
+    def test_recv(self, mock_mthd):
+        mock_mthd.side_effect = [bytearray(
+                b'\x67\x44\x66\x98'  # magic value
+                b'\x00\x00\x00\x00'  # error - none
+                b'\x50\xe4\x93\x01\x00\x88\xff\xff'),  # handle
+                bytearray(b'\x00' * 4096)]
+        d = dict()
+        d[0x50e493010088ffff] = 4096
+
+        obj = NBDResponseSocket(None, None, d).recv()
+
+        self.assertEqual(obj.handle, 0x50e493010088ffff)
+        self.assertEqual(obj.error, 0)
+        self.assertEqual(obj.data, bytearray(b'\x00' * 4096))
+        self.assertEqual(len(d), 0)
+
+    @mock.patch('fsresck.nbd.response.recvexactly')
+    def test_recv_with_invalid_magic(self, mock_mthd):
+        mock_mthd.side_effect = [bytearray(
+                b'\xff\xff\x66\x98'  # magic value
+                b'\x00\x00\x00\x00'  # error - none
+                b'\x50\xe4\x93\x01\x00\x88\xff\xff'),  # handle
+                bytearray(b'\x00' * 4096)]
+        d = dict()
+        d[0x50e493010088ffff] = 4096
+
+        sock = NBDResponseSocket(None, None, d)
+
+        with self.assertRaises(FSError):
+            sock.recv()
+
+    @mock.patch('fsresck.nbd.response.recvexactly')
+    def test_recv_with_error(self, mock_mthd):
+        mock_mthd.side_effect = [bytearray(
+                b'\x67\x44\x66\x98'  # magic value
+                b'\x00\x00\x00\x01'  # error - none
+                b'\x50\xe4\x93\x01\x00\x88\xff\xff'),  # handle
+                bytearray(b'\x00' * 4096)]
+        d = dict()
+        d[0x50e493010088ffff] = 4096
+
+        sock = NBDResponseSocket(None, None, d)
+
+        with self.assertRaises(FSError):
+            sock.recv()
+
+    def test_send(self):
+        sock_out = mock.MagicMock()
+
+        sock = NBDResponseSocket(None, sock_out, None)
+
+        response = NBDResponse(0, 44)
+        response.data = bytearray(b'\x00' * 1024)
+
+        sock.send(response)
+
+        sock_out.sendall.assert_called_with(bytearray(b'\x67\x44\x66\x98'
+                                                      b'\x00\x00\x00\x00'
+                                                      b'\x00\x00\x00\x00'
+                                                      b'\x00\x00\x00\x2c' +
+                                                      b'\x00' * 1024))
+    def test_send_without_data(self):
+        sock_out = mock.MagicMock()
+
+        sock = NBDResponseSocket(None, sock_out, None)
+
+        response = NBDResponse(0, 45)
+
+        sock.send(response)
+
+        sock_out.sendall.assert_called_with(bytearray(b'\x67\x44\x66\x98'
+                                                      b'\x00\x00\x00\x00'
+                                                      b'\x00\x00\x00\x00'
+                                                      b'\x00\x00\x00\x2d'))