From 63ac952803a0f6767ed64128f4e02b846647730b Mon Sep 17 00:00:00 2001
From: Pawel Marczewski <pawel@invisiblethingslab.com>
Date: Thu, 16 Jan 2020 14:41:00 +0100
Subject: [PATCH 1/2] Implement new admin.vm.ImportWithSize API call

This should allow importing a volume and changing the size at the
same time, without performing the resize operation on original
volume first.

The internal API has been renamed to internal.vm.volume.ImportBegin
to avoid confusion, and for symmetry with ImportEnd.

See QubesOS/qubes-issues#5239.
---
 Makefile                                      |  4 +-
 .../admin.vm.volume.Import.policy             | 13 +++++
 .../admin.vm.volume.ImportWithSize.policy     | 13 +++++
 qubes-rpc/admin.vm.volume.Import              | 25 ++++++--
 qubes/api/admin.py                            | 30 ----------
 qubes/api/internal.py                         | 57 ++++++++++++++++++-
 qubes/storage/__init__.py                     | 24 +++++---
 qubes/storage/file.py                         |  4 +-
 qubes/storage/lvm.py                          |  4 +-
 qubes/storage/reflink.py                      |  4 +-
 qubes/tests/api_admin.py                      | 51 ++++++++++++++---
 qubes/tests/storage_file.py                   | 30 +++++++++-
 qubes/tests/storage_lvm.py                    | 31 +++++++++-
 qubes/tests/storage_reflink.py                |  2 +-
 14 files changed, 226 insertions(+), 66 deletions(-)
 create mode 100644 qubes-rpc-policy/admin.vm.volume.Import.policy
 create mode 100644 qubes-rpc-policy/admin.vm.volume.ImportWithSize.policy

diff --git a/Makefile b/Makefile
index df64f81e..78197e83 100644
--- a/Makefile
+++ b/Makefile
@@ -191,6 +191,8 @@ endif
 	cp qubes-rpc-policy/qubes.GetDate.policy $(DESTDIR)/etc/qubes-rpc/policy/qubes.GetDate
 	cp qubes-rpc-policy/qubes.ConnectTCP.policy $(DESTDIR)/etc/qubes-rpc/policy/qubes.ConnectTCP
 	cp qubes-rpc-policy/admin.vm.Console.policy $(DESTDIR)/etc/qubes-rpc/policy/admin.vm.Console
+	cp qubes-rpc-policy/admin.vm.volume.Import.policy $(DESTDIR)/etc/qubes-rpc/policy/admin.vm.volume.Import
+	cp qubes-rpc-policy/admin.vm.volume.ImportWithSize.policy $(DESTDIR)/etc/qubes-rpc/policy/admin.vm.volume.ImportWithSize
 	cp qubes-rpc-policy/policy.RegisterArgument.policy $(DESTDIR)/etc/qubes-rpc/policy/policy.RegisterArgument
 	cp qubes-rpc/qubes.FeaturesRequest $(DESTDIR)/etc/qubes-rpc/
 	cp qubes-rpc/qubes.GetDate $(DESTDIR)/etc/qubes-rpc/
@@ -208,6 +210,7 @@ endif
 			$(DESTDIR)/etc/qubes-rpc/$$method || exit 1; \
 	done
 	install qubes-rpc/admin.vm.volume.Import $(DESTDIR)/etc/qubes-rpc/
+	ln -s admin.vm.volume.Import $(DESTDIR)/etc/qubes-rpc/admin.vm.volume.ImportWithSize
 	install qubes-rpc/admin.vm.Console $(DESTDIR)/etc/qubes-rpc/
 	PYTHONPATH=.:test-packages qubes-rpc-policy/generate-admin-policy \
 		--destdir=$(DESTDIR)/etc/qubes-rpc/policy \
@@ -260,4 +263,3 @@ msi:
 	candle -arch x64 -dversion=$(VERSION) installer.wxs
 	light -b destinstdir -o core-admin.msm installer.wixobj
 	rm -rf destinstdir
-
diff --git a/qubes-rpc-policy/admin.vm.volume.Import.policy b/qubes-rpc-policy/admin.vm.volume.Import.policy
new file mode 100644
index 00000000..c2b92162
--- /dev/null
+++ b/qubes-rpc-policy/admin.vm.volume.Import.policy
@@ -0,0 +1,13 @@
+## Note that policy parsing stops at the first match.
+## Anything not specifically allowed here (or in included file) will be denied.
+
+## Please use a single # to start your custom comments
+
+## Add your entries here, make sure to append ",target=dom0" to all allow/ask actions
+
+## Include a common file for all admin.* methods to ease setting up
+## Management VM.
+## To allow only specific actions, edit specific policy file, like this one. To
+## allow all of them, edit appropriate /etc/qubes-rpc/include/admin-*.
+
+$include:include/admin-local-rwx
diff --git a/qubes-rpc-policy/admin.vm.volume.ImportWithSize.policy b/qubes-rpc-policy/admin.vm.volume.ImportWithSize.policy
new file mode 100644
index 00000000..c2b92162
--- /dev/null
+++ b/qubes-rpc-policy/admin.vm.volume.ImportWithSize.policy
@@ -0,0 +1,13 @@
+## Note that policy parsing stops at the first match.
+## Anything not specifically allowed here (or in included file) will be denied.
+
+## Please use a single # to start your custom comments
+
+## Add your entries here, make sure to append ",target=dom0" to all allow/ask actions
+
+## Include a common file for all admin.* methods to ease setting up
+## Management VM.
+## To allow only specific actions, edit specific policy file, like this one. To
+## allow all of them, edit appropriate /etc/qubes-rpc/include/admin-*.
+
+$include:include/admin-local-rwx
diff --git a/qubes-rpc/admin.vm.volume.Import b/qubes-rpc/admin.vm.volume.Import
index f41ee987..db57c9e1 100755
--- a/qubes-rpc/admin.vm.volume.Import
+++ b/qubes-rpc/admin.vm.volume.Import
@@ -16,23 +16,36 @@
 #    2. Actual data import (done by this script, using dd)
 #    3. Report final result, produce final response to the caller (done by
 #       qubesd)
