Merge pull request #1660 from ogayot/no-probert-during-partitioning

storage: don't let probing discard storage config in the middle of partitioning

subiquity/common/apidef.py

@@ -277,6 +277,10 @@ class API:
 
         def POST(config: Payload[list]): ...
 
+        class dry_run_wait_probe:
+            """This endpoint only works in dry-run mode."""
+            def POST() -> None: ...
+
         class reset:
             def POST() -> StorageResponse: ...
 
@@ -302,6 +306,16 @@ class API:
             class reset:
                 def POST() -> StorageResponseV2: ...
 
+            class ensure_transaction:
+                """This call will ensure that a transaction is initiated.
+                During a transaction, storage probing runs are not permitted to
+                reset the partitioning configuration.
+                A transaction will also be initiated by any v2_storage POST
+                request that modifies the partitioning configuration (e.g.,
+                add_partition, edit_partition, ...) but ensure_transaction can
+                be called early if desired. """
+                def POST() -> None: ...
+
             class reformat_disk:
                 def POST(data: Payload[ReformatDisk]) \
                     -> StorageResponseV2: ...

subiquity/server/controllers/filesystem.py

@@ -22,7 +22,7 @@ import os
 import pathlib
 import select
 import time
-from typing import Dict, List, Optional
+from typing import Any, Dict, List, Optional
 
 from curtin.commands.extract import AbstractSourceHandler
 from curtin.storage_config import ptable_uuid_to_flag_entry
@@ -160,6 +160,10 @@ class FilesystemController(SubiquityController, FilesystemManipulator):
         self._role_to_device: Dict[str: _Device] = {}
         self._device_to_structure: Dict[_Device: snapdapi.OnVolume] = {}
         self.use_tpm: bool = False
+        self.locked_probe_data = False
+        # If probe data come in while we are doing partitioning, store it in
+        # this variable. It will be picked up on next reset will pick it up.
+        self.queued_probe_data: Optional[Dict[str, Any]] = None
 
     def is_core_boot_classic(self):
         return self._system is not None
@@ -680,9 +684,22 @@ class FilesystemController(SubiquityController, FilesystemManipulator):
 
     async def v2_reset_POST(self) -> StorageResponseV2:
         log.info("Resetting Filesystem model")
-        self.model.reset()
+        # From the API standpoint, it seems sound to set locked_probe_data back
+        # to False after a reset. But in practise, v2_reset_POST can be called
+        # during manual partitioning ; and we don't want to reenable automatic
+        # loading of probe data. Going forward, this could be controlled by an
+        # optional parameter maybe?
+        if self.queued_probe_data is not None:
+            log.debug("using newly obtained probe data")
+            self.model.load_probe_data(self.queued_probe_data)
+            self.queued_probe_data = None
+        else:
+            self.model.reset()
         return await self.v2_GET()
 
+    async def v2_ensure_transaction_POST(self) -> None:
+        self.locked_probe_data = True
+
     def get_capabilities(self):
         if self.is_core_boot_classic():
             classic_capabilities = []
@@ -772,16 +789,19 @@ class FilesystemController(SubiquityController, FilesystemManipulator):
     async def v2_guided_POST(self, data: GuidedChoiceV2) \
             -> GuidedStorageResponseV2:
         log.debug(data)
+        self.locked_probe_data = True
         await self.guided(data)
         return await self.v2_guided_GET()
 
     async def v2_reformat_disk_POST(self, data: ReformatDisk) \
             -> StorageResponseV2:
+        self.locked_probe_data = True
         self.reformat(self.model._one(id=data.disk_id), data.ptable)
         return await self.v2_GET()
 
     async def v2_add_boot_partition_POST(self, disk_id: str) \
             -> StorageResponseV2:
+        self.locked_probe_data = True
         disk = self.model._one(id=disk_id)
         if boot.is_boot_device(disk):
             raise ValueError('device already has bootloader partition')
@@ -793,6 +813,7 @@ class FilesystemController(SubiquityController, FilesystemManipulator):
     async def v2_add_partition_POST(self, data: AddPartitionV2) \
             -> StorageResponseV2:
         log.debug(data)
