AI-Studio/app/MindWork AI Studio/Tools/DataBackup/DataBackupWriter.cs

using System.IO.Compression;
using System.Security.Cryptography;
using System.Text;
using System.Text.Json;

namespace AIStudio.Tools.DataBackup;

/// <summary>
/// Writes a backup, cf. DataBackupFormat.
/// </summary>
/// <remarks>
/// Every file is streamed into the backup, so a backup of many gigabytes needs no more memory than
/// a small one. The caller writes the backup straight into the file the user chose: in the Flatpak,
/// the portal grants access to exactly that file, and to nothing next to it.
/// </remarks>
public static class DataBackupWriter
{
    private const int BUFFER_SIZE = 64 * 1024;

    /// <summary>
    /// The earliest time a ZIP file can store.
    /// </summary>
    private static readonly DateTimeOffset EARLIEST_ZIP_TIME = new(1980, 1, 1, 0, 0, 0, TimeSpan.Zero);

    /// <summary>
    /// A file to take along, and where it goes in the backup.
    /// </summary>
    private sealed record BackupFile(string EntryName, string SourcePath);

    /// <summary>
    /// Writes a backup into the given stream.
    /// </summary>
    /// <remarks>
    /// A file which is gone by the time it is written, e.g. a temporary chat AI Studio deleted
    /// meanwhile, is left out. A file which cannot be read fails the backup, though: a backup which
    /// quietly lacks something is worse than none.
    /// </remarks>
    /// <param name="target">Where to write the backup. It stays open.</param>
    /// <param name="sources">Where the files come from.</param>
    /// <param name="manifest">What to note in the manifest. The entries are filled in here.</param>
    /// <param name="progress">Receives how many files are written.</param>
    /// <param name="token">Cancels the backup.</param>
    /// <returns>The manifest as the backup holds it.</returns>
    public static async Task<DataBackupManifest> WriteAsync(Stream target, DataBackupSources sources, DataBackupManifest manifest, IProgress<DataBackupProgress>? progress, CancellationToken token)
    {
        // The provenance store follows the chats, cf. DataBackupSources.CopyProvenanceStoreAsync. Until it is copied, it counts as one file:
        var files = CollectFiles(sources, entry => !entry.ComesFromProvenanceCopy, string.Empty, token);
        var totalFiles = files.Count + 2;
        var entries = new List<DataBackupManifestEntry>(totalFiles);
        var buffer = new byte[BUFFER_SIZE];

        progress?.Report(new(DataBackupPhase.WRITING, 0, totalFiles));
        await using (var archive = await ZipArchive.CreateAsync(target, ZipArchiveMode.Create, leaveOpen: true, entryNameEncoding: null, token))
        {
            var settingsEntry = DataBackupContents.ENTRIES.First(entry => entry.Kind is DataBackupEntryKind.SETTINGS);
            await using (var settingsStream = new MemoryStream(Encoding.UTF8.GetBytes(sources.SettingsJson), writable: false))
                entries.Add(await WriteEntryAsync(archive, DataBackupFormat.EntryName(settingsEntry.Location, settingsEntry.RelativePath), settingsStream, DateTimeOffset.Now, buffer, token));

            progress?.Report(new(DataBackupPhase.WRITING, entries.Count, totalFiles));
            await WriteFilesAsync(archive, files, entries, totalFiles, progress, buffer, token);

            var provenanceCopyDirectory = sources.CopyProvenanceStoreAsync is null ? string.Empty : await sources.CopyProvenanceStoreAsync(token);
            var provenanceFiles = CollectFiles(sources, entry => entry.ComesFromProvenanceCopy, provenanceCopyDirectory, token);
            totalFiles = files.Count + 1 + provenanceFiles.Count;
            await WriteFilesAsync(archive, provenanceFiles, entries, totalFiles, progress, buffer, token);

            manifest = manifest with { Entries = entries };
            await using (var manifestStream = new MemoryStream(JsonSerializer.SerializeToUtf8Bytes(manifest, DataBackupFormat.JSON_OPTIONS), writable: false))
                await WriteEntryAsync(archive, DataBackupFormat.MANIFEST_ENTRY_NAME, manifestStream, DateTimeOffset.Now, buffer, token);
        }

        return manifest;
    }