-#    
+#
 #    This way we do not pass all the data through qubesd, but still can
 #    control the process from there in a meaningful way. Note that the last
 #    part (second call to qubesd) may perform all kind of verification (like
 #    a signature check on the data, or so) and can also prevent VM from
 #    starting (hooking also domain-pre-start event) from not verified image.
+#
+# Note that this script implements two calls:
+#   - admin.vm.volume.Import
+#   - admin.vm.volume.ImportWithSize
+# In the case of admin.vm.ImportWithSize, the first line of payload is then
+# data size in bytes. This is so that we can already notify qubesd to create a
+# volume with the new size.
 
 set -e
 
 # use temporary file, because env variables deal poorly with \0 inside
 tmpfile=$(mktemp)
 trap "rm -f $tmpfile" EXIT
-qubesd-query -e \
-        "$QREXEC_REMOTE_DOMAIN" \
-        "admin.vm.volume.Import" \
-        "$QREXEC_REQUESTED_TARGET" \
-        "$1" >$tmpfile
+
+requested_size=""
+if [[ ${0##*/} == admin.vm.volume.ImportWithSize ]]; then
+    read requested_size
+fi
+
+echo -n "$requested_size" | qubesd-query -c /var/run/qubesd.internal.sock \
+    "$QREXEC_REMOTE_DOMAIN" \
+    "internal.vm.volume.ImportBegin" \
+    "$QREXEC_REQUESTED_TARGET" \
+    "$1" >$tmpfile
 
 # exit if qubesd returned an error (not '0\0')
 if [ "$(head -c 2 $tmpfile | xxd -p)" != "3000" ]; then
diff --git a/qubes/api/admin.py b/qubes/api/admin.py
index 17a3ecd8..8a7da0ea 100644
--- a/qubes/api/admin.py
+++ b/qubes/api/admin.py
@@ -486,36 +486,6 @@ class QubesAdminAPI(qubes.api.AbstractQubesAPI):
         finally:  # even if calling qubes.ResizeDisk inside the VM failed
             self.app.save()
 
-    @qubes.api.method('admin.vm.volume.Import', no_payload=True,
-        scope='local', write=True)
-    @asyncio.coroutine
-    def vm_volume_import(self):
-        """Import volume data.
-
-        Note that this function only returns a path to where data should be
-        written, actual importing is done by a script in /etc/qubes-rpc
-        When the script finish importing, it will trigger
-        internal.vm.volume.ImportEnd (with either b'ok' or b'fail' as a
-        payload) and response from that call will be actually send to the
-        caller.
-        """
-        self.enforce(self.arg in self.dest.volumes.keys())
-
-        self.fire_event_for_permission()
-
-        if not self.dest.is_halted():
-            raise qubes.exc.QubesVMNotHaltedError(self.dest)
-
-        path = yield from self.dest.storage.import_data(self.arg)
-        self.enforce(' ' not in path)
-        size = self.dest.volumes[self.arg].size
-
-        # when we know the action is allowed, inform extensions that it will
-        # be performed
-        self.dest.fire_event('domain-volume-import-begin', volume=self.arg)
-
-        return '{} {}'.format(size, path)
-
     @qubes.api.method('admin.vm.volume.Set.revisions_to_keep',
         scope='local', write=True)
     @asyncio.coroutine
diff --git a/qubes/api/internal.py b/qubes/api/internal.py
index 09a05171..297edfcb 100644
--- a/qubes/api/internal.py
+++ b/qubes/api/internal.py
@@ -1,4 +1,4 @@
-# -*- encoding: utf8 -*-
+# -*- encoding: utf-8 -*-
 #
 # The Qubes OS Project, http://www.qubes-os.org
 #
@@ -56,6 +56,61 @@ class QubesInternalAPI(qubes.api.AbstractQubesAPI):
 
         return json.dumps(system_info)
 
+    @qubes.api.method('internal.vm.volume.ImportBegin',
+        scope='local', write=True)
+    @asyncio.coroutine
+    def vm_volume_import(self, untrusted_payload):
+        """Begin importing volume data. Payload is either size of new data
+        in bytes, or empty. If empty, the current volume's size will be used.
+        Returns size and path to where data should be written.
+
+        Triggered by scripts in /etc/qubes-rpc:
+        admin.vm.volume.Import, admin.vm.volume.ImportWithSize.
+
+        When the script finish importing, it will trigger
+        internal.vm.volume.ImportEnd (with either b'ok' or b'fail' as a
+        payload) and response from that call will be actually send to the
+        caller.
+        """
+        self.enforce(self.arg in self.dest.volumes.keys())
+
+        if untrusted_payload:
+            original_method = 'admin.vm.volume.ImportWithSize'
+        else:
+            original_method = 'admin.vm.volume.Import'
+        self.src.fire_event(
+            'admin-permission:' + original_method,
+            pre_event=True, dest=self.dest, arg=self.arg)
+
+        if not self.dest.is_halted():
+            raise qubes.exc.QubesVMNotHaltedError(self.dest)
+
+        requested_size = None
+        if untrusted_payload:
+            try:
+                untrusted_value = int(untrusted_payload.decode('ascii'))
+            except (UnicodeDecodeError, ValueError):
+                raise qubes.api.ProtocolError('Invalid value')
+            self.enforce(untrusted_value > 0)
+            requested_size = untrusted_value
+            del untrusted_value
+        del untrusted_payload
+
+        path = yield from self.dest.storage.import_data(
+            self.arg, requested_size)
+        self.enforce(' ' not in path)
+        if requested_size is None:
+            size = self.dest.volumes[self.arg].size
+        else:
+            size = requested_size
+
+        # when we know the action is allowed, inform extensions that it will
+        # be performed
+        self.dest.fire_event(
+            'domain-volume-import-begin', volume=self.arg, size=size)
+
+        return '{} {}'.format(size, path)
+
     @qubes.api.method('internal.vm.volume.ImportEnd')
     @asyncio.coroutine
     def vm_volume_import_end(self, untrusted_payload):
diff --git a/qubes/storage/__init__.py b/qubes/storage/__init__.py
index da518a99..f71c2a7b 100644
--- a/qubes/storage/__init__.py
+++ b/qubes/storage/__init__.py
@@ -188,7 +188,7 @@ class Volume:
         '''
         raise self._not_implemented("export")
 
-    def import_data(self):
+    def import_data(self, size):
         ''' Returns a path to overwrite volume data.
 
             This method is called after volume was already :py:meth:`create`-ed.
@@ -199,6 +199,8 @@ class Volume:
             on the fly), the returned path may be a pipe.
 
             This can be implemented as a coroutine.
+
+            :param int size: size of new data in bytes
         '''
         raise self._not_implemented("import_data")
 
@@ -624,14 +626,22 @@ class Storage:
         return self.vm.volumes[volume].export()
 
     @asyncio.coroutine
-    def import_data(self, volume):
-        ''' Helper function to import volume data (pool.import_data(volume))'''
+    def import_data(self, volume, size):
+        '''
+        Helper function to import volume data (pool.import_data(volume)).
+
+        :size: new size in bytes, or None if using old size
+        '''
+
         assert isinstance(volume, (Volume, str)), \
             "You need to pass a Volume or pool name as str"
-        if isinstance(volume, Volume):
-            ret = volume.import_data()
-        else:
-            ret = self.vm.volumes[volume].import_data()
+        if isinstance(volume, str):
+            volume = self.vm.volumes[volume]
+
+        if size is None:
+            size = volume.size
+
+        ret = volume.import_data(size)
         return (yield from qubes.utils.coro_maybe(ret))
 
     @asyncio.coroutine
diff --git a/qubes/storage/file.py b/qubes/storage/file.py
index cd3ee9d7..c4ef05a5 100644
--- a/qubes/storage/file.py
+++ b/qubes/storage/file.py
@@ -283,12 +283,12 @@ class FileVolume(qubes.storage.Volume):
             copy_file(src_volume.export(), self.path)
         return self
 
-    def import_data(self):
+    def import_data(self, size):
         if not self.save_on_stop:
             raise qubes.storage.StoragePoolException(
                 "Can not import into save_on_stop=False volume {!s}".format(
                     self))
-        create_sparse_file(self.path_import, self.size)
+        create_sparse_file(self.path_import, size)
         return self.path_import
 
     def import_data_end(self, success):
diff --git a/qubes/storage/lvm.py b/qubes/storage/lvm.py
index b0a3b52d..1e2f0e9d 100644
--- a/qubes/storage/lvm.py
+++ b/qubes/storage/lvm.py
@@ -560,7 +560,7 @@ class ThinVolume(qubes.storage.Volume):
 
     @locked
     @asyncio.coroutine
-    def import_data(self):
+    def import_data(self, size):
         ''' Returns an object that can be `open()`. '''
         if self.is_dirty():
             raise qubes.storage.StoragePoolException(
@@ -569,7 +569,7 @@ class ThinVolume(qubes.storage.Volume):
         self.abort_if_import_in_progress()
         # pylint: disable=protected-access
         cmd = ['create', self.pool._pool_id, self._vid_import.split('/')[1],
-               str(self.size)]
+               str(size)]
         yield from qubes_lvm_coro(cmd, self.log)
         yield from reset_cache_coro()
         devpath = '/dev/' + self._vid_import
diff --git a/qubes/storage/reflink.py b/qubes/storage/reflink.py
index 795a7ffc..a7aa320c 100644
--- a/qubes/storage/reflink.py
+++ b/qubes/storage/reflink.py
@@ -289,11 +289,11 @@ class ReflinkVolume(qubes.storage.Volume):
 
     @_coroutinized
     @_locked
-    def import_data(self):
+    def import_data(self, size):
         if not self.save_on_stop:
             raise NotImplementedError(
                 'Cannot import_data: {} is not save_on_stop'.format(self.vid))
-        _create_sparse_file(self._path_import, self._get_size())
+        _create_sparse_file(self._path_import, size)
         return self._path_import
 
     def _import_data_end(self, success):
diff --git a/qubes/tests/api_admin.py b/qubes/tests/api_admin.py
index 0cd6faa2..185b487f 100644
--- a/qubes/tests/api_admin.py
+++ b/qubes/tests/api_admin.py
@@ -1,4 +1,4 @@
-# -*- encoding: utf8 -*-
+# -*- encoding: utf-8 -*-
 #
 # The Qubes OS Project, http://www.qubes-os.org
 #
@@ -34,6 +34,7 @@ import qubes
 import qubes.devices
 import qubes.firewall
 import qubes.api.admin
+import qubes.api.internal
 import qubes.tests
 import qubes.storage
 
@@ -113,6 +114,13 @@ class AdminAPITestCase(qubes.tests.QubesTestCase):
             'admin-permission:' + method.decode('ascii'))
         return response
 
+    def call_internal_mgmt_func(self, method, dest, arg=b'', payload=b''):
+        mgmt_obj = qubes.api.internal.QubesInternalAPI(self.app, b'dom0', method, dest, arg)
+        loop = asyncio.get_event_loop()
+        response = loop.run_until_complete(
+            mgmt_obj.execute(untrusted_payload=payload))
+        return response
+
 
 class TC_00_VMs(AdminAPITestCase):
     def test_000_vm_list(self):
@@ -207,11 +215,9 @@ qid default=False type=int 2
 qrexec_timeout default=True type=int 60
 updateable default=True type=bool False
 kernelopts default=False type=str opt1\\nopt2\\nopt3\\\\opt4
-netvm default=True type=vm 
-'''
+netvm default=True type=vm \n'''
         self.assertEqual(value, expected)
 
-
     def test_030_vm_property_set_vm(self):
         netvm = self.app.add_new_vm('AppVM', label='red', name='test-net',
             template='test-template', provides_network=True)
@@ -1747,9 +1753,12 @@ netvm default=True type=vm
         self.assertFalse(mock_remove.called)
         self.assertFalse(self.app.save.called)
 
+    # Import tests
+    # (internal methods, normally called from qubes-rpc script)
+
     def test_510_vm_volume_import(self):
-        value = self.call_mgmt_func(b'admin.vm.volume.Import', b'test-vm1',
-            b'private')
+        value = self.call_internal_mgmt_func(
+            b'internal.vm.volume.ImportBegin', b'test-vm1', b'private')
         self.assertEqual(value, '{} {}'.format(
             2*2**30, '/tmp/qubes-test-dir/appvms/test-vm1/private-import.img'))
         self.assertFalse(self.app.save.called)
@@ -1758,8 +1767,34 @@ netvm default=True type=vm
         with unittest.mock.patch.object(
                 self.vm, 'get_power_state', lambda: 'Running'):
             with self.assertRaises(qubes.exc.QubesVMNotHaltedError):
-                self.call_mgmt_func(b'admin.vm.volume.Import', b'test-vm1',
-                    b'private')
+                self.call_internal_mgmt_func(
+                    b'internal.vm.volume.ImportBegin', b'test-vm1', b'private')
+
+    def test_512_vm_volume_import_with_size(self):
+        new_size = 4 * 2**30
+        file_name = '/tmp/qubes-test-dir/appvms/test-vm1/private-import.img'
+
+        value = self.call_internal_mgmt_func(
+            b'internal.vm.volume.ImportBegin', b'test-vm1',
+            b'private', payload=str(new_size).encode())
+        self.assertEqual(value, '{} {}'.format(
+            new_size, file_name))
+        self.assertFalse(self.app.save.called)
+
+        self.assertEqual(os.stat(file_name).st_size, new_size)
+
+    def test_515_vm_volume_import_fire_event(self):
+        self.call_internal_mgmt_func(
+            b'internal.vm.volume.ImportBegin', b'test-vm1', b'private')
+        self.assertEventFired(
+            self.emitter, 'admin-permission:admin.vm.volume.Import')
+
+    def test_516_vm_volume_import_fire_event_with_size(self):
+        self.call_internal_mgmt_func(
+            b'internal.vm.volume.ImportBegin', b'test-vm1', b'private',
+            b'123')
+        self.assertEventFired(
+            self.emitter, 'admin-permission:admin.vm.volume.ImportWithSize')
 
     def setup_for_clone(self):
         self.pool = unittest.mock.MagicMock()
diff --git a/qubes/tests/storage_file.py b/qubes/tests/storage_file.py
index 9742ffa1..bc3df8cb 100644
--- a/qubes/tests/storage_file.py
+++ b/qubes/tests/storage_file.py
@@ -332,7 +332,7 @@ class TC_01_FileVolumes(qubes.tests.QubesTestCase):
         vm = qubes.tests.storage.TestVM(self)
         volume = self.app.get_pool(self.POOL_NAME).init_volume(vm, config)
         volume.create()
-        import_path = volume.import_data()
+        import_path = volume.import_data(volume.size)
         self.assertNotEqual(volume.path, import_path)
         with open(import_path, 'w+') as import_file:
             import_file.write('test')
@@ -353,7 +353,7 @@ class TC_01_FileVolumes(qubes.tests.QubesTestCase):
         vm = qubes.tests.storage.TestVM(self)
         volume = self.app.get_pool(self.POOL_NAME).init_volume(vm, config)
         volume.create()
-        import_path = volume.import_data()
+        import_path = volume.import_data(volume.size)
         self.assertNotEqual(volume.path, import_path)
         with open(import_path, 'w+') as import_file:
             import_file.write('test')
@@ -376,7 +376,7 @@ class TC_01_FileVolumes(qubes.tests.QubesTestCase):
         volume.create()
         with open(volume.path, 'w') as vol_file:
             vol_file.write('test data')
-        import_path = volume.import_data()
+        import_path = volume.import_data(volume.size)
         self.assertNotEqual(volume.path, import_path)
         with open(import_path, 'w+'):
             pass
@@ -402,6 +402,30 @@ class TC_01_FileVolumes(qubes.tests.QubesTestCase):
         self.assertEqual(os.path.getsize(volume.path), new_size)
         self.assertEqual(volume.size, new_size)
 
+    def test_024_import_data_with_new_size(self):
+        config = {
+            'name': 'root',
+            'pool': self.POOL_NAME,
+            'save_on_stop': True,
+            'rw': True,
+            'size': 1024 * 1024,
+        }
+        new_size = 2 * 1024 * 1024
+
+        vm = qubes.tests.storage.TestVM(self)
+        volume = self.app.get_pool(self.POOL_NAME).init_volume(vm, config)
+        volume.create()
+        import_path = volume.import_data(new_size)
+        self.assertNotEqual(volume.path, import_path)
+        with open(import_path, 'r+b') as import_file:
+            import_file.write(b'test')
+        volume.import_data_end(True)
+        self.assertFalse(os.path.exists(import_path), import_path)
+        with open(volume.path, 'rb') as volume_file:
+            volume_data = volume_file.read()
+        self.assertEqual(volume_data.strip(b'\0'), b'test')
+        self.assertEqual(len(volume_data), new_size)
+
     def _get_loop_size(self, path):
         sudo = [] if os.getuid() == 0 else ['sudo']
         try:
diff --git a/qubes/tests/storage_lvm.py b/qubes/tests/storage_lvm.py
index 563fdf9f..5646ac9f 100644
--- a/qubes/tests/storage_lvm.py
+++ b/qubes/tests/storage_lvm.py
@@ -698,7 +698,8 @@ class TC_00_ThinPool(ThinPoolBase):
         self.loop.run_until_complete(volume.create())
         current_uuid = self._get_lv_uuid(volume.path)
         self.assertFalse(volume.is_dirty())
-        import_path = self.loop.run_until_complete(volume.import_data())
+        import_path = self.loop.run_until_complete(
+            volume.import_data(volume.size))
         import_uuid = self._get_lv_uuid(import_path)
         self.assertNotEqual(current_uuid, import_uuid)
         # success - commit data
@@ -729,7 +730,8 @@ class TC_00_ThinPool(ThinPoolBase):
         self.loop.run_until_complete(volume.create())
         current_uuid = self._get_lv_uuid(volume.path)
         self.assertFalse(volume.is_dirty())
-        import_path = self.loop.run_until_complete(volume.import_data())
+        import_path = self.loop.run_until_complete(
+            volume.import_data(volume.size))
         import_uuid = self._get_lv_uuid(import_path)
         self.assertNotEqual(current_uuid, import_uuid)
         # fail - discard data
@@ -860,7 +862,8 @@ class TC_00_ThinPool(ThinPoolBase):
             'sudo', 'dd', 'if=/dev/urandom', 'of=' + volume.path, 'count=1', 'bs=1M'
         ))
         self.loop.run_until_complete(p.wait())
-        import_path = self.loop.run_until_complete(volume.import_data())
+        import_path = self.loop.run_until_complete(
+            volume.import_data(volume.size))
         self.assertNotEqual(volume.path, import_path)
         p = self.loop.run_until_complete(asyncio.create_subprocess_exec(
             'sudo', 'touch', import_path))
@@ -874,6 +877,28 @@ class TC_00_ThinPool(ThinPoolBase):
         volume_data, _ = self.loop.run_until_complete(p.communicate())
         self.assertEqual(volume_data.strip(b'\0'), b'')
 
+    def test_035_import_data_new_size(self):
+        ''' Test volume import'''
+        config = {
+            'name': 'root',
+            'pool': self.pool.name,
+            'save_on_stop': True,
+            'rw': True,
+            'revisions_to_keep': 2,
+            'size': qubes.config.defaults['root_img_size'],
+        }
+        new_size = 2 * qubes.config.defaults['root_img_size']
+
+        vm = qubes.tests.storage.TestVM(self)
+        volume = self.app.get_pool(self.pool.name).init_volume(vm, config)
+        self.loop.run_until_complete(volume.create())
+        self.loop.run_until_complete(
+            volume.import_data(new_size))
+        self.loop.run_until_complete(volume.import_data_end(True))
+        self.assertEqual(volume.size, new_size)
+
+        self.loop.run_until_complete(volume.remove())
+
     def test_040_volatile(self):
         '''Volatile volume test'''
         config = {
diff --git a/qubes/tests/storage_reflink.py b/qubes/tests/storage_reflink.py
index 9da68638..5483cb96 100644
--- a/qubes/tests/storage_reflink.py
+++ b/qubes/tests/storage_reflink.py
@@ -128,7 +128,7 @@ class TC_10_ReflinkPool(qubes.tests.QubesTestCase):
         self.loop.run_until_complete(volume.create())
         with open(volume.export(), 'w') as vol_file:
             vol_file.write('test data')
-        import_path = self.loop.run_until_complete(volume.import_data())
+        import_path = self.loop.run_until_complete(volume.import_data(volume.size))
         self.assertNotEqual(volume.path, import_path)
         with open(import_path, 'w+'):
             pass

From e9b97e42b1dffe7b1d7317d73fcfa912f724d016 Mon Sep 17 00:00:00 2001
From: Pawel Marczewski <pawel@invisiblethingslab.com>
Date: Mon, 20 Jan 2020 18:14:05 +0100
Subject: [PATCH 2/2] import: check exact size of copied data

The import will error out if there is not enough data, or too
much data provided.
---
 qubes-rpc/admin.vm.volume.Import |  45 ++++++--
 qubes/api/internal.py            |   9 +-
 qubes/tests/__init__.py          |   1 +
 qubes/tests/api_admin.py         |  22 ++++
 qubes/tests/rpc_import.py        | 180 +++++++++++++++++++++++++++++++
 rpm_spec/core-dom0.spec.in       |   1 +
 6 files changed, 247 insertions(+), 11 deletions(-)
 create mode 100644 qubes/tests/rpc_import.py

diff --git a/qubes-rpc/admin.vm.volume.Import b/qubes-rpc/admin.vm.volume.Import
index db57c9e1..88704596 100755
--- a/qubes-rpc/admin.vm.volume.Import
+++ b/qubes-rpc/admin.vm.volume.Import
@@ -1,4 +1,4 @@
-#!/bin/sh
+#!/bin/bash
 #
 # This Admin API call is implemented as a custom script, instead of dumb
 # passthrough to qubesd because it may get huge amount of data (whole root.img
@@ -32,39 +32,64 @@
 
 set -e
 
+# make dd output consistent
+export LC_ALL=C
+
 # use temporary file, because env variables deal poorly with \0 inside
 tmpfile=$(mktemp)
-trap "rm -f $tmpfile" EXIT
+trap 'rm -f $tmpfile' EXIT
 
 requested_size=""
-if [[ ${0##*/} == admin.vm.volume.ImportWithSize ]]; then
-    read requested_size
+if [[ "${0##*/}" == admin.vm.volume.ImportWithSize ]]; then
+    read -r requested_size
 fi
 
 echo -n "$requested_size" | qubesd-query -c /var/run/qubesd.internal.sock \
     "$QREXEC_REMOTE_DOMAIN" \
     "internal.vm.volume.ImportBegin" \
     "$QREXEC_REQUESTED_TARGET" \
-    "$1" >$tmpfile
+    "$1" >"$tmpfile"
 
 # exit if qubesd returned an error (not '0\0')
-if [ "$(head -c 2 $tmpfile | xxd -p)" != "3000" ]; then
+if [ "$(head -c 2 "$tmpfile" | xxd -p)" != "3000" ]; then
     cat "$tmpfile"
     exit 1
 fi
 size=$(tail -c +3 "$tmpfile"|cut -d ' ' -f 1)
 path=$(tail -c +3 "$tmpfile"|cut -d ' ' -f 2)
 
+error=""
+
 # now process stdin into this path
-if sudo dd bs=128K of="$path" count="$size" iflag=count_bytes,fullblock \
-        conv=sparse,notrunc,nocreat,fdatasync status=none; then
+if ! sudo dd bs=128K of="$path" count="$size" iflag=count_bytes,fullblock \
+        conv=sparse,notrunc,nocreat,fdatasync 2>"$tmpfile"; then
+    error="error copying data"
+fi
+
+# Examine dd's output and check if number of bytes copied matches
+if [ -z "$error" ]; then
+    bytes_copied=$(tail -n1 "$tmpfile" | cut -d ' ' -f 1)
+    if [ "$bytes_copied" -ne "$size" ]; then
+        error="not enough data (copied $bytes_copied bytes, expected $size bytes)"
+    fi
+fi
+
+# Check if there is nothing more to be read from stdin
+if [ -z "$error" ]; then
+    if ! dd of="$tmpfile" bs=1 count=1 status=none || \
+            [ "$(stat -c %s "$tmpfile")" -ne 0 ]; then
+        error="too much data (expected $size bytes)"
+    fi
+fi
+
+if [ -z "$error" ]; then
     status="ok"
 else
-    status="fail"
+    status="fail\n$error"
 fi
 
 # send status notification to qubesd, and pass its response to the caller
-echo -n "$status" | qubesd-query -c /var/run/qubesd.internal.sock \
+echo -ne "$status" | qubesd-query -c /var/run/qubesd.internal.sock \
     "$QREXEC_REMOTE_DOMAIN" \
     "internal.vm.volume.ImportEnd" \
     "$QREXEC_REQUESTED_TARGET" \
diff --git a/qubes/api/internal.py b/qubes/api/internal.py
index 297edfcb..c4fbb9bc 100644
--- a/qubes/api/internal.py
+++ b/qubes/api/internal.py
@@ -118,6 +118,8 @@ class QubesInternalAPI(qubes.api.AbstractQubesAPI):
         This is second half of admin.vm.volume.Import handling. It is called
         when actual import is finished. Response from this method is sent do
         the client (as a response for admin.vm.volume.Import call).
+
+        The payload is either 'ok', or 'fail\n<error message>'.
         '''
         self.enforce(self.arg in self.dest.volumes.keys())
         success = untrusted_payload == b'ok'
@@ -134,7 +136,12 @@ class QubesInternalAPI(qubes.api.AbstractQubesAPI):
             success=success)
 
         if not success:
-            raise qubes.exc.QubesException('Data import failed')
+            error = ''
+            parts = untrusted_payload.decode('ascii').split('\n', 1)
+            if len(parts) > 1:
+                error = parts[1]
+            raise qubes.exc.QubesException(
+                'Data import failed: {}'.format(error))
 
     @qubes.api.method('internal.SuspendPre', no_payload=True)
     @asyncio.coroutine
diff --git a/qubes/tests/__init__.py b/qubes/tests/__init__.py
index 06260466..2ed87dad 100644
--- a/qubes/tests/__init__.py
+++ b/qubes/tests/__init__.py
@@ -1403,6 +1403,7 @@ def load_tests(loader, tests, pattern): # pylint: disable=unused-argument
             'qubes.tests.api_admin',
             'qubes.tests.api_misc',
             'qubes.tests.api_internal',
+            'qubes.tests.rpc_import',
             ):
         tests.addTests(loader.loadTestsFromName(modname))
 