+        self.locked_probe_data = True
         if data.partition.boot is not None:
             raise ValueError('add_partition does not support changing boot')
         disk = self.model._one(id=data.disk_id)
@@ -814,6 +835,7 @@ class FilesystemController(SubiquityController, FilesystemManipulator):
     async def v2_delete_partition_POST(self, data: ModifyPartitionV2) \
             -> StorageResponseV2:
         log.debug(data)
+        self.locked_probe_data = True
         disk = self.model._one(id=data.disk_id)
         partition = self.get_partition(disk, data.partition.number)
         self.delete_partition(partition)
@@ -822,6 +844,7 @@ class FilesystemController(SubiquityController, FilesystemManipulator):
     async def v2_edit_partition_POST(self, data: ModifyPartitionV2) \
             -> StorageResponseV2:
         log.debug(data)
+        self.locked_probe_data = True
         disk = self.model._one(id=data.disk_id)
         partition = self.get_partition(disk, data.partition.number)
         if data.partition.size not in (None, partition.size) \
@@ -843,6 +866,15 @@ class FilesystemController(SubiquityController, FilesystemManipulator):
         self.partition_disk_handler(disk, spec, partition=partition)
         return await self.v2_GET()
 
+    async def dry_run_wait_probe_POST(self) -> None:
+        if not self.app.opts.dry_run:
+            raise NotImplementedError
+
+        # This will start the probe task if not yet started.
+        self.ensure_probing()
+
+        await self._probe_task.task
+
     @with_context(name='probe_once', description='restricted={restricted}')
     async def _probe_once(self, *, context, restricted):
         if restricted:
@@ -867,7 +899,11 @@ class FilesystemController(SubiquityController, FilesystemManipulator):
         with open(fpath, 'w') as fp:
             json.dump(storage, fp, indent=4)
         self.app.note_file_for_apport(key, fpath)
-        self.model.load_probe_data(storage)
+        if not self.locked_probe_data:
+            self.queued_probe_data = None
+            self.model.load_probe_data(storage)
+        else:
+            self.queued_probe_data = storage
 
     @with_context()
     async def _probe(self, *, context=None):
@@ -1022,6 +1058,14 @@ class FilesystemController(SubiquityController, FilesystemManipulator):
         loop = asyncio.get_running_loop()
         loop.remove_reader(self._monitor.fileno())
 
+    def ensure_probing(self):
+        try:
+            self._probe_task.start_sync()
+        except TaskAlreadyRunningError:
+            log.debug('Skipping run of Probert - probe run already active')
+        else:
+            log.debug('Triggered Probert run on udev event')
+
     def _udev_event(self):
         cp = run_command(['udevadm', 'settle', '-t', '0'])
         if cp.returncode != 0:
@@ -1038,12 +1082,7 @@ class FilesystemController(SubiquityController, FilesystemManipulator):
         while select.select([self._monitor.fileno()], [], [], 0)[0]:
             action, dev = self._monitor.receive_device()
             log.debug("_udev_event %s %s", action, dev)
-        try:
-            self._probe_task.start_sync()
-        except TaskAlreadyRunningError:
-            log.debug('Skipping run of Probert - probe run already active')
-        else:
-            log.debug('Triggered Probert run on udev event')
+        self.ensure_probing()
 
     def make_autoinstall(self):
         rendered = self.model.render()

subiquity/server/controllers/tests/test_filesystem.py

@@ -27,14 +27,21 @@ from subiquitycore.utils import matching_dicts
 from subiquitycore.tests.util import random_string
 
 from subiquity.common.filesystem import gaps
+from subiquity.common.filesystem.actions import DeviceAction
 from subiquity.common.types import (
+    AddPartitionV2,
     Bootloader,
+    Gap,
+    GapUsable,
     GuidedCapability,
     GuidedChoiceV2,
     GuidedStorageTargetReformat,
     GuidedStorageTargetResize,
     GuidedStorageTargetUseGap,
+    ModifyPartitionV2,
+    Partition,
     ProbeStatus,
+    ReformatDisk,
     SizingPolicy,
     )
 from subiquity.models.filesystem import dehumanize_size