    /// <summary>
    /// Collects the files of the data folder which belong into a backup, cf. DataBackupContents.
    /// </summary>
    /// <remarks>
    /// The settings are no part of it: they come from memory. The indexes and the provenance store
    /// come from their copies.
    /// </remarks>
    /// <param name="sources">Where the files come from.</param>
    /// <param name="include">Which entries of the list to collect.</param>
    /// <param name="provenanceCopyDirectory">The copy of the provenance store, or an empty string to take it from the data folder.</param>
    /// <param name="token">Cancels collecting.</param>
    private static List<BackupFile> CollectFiles(DataBackupSources sources, Func<DataBackupEntry, bool> include, string provenanceCopyDirectory, CancellationToken token)
    {
        var excludedPaths = DataBackupContents.ENTRIES
            .Where(entry => !entry.IsIncluded && entry.Location is DataBackupLocation.DATA_DIRECTORY)
            .Select(entry => entry.RelativePath)
            .Concat(sources.OrganizationPluginFolders)
            .ToList();

        var files = new List<BackupFile>();
        foreach (var entry in DataBackupContents.ENTRIES.Where(entry => entry.IsIncluded && entry.Location is DataBackupLocation.DATA_DIRECTORY && include(entry)))
        {
            var root = entry switch
            {
                { ComesFromIndexCopies: true } => sources.IndexCopiesDirectory,
                { ComesFromProvenanceCopy: true } when !string.IsNullOrWhiteSpace(provenanceCopyDirectory) => provenanceCopyDirectory,
                _ => sources.DataDirectory,
            };

            if (string.IsNullOrWhiteSpace(root))
                continue;

            var path = Path.Join(root, entry.RelativePath.Replace(DataBackupContents.SEPARATOR, Path.DirectorySeparatorChar));
            if (!entry.IsFolder)
            {
                if (File.Exists(path))
                    files.Add(new(DataBackupFormat.EntryName(entry.Location, entry.RelativePath), path));

                continue;
            }

            if (!Directory.Exists(path))
                continue;

            //
            // Hidden files count as well, since on Linux and macOS a dot makes a file hidden. Links
            // are not followed, so nothing outside the folder gets into the backup. A folder which
            // cannot be read fails the backup instead of leaving a gap:
            //
            var options = new EnumerationOptions { RecurseSubdirectories = true, AttributesToSkip = FileAttributes.ReparsePoint, IgnoreInaccessible = false };
            foreach (var filePath in Directory.EnumerateFiles(path, "*", options))
            {
                token.ThrowIfCancellationRequested();

                var relativePath = Path.GetRelativePath(root, filePath).Replace(Path.DirectorySeparatorChar, DataBackupContents.SEPARATOR);
                if (excludedPaths.Any(excludedPath => DataBackupContents.Covers(excludedPath, relativePath)))
                    continue;

                files.Add(new(DataBackupFormat.EntryName(DataBackupLocation.DATA_DIRECTORY, relativePath), filePath));
            }
        }

        return files;
    }

    /// <summary>
    /// Writes files into the backup, and notes each one in the entries of the manifest.
    /// </summary>
    private static async Task WriteFilesAsync(ZipArchive archive, IReadOnlyList<BackupFile> files, List<DataBackupManifestEntry> entries, int totalFiles, IProgress<DataBackupProgress>? progress, byte[] buffer, CancellationToken token)
    {
        foreach (var file in files)
        {
            token.ThrowIfCancellationRequested();

            FileStream source;
            try
            {
                source = new FileStream(file.SourcePath, FileMode.Open, FileAccess.Read, FileShare.ReadWrite | FileShare.Delete, BUFFER_SIZE, FileOptions.Asynchronous | FileOptions.SequentialScan);
            }
            catch (Exception exception) when (exception is FileNotFoundException or DirectoryNotFoundException)
            {
                continue;
            }

            await using (source)
                entries.Add(await WriteEntryAsync(archive, file.EntryName, source, File.GetLastWriteTime(file.SourcePath), buffer, token));

            progress?.Report(new(DataBackupPhase.WRITING, entries.Count, totalFiles));
        }
    }

    /// <summary>
    /// Writes one entry, and notes its size and hash while it streams through.
    /// </summary>
    private static async Task<DataBackupManifestEntry> WriteEntryAsync(ZipArchive archive, string entryName, Stream source, DateTimeOffset lastWriteTime, byte[] buffer, CancellationToken token)
    {
        var compressionLevel = DataBackupFormat.ALREADY_COMPRESSED_EXTENSIONS.Contains(Path.GetExtension(entryName))
            ? CompressionLevel.NoCompression
            : CompressionLevel.Optimal;

        var entry = archive.CreateEntry(entryName, compressionLevel);

        // The time of the last change tells, e.g., the order of the chats in a workspace:
        entry.LastWriteTime = lastWriteTime < EARLIEST_ZIP_TIME ? DateTimeOffset.Now : lastWriteTime;

        using var hash = IncrementalHash.CreateHash(HashAlgorithmName.SHA256);
        long size = 0;
        await using (var entryStream = await entry.OpenAsync(token))
        {
            int read;
            while ((read = await source.ReadAsync(buffer, token)) > 0)
            {
                hash.AppendData(buffer, 0, read);
                await entryStream.WriteAsync(buffer.AsMemory(0, read), token);
                size += read;
            }
        }

        return new(entryName, size, Convert.ToHexStringLower(hash.GetHashAndReset()));
    }
}