diff --git a/qubes/tests/api_admin.py b/qubes/tests/api_admin.py
index 185b487f..c9423fea 100644
--- a/qubes/tests/api_admin.py
+++ b/qubes/tests/api_admin.py
@@ -1796,6 +1796,28 @@ netvm default=True type=vm \n'''
         self.assertEventFired(
             self.emitter, 'admin-permission:admin.vm.volume.ImportWithSize')
 
+    def test_510_vm_volume_import_end_success(self):
+        import_data_end_mock, self.vm.storage.import_data_end = \
+            self.coroutine_mock()
+        self.call_internal_mgmt_func(
+            b'internal.vm.volume.ImportEnd', b'test-vm1', b'private',
+            payload=b'ok')
+        self.assertEqual(import_data_end_mock.mock_calls, [
+            unittest.mock.call('private', success=True)
+        ])
+
+    def test_510_vm_volume_import_end_failure(self):
+        import_data_end_mock, self.vm.storage.import_data_end = \
+            self.coroutine_mock()
+        with self.assertRaisesRegexp(
+                qubes.exc.QubesException, 'error message'):
+            self.call_internal_mgmt_func(
+                b'internal.vm.volume.ImportEnd', b'test-vm1', b'private',
+                payload=b'fail\nerror message')
+        self.assertEqual(import_data_end_mock.mock_calls, [
+            unittest.mock.call('private', success=False)
+        ])
+
     def setup_for_clone(self):
         self.pool = unittest.mock.MagicMock()
         self.app.pools['test'] = self.pool
diff --git a/qubes/tests/rpc_import.py b/qubes/tests/rpc_import.py
new file mode 100644
index 00000000..5d9f2666
--- /dev/null
+++ b/qubes/tests/rpc_import.py
@@ -0,0 +1,180 @@
+
+#
+# The Qubes OS Project, https://www.qubes-os.org/
+#
+# Copyright (C) 2020 Paweł Marczewski <pawel@invisiblethingslab.com>
+#
+# This library is free software; you can redistribute it and/or
+# modify it under the terms of the GNU Lesser General Public
+# License as published by the Free Software Foundation; either
+# version 2.1 of the License, or (at your option) any later version.
+#
+# This library 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
+# Lesser General Public License for more details.
+#
+# You should have received a copy of the GNU Lesser General Public
+# License along with this library; if not, see <https://www.gnu.org/licenses/>.
+#
+
+import unittest
+import tempfile
+import shutil
+import os
+import subprocess
+
+
+import qubes.tests
+
+
+class TestRpcImport(qubes.tests.QubesTestCase):
+    '''
+    Tests for qubes-rpc/admin.vm.volume.Import script.
+
+    It is a shell script that calls internal API methods via qubesd-query.
+    These tests mock all the calls.
+    '''
+
+
+    QUBESD_QUERY = '''\
+#!/bin/sh -e
+
+method=$4
+echo "$@" > command-$method
+cat > payload-$method
+cat response-$method
+'''
+
+    RPC_FILE_PATH = os.path.abspath(os.path.join(
+        os.path.dirname(__file__),
+        '../../qubes-rpc/admin.vm.volume.Import'))
+
+    def setUp(self):
+        self.tmpdir = tempfile.mkdtemp()
+        self.addCleanup(shutil.rmtree, self.tmpdir)
+        with open(os.path.join(self.tmpdir, 'qubesd-query'), 'w') \
+             as qubesd_query_f:
+            qubesd_query_f.write(self.QUBESD_QUERY)
+        os.chmod(os.path.join(self.tmpdir, 'qubesd-query'), 0o700)
+
+        shutil.copy(
+            self.RPC_FILE_PATH,
+            os.path.join(self.tmpdir, 'admin.vm.volume.Import'))
+        shutil.copy(
+            self.RPC_FILE_PATH,
+            os.path.join(self.tmpdir, 'admin.vm.volume.ImportWithSize'))
+
+    # pylint:disable=invalid-name
+    def mockMethod(self, method, response):
+        with open(os.path.join(self.tmpdir, 'response-' + method), 'wb') \
+             as response_f:
+            response_f.write(response)
+
+    # pylint:disable=invalid-name
+    def assertMethodCalled(self, method, arg, expected_payload=b''):
+        try:
+            with open(os.path.join(self.tmpdir, 'command-' + method), 'rb') \
+                 as command_f:
+                command = command_f.read()
+            with open(os.path.join(self.tmpdir, 'payload-' + method), 'rb') \
+                 as payload_f:
+                payload = payload_f.read()
+        except FileNotFoundError:
+            self.fail('{} was not called'.format(method))
+
+        self.assertListEqual(command.decode().split(), [
+            '-c', '/var/run/qubesd.internal.sock',
+            'remote', method, 'target', arg
+        ])
+        self.assertEqual(payload, expected_payload)
+
+    # pylint:disable=invalid-name
+    def assertFileData(self, path, expected_data):
+        with open(path, 'rb') as data_f:
+            data = data_f.read()
+        self.assertEquals(data, expected_data)
+
+    def setup_import(self, size):
+        self.target = os.path.join(self.tmpdir, 'target')
+        os.mknod(self.target)
+
+        self.mockMethod(
+            'internal.vm.volume.ImportBegin',
+            '\x30\x00{} {}'.format(size, self.target).encode())
+
+        self.mockMethod(
+            'internal.vm.volume.ImportEnd',
+            b'\x30\x00import-end')
+
+    def run_rpc(self, command, arg, data):
+        with open(os.path.join(self.tmpdir, 'data'), 'w+b') as data_f:
+            data_f.write(data)
+            data_f.seek(0)
+            env = {
+                'PATH': self.tmpdir + ':' + os.getenv('PATH'),
+                'QREXEC_REMOTE_DOMAIN': 'remote',
+                'QREXEC_REQUESTED_TARGET': 'target',
+            }
+            output = subprocess.check_output(
+                [command, arg],
+                env=env,
+                cwd=self.tmpdir,
+                stdin=data_f
+            )
+
+        self.assertEqual(output, b'\x30\x00import-end')
+
+    def test_00_import(self):
+        data = b'abcd' * 1024
+        size = len(data)
+        self.setup_import(size)
+
+        self.run_rpc('admin.vm.volume.Import', 'volume', data)
+
+        self.assertMethodCalled('internal.vm.volume.ImportBegin', 'volume')
+        self.assertMethodCalled('internal.vm.volume.ImportEnd', 'volume',
+                                b'ok')
+        self.assertFileData(self.target, data)
+
+    def test_01_import_with_size(self):
+        data = b'abcd' * 1024
+        size = len(data)
+        self.setup_import(size)
+
+        self.run_rpc('admin.vm.volume.ImportWithSize', 'volume',
+                     str(size).encode() + b'\n' + data)
+
+        self.assertMethodCalled('internal.vm.volume.ImportBegin', 'volume',
+                                str(size).encode())
+        self.assertMethodCalled('internal.vm.volume.ImportEnd', 'volume',
+                                b'ok')
+        self.assertFileData(self.target, data)
+
+    def test_02_import_not_enough_data(self):
+        data = b'abcd' * 1024
+        size = len(data) + 1
+        self.setup_import(size)
+
+        self.run_rpc('admin.vm.volume.Import', 'volume', data)
+
+        self.assertMethodCalled('internal.vm.volume.ImportBegin', 'volume')
+        self.assertMethodCalled(
+            'internal.vm.volume.ImportEnd', 'volume',
+            b'fail\n' +
+            ('not enough data (copied {} bytes, expected {} bytes)'
+             .format(len(data), size).encode()))
+
+    def test_03_import_too_much_data(self):
+        data = b'abcd' * 1024
+        size = len(data) - 1
+        self.setup_import(size)
+
+        output = self.run_rpc('admin.vm.volume.Import', 'volume', data)
+
+        self.assertMethodCalled('internal.vm.volume.ImportBegin', 'volume')
+        self.assertMethodCalled(
+            'internal.vm.volume.ImportEnd', 'volume',
+            b'fail\n' +
+            ('too much data (expected {} bytes)'
+             .format(size).encode()))
diff --git a/rpm_spec/core-dom0.spec.in b/rpm_spec/core-dom0.spec.in
index 76ade37e..76339017 100644
--- a/rpm_spec/core-dom0.spec.in
+++ b/rpm_spec/core-dom0.spec.in
@@ -300,6 +300,7 @@ fi
 %{python3_sitelib}/qubes/tests/ext.py
 %{python3_sitelib}/qubes/tests/firewall.py
 %{python3_sitelib}/qubes/tests/init.py
+%{python3_sitelib}/qubes/tests/rpc_import.py
 %{python3_sitelib}/qubes/tests/storage.py
 %{python3_sitelib}/qubes/tests/storage_file.py
 %{python3_sitelib}/qubes/tests/storage_reflink.py