@@ -61,12 +68,16 @@ default_capabilities = [
 
 
 class TestSubiquityControllerFilesystem(IsolatedAsyncioTestCase):
+    MOCK_PREFIX = 'subiquity.server.controllers.filesystem.'
+
     def setUp(self):
         self.app = make_app()
         self.app.opts.bootloader = 'UEFI'
         self.app.report_start_event = mock.Mock()
         self.app.report_finish_event = mock.Mock()
         self.app.prober = mock.Mock()
+        self.app.block_log_dir = '/inexistent'
+        self.app.note_file_for_apport = mock.Mock()
         self.fsc = FilesystemController(app=self.app)
         self.fsc._configured = True
 
@@ -88,6 +99,197 @@ class TestSubiquityControllerFilesystem(IsolatedAsyncioTestCase):
         actual = self.app.prober.get_storage.call_args.args[0]
         self.assertTrue({'defaults', 'os'} <= actual)
 
+    async def test_probe_once_fs_configured(self):
+        self.fsc._configured = True
+        self.fsc.queued_probe_data = None
+        with mock.patch.object(self.fsc.model, 'load_probe_data') as load:
+            await self.fsc._probe_once(restricted=True)
+        self.assertIsNone(self.fsc.queued_probe_data)
+        load.assert_not_called()
+
+    @mock.patch('subiquity.server.controllers.filesystem.open',
+                mock.mock_open())
+    async def test_probe_once_locked_probe_data(self):
+        self.fsc._configured = False
+        self.fsc.locked_probe_data = True
+        self.fsc.queued_probe_data = None
+        self.app.prober.get_storage = mock.Mock(return_value={})
+        with mock.patch.object(self.fsc.model, 'load_probe_data') as load:
+            await self.fsc._probe_once(restricted=True)
+        self.assertEqual(self.fsc.queued_probe_data, {})
+        load.assert_not_called()
+
+    @mock.patch('subiquity.server.controllers.filesystem.open',
+                mock.mock_open())
+    async def test_probe_once_unlocked_probe_data(self):
+        self.fsc._configured = False
+        self.fsc.locked_probe_data = False
+        self.fsc.queued_probe_data = None
+        self.app.prober.get_storage = mock.Mock(return_value={})
+        with mock.patch.object(self.fsc.model, 'load_probe_data') as load:
+            await self.fsc._probe_once(restricted=True)
+        self.assertIsNone(self.fsc.queued_probe_data, {})
+        load.assert_called_once_with({})
+
+    async def test_v2_reset_POST_no_queued_data(self):
+        self.fsc.queued_probe_data = None
+        with mock.patch.object(self.fsc.model, 'load_probe_data') as load:
+            await self.fsc.v2_reset_POST()
+        load.assert_not_called()
+
+    async def test_v2_reset_POST_queued_data(self):
+        self.fsc.queued_probe_data = {}
+        with mock.patch.object(self.fsc.model, 'load_probe_data') as load:
+            await self.fsc.v2_reset_POST()
+        load.assert_called_once_with({})
+
+    async def test_v2_ensure_transaction_POST(self):
+        self.fsc.locked_probe_data = False
+        await self.fsc.v2_ensure_transaction_POST()
+        self.assertTrue(self.fsc.locked_probe_data)
+
+    async def test_v2_reformat_disk_POST(self):
+        self.fsc.locked_probe_data = False
+        with mock.patch.object(self.fsc, 'reformat') as reformat:
+            await self.fsc.v2_reformat_disk_POST(
+                    ReformatDisk(disk_id='dev-sda'))
+        self.assertTrue(self.fsc.locked_probe_data)
+        reformat.assert_called_once()
+
+    @mock.patch(MOCK_PREFIX + 'boot.is_boot_device',
+                mock.Mock(return_value=True))
+    async def test_v2_add_boot_partition_POST_existing_bootloader(self):
+        self.fsc.locked_probe_data = False
+        with mock.patch.object(self.fsc, 'add_boot_disk') as add_boot_disk:
+            with self.assertRaisesRegex(ValueError,
+                                        r'already\ has\ bootloader'):
+                await self.fsc.v2_add_boot_partition_POST('dev-sda')
+        self.assertTrue(self.fsc.locked_probe_data)
+        add_boot_disk.assert_not_called()
+
+    @mock.patch(MOCK_PREFIX + 'boot.is_boot_device',
+                mock.Mock(return_value=False))
+    @mock.patch(MOCK_PREFIX + 'DeviceAction.supported',
+                mock.Mock(return_value=[]))
+    async def test_v2_add_boot_partition_POST_not_supported(self):
+        self.fsc.locked_probe_data = False
+        with mock.patch.object(self.fsc, 'add_boot_disk') as add_boot_disk:
+            with self.assertRaisesRegex(ValueError,
+                                        r'does\ not\ support\ boot'):
+                await self.fsc.v2_add_boot_partition_POST('dev-sda')
+        self.assertTrue(self.fsc.locked_probe_data)
+        add_boot_disk.assert_not_called()
+
+    @mock.patch(MOCK_PREFIX + 'boot.is_boot_device',
+                mock.Mock(return_value=False))
+    @mock.patch(MOCK_PREFIX + 'DeviceAction.supported',
+                mock.Mock(return_value=[DeviceAction.TOGGLE_BOOT]))
+    async def test_v2_add_boot_partition_POST(self):
+        self.fsc.locked_probe_data = False
+        with mock.patch.object(self.fsc, 'add_boot_disk') as add_boot_disk:
+            await self.fsc.v2_add_boot_partition_POST('dev-sda')
+        self.assertTrue(self.fsc.locked_probe_data)
+        add_boot_disk.assert_called_once()
+
+    async def test_v2_add_partition_POST_changing_boot(self):
+        self.fsc.locked_probe_data = False
+        data = AddPartitionV2(
+                disk_id='dev-sda',
+                partition=Partition(
+                    format='ext4',
+                    mount='/',
+                    boot=True,
+                ), gap=Gap(
+                    offset=1 << 20,
+                    size=1000 << 20,
+                    usable=GapUsable.YES,
+                ))
+        with mock.patch.object(self.fsc, 'create_partition') as create_part:
+            with self.assertRaisesRegex(ValueError,
+                                        r'does\ not\ support\ changing\ boot'):
+                await self.fsc.v2_add_partition_POST(data)
+        self.assertTrue(self.fsc.locked_probe_data)
+        create_part.assert_not_called()
+
+    async def test_v2_add_partition_POST_too_large(self):
+        self.fsc.locked_probe_data = False
+        data = AddPartitionV2(
+                disk_id='dev-sda',
+                partition=Partition(
+                    format='ext4',
+                    mount='/',
+                    size=2000 << 20,
+                ), gap=Gap(
+                    offset=1 << 20,
+                    size=1000 << 20,
+                    usable=GapUsable.YES,
+                ))
+        with mock.patch.object(self.fsc, 'create_partition') as create_part:
+            with self.assertRaisesRegex(ValueError, r'too\ large'):
+                await self.fsc.v2_add_partition_POST(data)
+        self.assertTrue(self.fsc.locked_probe_data)
+        create_part.assert_not_called()
+
+    @mock.patch(MOCK_PREFIX + 'gaps.at_offset')
+    async def test_v2_add_partition_POST(self, at_offset):
+        at_offset.split = mock.Mock(return_value=[mock.Mock()])
+        self.fsc.locked_probe_data = False
+        data = AddPartitionV2(
+                disk_id='dev-sda',
+                partition=Partition(
+                    format='ext4',
+                    mount='/',
+                ), gap=Gap(
+                    offset=1 << 20,
+                    size=1000 << 20,
+                    usable=GapUsable.YES,
+                ))
+        with mock.patch.object(self.fsc, 'create_partition') as create_part:
+            await self.fsc.v2_add_partition_POST(data)
+        self.assertTrue(self.fsc.locked_probe_data)
+        create_part.assert_called_once()
+
+    async def test_v2_delete_partition_POST(self):
+        self.fsc.locked_probe_data = False
+        data = ModifyPartitionV2(
+                disk_id='dev-sda',
+                partition=Partition(number=1),
+                )
+        with mock.patch.object(self.fsc, 'delete_partition') as del_part:
+            with mock.patch.object(self.fsc, 'get_partition'):
+                await self.fsc.v2_delete_partition_POST(data)
+        self.assertTrue(self.fsc.locked_probe_data)
+        del_part.assert_called_once()
+
+    async def test_v2_edit_partition_POST_change_boot(self):
+        self.fsc.locked_probe_data = False
+        data = ModifyPartitionV2(
+                disk_id='dev-sda',
+                partition=Partition(number=1, boot=True),
+                )
+        existing = Partition(number=1, size=1000 << 20, boot=False)
+        with mock.patch.object(self.fsc, 'partition_disk_handler') as handler:
+            with mock.patch.object(self.fsc, 'get_partition',
+                                   return_value=existing):
+                with self.assertRaisesRegex(ValueError, r'changing\ boot'):
+                    await self.fsc.v2_edit_partition_POST(data)
+        self.assertTrue(self.fsc.locked_probe_data)
+        handler.assert_not_called()
+
+    async def test_v2_edit_partition_POST(self):
+        self.fsc.locked_probe_data = False
+        data = ModifyPartitionV2(
+                disk_id='dev-sda',
+                partition=Partition(number=1),
+                )
+        existing = Partition(number=1, size=1000 << 20, boot=False)
+        with mock.patch.object(self.fsc, 'partition_disk_handler') as handler:
+            with mock.patch.object(self.fsc, 'get_partition',
+                                   return_value=existing):
+                await self.fsc.v2_edit_partition_POST(data)
+        self.assertTrue(self.fsc.locked_probe_data)
+        handler.assert_called_once()
+
 
 class TestGuided(IsolatedAsyncioTestCase):
     boot_expectations = [

subiquity/tests/api/test_api.py

@@ -1533,6 +1533,49 @@ class TestRegression(TestAPI):
                 self.assertLessEqual(cur['end'] + (1 << 20), nxt['offset'],
                                      f'partition overlap {cur} {nxt}')
 
+    @timeout()
+    async def test_probert_result_during_partitioning(self):
+        '''If a probert run finished during manual partition, we used to
+        load the probing data, essentially discarding changes made by the user
+        so far.  This test creates a new partition, simulates the end of a
+        probert run, and then tries to edit the previously created partition.
+        The edit operation would fail in earlier versions, because the new
+        partition would be discarded.
+        '''
+        cfg = 'examples/simple.json'
+        extra = ['--storage-version', '2']
+        async with start_server(cfg, extra_args=extra) as inst:
+            names = ['locale', 'keyboard', 'source', 'network', 'proxy',
+                     'mirror']
+            await inst.post('/meta/mark_configured', endpoint_names=names)
+            resp = await inst.get('/storage/v2')
+            [d] = resp['disks']
+            [g] = d['partitions']
+            data = {
+                'disk_id': 'disk-sda',
+                'gap': g,
+                'partition': {
+                    'size': -1,
+                    'mount': '/',
+                    'format': 'ext4',
+                }
+            }
+            add_resp = await inst.post('/storage/v2/add_partition', data)
+            [sda] = add_resp['disks']
+            [root] = match(sda['partitions'], mount='/')
+
+            # Now let's make sure we get the results from a probert run to kick
+            # in.
+            await inst.post('/storage/dry_run_wait_probe')
+            data = {
+                'disk_id': 'disk-sda',
+                'partition': {
+                    'number': root['number'],
+                }
+            }
+            # We should be able to modify the created partition.
+            await inst.post('/storage/v2/edit_partition', data)
+
 
 class TestCancel(TestAPI):
     @timeout()