zcbot/windows-node/Zcbot.WindowsNode/NodeDataRoot.cs

using Microsoft.Win32;
using System.Security;
using System.Security.Cryptography;

namespace Zcbot.WindowsNode;

internal sealed record NodeDataRootSelection(
    string RootDirectory,
    bool IsEnvironmentManaged,
    string Source);

internal sealed record NodeDataMigrationResult(
    string SourceDirectory,
    string TargetDirectory,
    int FileCount,
    long TotalBytes);

internal static class NodeDataRootSettings
{
    internal const string EnvironmentVariableName = "ZCBOT_WINDOWS_NODE_DATA_DIR";
    internal const string RegistryPath = @"Software\Zcbot\WindowsNode";
    internal const string RegistryValueName = "DataRoot";

    internal static NodeDataRootSelection Resolve()
    {
        var managed = Environment.GetEnvironmentVariable(EnvironmentVariableName)?.Trim();
        if (!string.IsNullOrWhiteSpace(managed))
        {
            return new NodeDataRootSelection(
                Normalize(managed), IsEnvironmentManaged: true, "environment");
        }

        using var key = Registry.CurrentUser.OpenSubKey(RegistryPath, writable: false);
        var configured = (key?.GetValue(RegistryValueName) as string)?.Trim();
        if (!string.IsNullOrWhiteSpace(configured))
        {
            return new NodeDataRootSelection(
                Normalize(configured), IsEnvironmentManaged: false, "user");
        }

        return new NodeDataRootSelection(DefaultRoot(), IsEnvironmentManaged: false, "default");
    }

    internal static void SaveUserRoot(string root)
    {
        var normalized = Normalize(root);
        using var key = Registry.CurrentUser.CreateSubKey(RegistryPath, writable: true);
        key.SetValue(RegistryValueName, normalized, RegistryValueKind.String);
    }

    internal static string Normalize(string root)
    {
        if (string.IsNullOrWhiteSpace(root) || !Path.IsPathFullyQualified(root))
        {
            throw new NodeConfigurationException("Data root must be an absolute local path.");
        }
        var normalized = Path.TrimEndingDirectorySeparator(Path.GetFullPath(root.Trim()));
        if (normalized.StartsWith(@"\\", StringComparison.Ordinal))
        {
            throw new NodeConfigurationException("Data root must be on a local drive.");
        }
        if (string.Equals(normalized, Path.GetPathRoot(normalized), StringComparison.OrdinalIgnoreCase))
        {
            throw new NodeConfigurationException("Data root cannot be the root of a drive.");
        }
        return normalized;
    }

    private static string DefaultRoot() => Path.Combine(
        Environment.GetFolderPath(Environment.SpecialFolder.CommonApplicationData),
        "Zcbot", "WindowsNode");
}

internal static class NodeDataMigrator
{
    internal static Task<NodeDataMigrationResult> MigrateAsync(
        string sourceRoot,
        string targetRoot,
        CancellationToken cancellationToken) => Task.Run(
            () => Migrate(sourceRoot, targetRoot, cancellationToken), cancellationToken);

    private static NodeDataMigrationResult Migrate(
        string sourceRoot,
        string targetRoot,
        CancellationToken cancellationToken)
    {
        var source = NodeDataRootSettings.Normalize(sourceRoot);
        var target = NodeDataRootSettings.Normalize(targetRoot);
        if (string.Equals(source, target, StringComparison.OrdinalIgnoreCase))
        {
            throw new NodeConfigurationException("The new data root is the same as the current data root.");
        }
        EnsureSeparateTrees(source, target);
        var targetDrive = EnsureLocalFixedDrive(target);

        if (Directory.Exists(target) && Directory.EnumerateFileSystemEntries(target).Any())
        {
            throw new NodeConfigurationException("The selected data root must be empty.");
        }
        if (File.Exists(target))
        {
            throw new NodeConfigurationException("The selected data root is an existing file.");
        }

        var requiredBytes = Directory.Exists(source)
            ? EnumerateFiles(source)
                .Sum(path => new FileInfo(path).Length)
            : 0L;
        if (targetDrive.AvailableFreeSpace < requiredBytes)
        {
            throw new NodeConfigurationException(
                $"The target drive needs at least {requiredBytes} free bytes for migration.");
        }

        var parent = Directory.GetParent(target)?.FullName
            ?? throw new NodeConfigurationException("The selected data root has no parent directory.");
        Directory.CreateDirectory(parent);
        var staging = Path.Combine(parent, $".{Path.GetFileName(target)}.migration-{Guid.NewGuid():N}");
        Directory.CreateDirectory(staging);
        NodeConfigStore.RestrictDirectory(staging);

        try
        {
            var copied = Directory.Exists(source)
                ? CopyTree(source, staging, cancellationToken)
                : (FileCount: 0, TotalBytes: 0L);
            if (Directory.Exists(source))
            {
                VerifyTree(source, staging, cancellationToken);
            }
            if (Directory.Exists(target))
            {
                Directory.Delete(target, recursive: false);
            }
            Directory.Move(staging, target);
            return new NodeDataMigrationResult(source, target, copied.FileCount, copied.TotalBytes);
        }
        catch
        {
            try
            {
                if (Directory.Exists(staging))
                {
                    Directory.Delete(staging, recursive: true);
                }
            }
            catch (IOException)
            {
            }
            catch (UnauthorizedAccessException)
            {
            }
            throw;
        }
    }

