Desktop: let a lock pressed during an unlock win, instead of the unlock finishing behind it

src/DodoSSH.Client.Shell/ViewModels/MainWindowViewModel.cs

@@ -2963,13 +2963,24 @@ internal sealed partial class MainWindowViewModel : ObservableObject, IAsyncDisp
     /// </remarks>
     private async Task AdoptAsync(VaultSession session, CancellationToken cancellationToken)
     {
+        // ◆ Which lock this unlock started under. Every await below is a moment the user can press Lock, and
+        // LockAsync can only close what has been handed over by then — so an unlock still attaching its
+        // stores used to finish after the lock and open the vault behind the lock screen.
+        var generation = locks;
+
         await AttachStoresAsync(session, cancellationToken).ConfigureAwait(true);
 
         // Before the vault view model, because that is what reads it — and read at all rather than defaulted
         // to "everything shown", because a vault somebody set aside last week should still be set aside.
         visibility = await VaultVisibility.LoadAsync(session, cancellationToken).ConfigureAwait(true);
 
-        Vault = new VaultViewModel(
+        if (generation != locks)
+        {
+            await AbandonAsync(session).ConfigureAwait(true);
+            return;
+        }
+
+        var opened = Vault = new VaultViewModel(
             session,
             workspace,
             knownHosts,
@@ -2993,9 +3004,30 @@ internal sealed partial class MainWindowViewModel : ObservableObject, IAsyncDisp
         // are both !CanRegisterDevice, and only the second has anything to withdraw.
         CanForgetDevice = session.Profile.DeviceWrappedPrivateKey is not null;
 
-        await Vault.LoadAsync(cancellationToken).ConfigureAwait(true);
+        await OpenAsync(opened, generation, cancellationToken).ConfigureAwait(true);
+    }
 
-        AfterTheFirstLoad(Vault);
+    /// <summary>Loads the vault just handed over, and starts it syncing — unless a lock has closed it since.</summary>
+    /// <remarks>
+    /// From here the vault belongs to the shell, so a lock closes it itself: there is nothing to undo, only
+    /// nothing more to do to a vault that is gone. Before these checks, a lock pressed while the first load
+    /// ran left this to carry on into a disposed session and a vault the shell no longer held.
+    /// </remarks>
+    private async Task OpenAsync(VaultViewModel opened, int generation, CancellationToken cancellationToken)
+    {
+        if (generation != locks)
+        {
+            return;
+        }
+
+        await opened.LoadAsync(cancellationToken).ConfigureAwait(true);
+
+        if (generation != locks)
+        {
+            return;
+        }
+
+        AfterTheFirstLoad(opened);
 
         // After the first load, so the list is on screen before anything talks to a server. The loop is
         // started from the UI thread deliberately: every pass resumes here, which is what keeps the
@@ -3010,7 +3042,21 @@ internal sealed partial class MainWindowViewModel : ObservableObject, IAsyncDisp
         // and unlocking must never wait on one. Everything the unlock screen promises about working
         // offline stops being true the moment the passphrase leads to a socket. So the vault opens, and
         // the titlebar says OFFLINE until the round trip this starts has an answer.
-        Vault.StartAutoSync();
+        opened.StartAutoSync();
+    }
+
+    /// <summary>Puts away an unlock that a lock overtook before the vault was handed over.</summary>
+    /// <remarks>
+    /// Everything <see cref="LockAsync"/> would have closed had it been here to close: the stores attached to
+    /// this session, the switches read from it, and the session itself, whose keys nothing else now holds.
+    /// </remarks>
+    private async Task AbandonAsync(VaultSession session)
+    {
+        knownHosts.Close();
+        connectionLog.Close();
+        visibility = null;
+
+        await session.DisposeAsync().ConfigureAwait(true);
     }
 
     /// <summary>Everything an unlock does that needs the vault's lists to have been read.</summary>
@@ -3272,6 +3318,10 @@ internal sealed partial class MainWindowViewModel : ObservableObject, IAsyncDisp
     [RelayCommand]
     private async Task LockAsync()
     {
+        // Counted first, so an unlock still on its way in sees it at its next step and stands down rather than
+        // opening the vault after this has run. See AdoptAsync.
+        locks++;
+
         // First, and before the session it read from goes: a synchronisation pass may be in flight, and it
         // ends by refreshing this store. Detaching now makes that refresh a no-op instead of a set of pins
         // reappearing behind a lock screen.
@@ -3319,6 +3369,9 @@ internal sealed partial class MainWindowViewModel : ObservableObject, IAsyncDisp
     [ObservableProperty]
     private bool isConfirmingSignOut;
 
+    /// <summary>How many times this shell has been locked, for an unlock to notice it was overtaken.</summary>
+    private int locks;
+
     /// <summary>
     /// What signing out costs, on this machine, right now.
     /// </summary>

tests/DodoSSH.Client.App.Tests/ShellFlowTests.cs

@@ -639,6 +639,36 @@ public sealed partial class ShellFlowTests : IAsyncLifetime
         server.LiveRowCount.ShouldBe(1, "a save should reach the server without pressing Sync");
     }
 
+    /// <remarks>
+    /// Lock pressed while an unlock is still loading. The unlock is held at the point the vault has been handed
+    /// over and is about to be read — the device-key probe — and the lock lands there. The unlock used to carry
+    /// on regardless, into a vault the lock had disposed and a shell that no longer held one, and put an
+    /// exception's message on the lock screen. An earlier lock, before the hand-over, would have left the vault
+    /// opening behind the lock screen; both are the same check.
+    /// </remarks>
+    [Fact]
+    public async Task LockingWhileAnUnlockIsStillLoading_LeavesItLocked()
+    {
+        await EnrolledAndConfirmedAsync();
+
+        var release = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
+        deviceKeys.HoldAvailability = release.Task;
+
+        shell.Passphrase = Passphrase;
+        var unlocking = shell.UnlockCommand.ExecuteAsync(null);
+
+        await EventuallyAsync(() => shell.Vault is not null, "the unlock should reach the hand-over and wait there");
+
+        await shell.LockCommand.ExecuteAsync(null);
+
+        release.SetResult();
+        await unlocking;
+
+        shell.State.ShouldBe(ShellState.Locked);
+        shell.Vault.ShouldBeNull();
+        shell.StatusMessage.ShouldBe("Locked.");
+    }
+
     [Fact]
     public async Task AnAutomaticPass_SaysNothingWhenThereIsNothingToDo()
     {

tests/DodoSSH.Client.Session.Tests/FakeDeviceKeyStore.cs

@@ -25,8 +25,18 @@ internal sealed class FakeDeviceKeyStore : IDeviceKeyStore
     internal void Overwrite(byte[] replacement) => key = replacement;
 
     /// <inheritdoc />
-    public ValueTask<bool> IsAvailableAsync(CancellationToken cancellationToken) =>
-        ValueTask.FromResult(IsAvailable);
+    public async ValueTask<bool> IsAvailableAsync(CancellationToken cancellationToken)
+    {
+        if (HoldAvailability is { } hold)
+        {
+            await hold.WaitAsync(cancellationToken).ConfigureAwait(false);
+        }
+
+        return IsAvailable;
+    }
+
+    /// <summary>When set, availability is not answered until it completes: an unlock held part-way.</summary>
+    internal Task? HoldAvailability { get; set; }
 
     /// <inheritdoc />
     public ValueTask SaveAsync(ReadOnlyMemory<byte> devicePrivateKey, CancellationToken cancellationToken)