    private static (int FileCount, long TotalBytes) CopyTree(
        string source,
        string target,
        CancellationToken cancellationToken)
    {
        var fileCount = 0;
        long totalBytes = 0;
        foreach (var directory in EnumerateDirectories(source))
        {
            cancellationToken.ThrowIfCancellationRequested();
            Directory.CreateDirectory(Path.Combine(target, Path.GetRelativePath(source, directory)));
        }
        foreach (var file in EnumerateFiles(source))
        {
            cancellationToken.ThrowIfCancellationRequested();
            var destination = Path.Combine(target, Path.GetRelativePath(source, file));
            Directory.CreateDirectory(Path.GetDirectoryName(destination)!);
            File.Copy(file, destination, overwrite: false);
            var length = new FileInfo(file).Length;
            fileCount++;
            totalBytes += length;
        }
        return (fileCount, totalBytes);
    }

    private static void VerifyTree(
        string source,
        string target,
        CancellationToken cancellationToken)
    {
        var sourceFiles = EnumerateFiles(source)
            .Select(path => Path.GetRelativePath(source, path))
            .OrderBy(path => path, StringComparer.OrdinalIgnoreCase)
            .ToArray();
        var targetFiles = EnumerateFiles(target)
            .Select(path => Path.GetRelativePath(target, path))
            .OrderBy(path => path, StringComparer.OrdinalIgnoreCase)
            .ToArray();
        if (!sourceFiles.SequenceEqual(targetFiles, StringComparer.OrdinalIgnoreCase))
        {
            throw new IOException("Data root migration verification found a different file set.");
        }
        foreach (var relativePath in sourceFiles)
        {
            cancellationToken.ThrowIfCancellationRequested();
            var sourcePath = Path.Combine(source, relativePath);
            var targetPath = Path.Combine(target, relativePath);
            if (new FileInfo(sourcePath).Length != new FileInfo(targetPath).Length
                || !HashFile(sourcePath).SequenceEqual(HashFile(targetPath)))
            {
                throw new IOException($"Data root migration verification failed: {relativePath}");
            }
        }
    }

    private static byte[] HashFile(string path)
    {
        using var stream = new FileStream(
            path, FileMode.Open, FileAccess.Read, FileShare.Read, 1024 * 1024,
            FileOptions.SequentialScan);
        return SHA256.HashData(stream);
    }

    private static IEnumerable<string> EnumerateDirectories(string root)
    {
        var pending = new Stack<string>();
        pending.Push(root);
        while (pending.Count > 0)
        {
            var current = pending.Pop();
            foreach (var directory in Directory.EnumerateDirectories(
                current, "*", SearchOption.TopDirectoryOnly))
            {
                RejectReparsePoint(directory);
                pending.Push(directory);
                yield return directory;
            }
        }
    }

    private static IEnumerable<string> EnumerateFiles(string root)
    {
        RejectReparsePoint(root);
        foreach (var file in Directory.EnumerateFiles(root, "*", SearchOption.TopDirectoryOnly))
        {
            RejectReparsePoint(file);
            yield return file;
        }
        foreach (var directory in EnumerateDirectories(root))
        {
            foreach (var file in Directory.EnumerateFiles(
                directory, "*", SearchOption.TopDirectoryOnly))
            {
                RejectReparsePoint(file);
                yield return file;
            }
        }
    }

    private static void EnsureSeparateTrees(string source, string target)
    {
        var separator = Path.DirectorySeparatorChar.ToString();
        var sourcePrefix = source + separator;
        var targetPrefix = target + separator;
        if (target.StartsWith(sourcePrefix, StringComparison.OrdinalIgnoreCase)
            || source.StartsWith(targetPrefix, StringComparison.OrdinalIgnoreCase))
        {
            throw new NodeConfigurationException(
                "The old and new data roots cannot contain one another.");
        }
    }

    private static DriveInfo EnsureLocalFixedDrive(string target)
    {
        var root = Path.GetPathRoot(target)
            ?? throw new NodeConfigurationException("The selected data root has no drive.");
        var drive = new DriveInfo(root);
        if (!drive.IsReady || drive.DriveType != DriveType.Fixed)
        {
            throw new NodeConfigurationException("Data root must be on a fixed local drive.");
        }
        return drive;
    }

    private static void RejectReparsePoint(string path)
    {
        if ((File.GetAttributes(path) & FileAttributes.ReparsePoint) != 0)
        {
            throw new IOException($"Data root migration does not follow links: {path}");
        }
    }
}