using System; using System.Collections.Generic; using System.Collections.Immutable; using System.IO; using System.Linq; using System.Reactive.Linq; using System.Reactive.Subjects; using System.Text; using System.Threading.Tasks; using Alphaleonis.Win32.Filesystem; using ICSharpCode.SharpZipLib.Zip.Compression.Streams; using Wabbajack.Common; using Wabbajack.Common.StatusFeed.Errors; using Wabbajack.VirtualFileSystem.ExtractedFiles; using Directory = Alphaleonis.Win32.Filesystem.Directory; using File = System.IO.File; using FileInfo = Alphaleonis.Win32.Filesystem.FileInfo; using Path = Alphaleonis.Win32.Filesystem.Path; namespace Wabbajack.VirtualFileSystem { public class Context { private static Task _cleanupTask; static Context() { Utils.Log("Cleaning VFS, this may take a bit of time"); _cleanupTask = Utils.DeleteDirectory(StagingFolder); } public const ulong FileVersion = 0x03; public const string Magic = "WABBAJACK VFS FILE"; private static readonly AbsolutePath StagingFolder = ((RelativePath)"vfs_staging").RelativeToWorkingDirectory(); public IndexRoot Index { get; private set; } = IndexRoot.Empty; /// /// A stream of tuples of ("Update Title", 0.25) which represent the name of the current task /// and the current progress. /// public IObservable<(string, float)> ProgressUpdates => _progressUpdates; private readonly Subject<(string, float)> _progressUpdates = new Subject<(string, float)>(); public StatusUpdateTracker UpdateTracker { get; set; } = new StatusUpdateTracker(1); public WorkQueue Queue { get; } public bool UseExtendedHashes { get; set; } public bool FavorPerfOverRAM { get; set; } public Context(WorkQueue queue, bool extendedHashes = false) { Queue = queue; UseExtendedHashes = extendedHashes; } public static TemporaryDirectory GetTemporaryFolder() { return new TemporaryDirectory(((RelativePath)Guid.NewGuid().ToString()).RelativeTo(StagingFolder)); } public async Task AddRoot(AbsolutePath root) { await _cleanupTask; var filtered = Index.AllFiles.Where(file => file.IsNative && ((AbsolutePath) file.Name).Exists).ToList(); var byPath = filtered.ToImmutableDictionary(f => f.Name); var filesToIndex = root.EnumerateFiles().Distinct().ToList(); var allFiles = await filesToIndex .PMap(Queue, async f => { if (byPath.TryGetValue(f, out var found)) { if (found.LastModified == f.LastModifiedUtc.AsUnixTime() && found.Size == f.Size) return found; } return await VirtualFile.Analyze(this, null, new NativeFileStreamFactory(f), f, 0); }); var newIndex = await IndexRoot.Empty.Integrate(filtered.Concat(allFiles).ToList()); lock (this) { Index = newIndex; } return newIndex; } public async Task AddRoots(List roots) { await _cleanupTask; var native = Index.AllFiles.Where(file => file.IsNative).ToDictionary(file => file.FullPath.Base); var filtered = Index.AllFiles.Where(file => ((AbsolutePath)file.Name).Exists).ToList(); var filesToIndex = roots.SelectMany(root => root.EnumerateFiles()).ToList(); var allFiles = await filesToIndex .PMap(Queue, async f => { Utils.Status($"Indexing {Path.GetFileName((string)f)}"); if (native.TryGetValue(f, out var found)) { if (found.LastModified == f.LastModifiedUtc.AsUnixTime() && found.Size == f.Size) return found; } return await VirtualFile.Analyze(this, null, new NativeFileStreamFactory(f), f, 0); }); var newIndex = await IndexRoot.Empty.Integrate(filtered.Concat(allFiles).ToList()); lock (this) { Index = newIndex; } return newIndex; } class Box { public T Value { get; set; } } private Func, IObservable> ProgressUpdater(string s, float totalCount) { if (totalCount == 0) totalCount = 1; var box = new Box(); return sub => sub.Select(itm => { box.Value += 1; _progressUpdates.OnNext((s, box.Value / totalCount)); return itm; }); } public async Task WriteToFile(AbsolutePath filename) { await using var fs = await filename.Create(); await using var bw = new BinaryWriter(fs, Encoding.UTF8, true); fs.SetLength(0); bw.Write(Encoding.ASCII.GetBytes(Magic)); bw.Write(FileVersion); bw.Write((ulong) Index.AllFiles.Count); await (await Index.AllFiles .PMap(Queue, f => { var ms = new MemoryStream(); using var ibw = new BinaryWriter(ms, Encoding.UTF8, true); f.Write(ibw); return ms; })) .DoAsync(async ms => { var size = ms.Position; ms.Position = 0; bw.Write((ulong) size); await ms.CopyToAsync(fs); }); Utils.Log($"Wrote {fs.Position.ToFileSizeString()} file as vfs cache file {filename}"); } public async Task IntegrateFromFile(AbsolutePath filename) { try { await using var fs = await filename.OpenRead(); using var br = new BinaryReader(fs, Encoding.UTF8, true); var magic = Encoding.ASCII.GetString(br.ReadBytes(Encoding.ASCII.GetBytes(Magic).Length)); var fileVersion = br.ReadUInt64(); if (fileVersion != FileVersion || magic != Magic) throw new InvalidDataException("Bad Data Format"); var numFiles = br.ReadUInt64(); var files = Enumerable.Range(0, (int) numFiles) .Select(idx => { var size = br.ReadUInt64(); var bytes = new byte[size]; br.BaseStream.Read(bytes, 0, (int) size); return VirtualFile.Read(this, bytes); }).ToList(); var newIndex = await Index.Integrate(files); lock (this) { Index = newIndex; } } catch (IOException) { await filename.DeleteAsync(); } } /// /// Extract the set of files and call the callback for each, handing it a stream factory and the virtual file, /// top level archives (native archives) will be processed in parallel. Duplicate files will not be /// /// < /// /// /// public async Task Extract(WorkQueue queue, HashSet files, Func callback, AbsolutePath? tempFolder = null, StatusUpdateTracker updateTracker = null) { var top = new VirtualFile(); var filesByParent = files.SelectMany(f => f.FilesInFullPath) .Distinct() .GroupBy(f => f.Parent ?? top) .ToDictionary(f => f.Key); async Task HandleFile(VirtualFile file, IExtractedFile sfn) { if (filesByParent.ContainsKey(file)) sfn.CanMove = false; if (files.Contains(file)) await callback(file, sfn); if (filesByParent.TryGetValue(file, out var children)) { var fileNames = children.ToDictionary(c => c.RelativeName); try { await FileExtractor2.GatheringExtract(queue, sfn, r => fileNames.ContainsKey(r), async (rel, csf) => { await HandleFile(fileNames[rel], csf); return 0; }, tempFolder: tempFolder, onlyFiles: fileNames.Keys.ToHashSet()); } catch (_7zipReturnError) { await using var stream = await sfn.GetStream(); var hash = await stream.xxHashAsync(); if (hash != file.Hash) { throw new Exception($"File {file.FullPath} is corrupt, please delete it and retry the installation"); } throw; } } } updateTracker ??= new StatusUpdateTracker(1); await filesByParent[top].PMap(queue, updateTracker, async file => await HandleFile(file, new ExtractedNativeFile(file.AbsoluteName) {CanMove = false})); } #region KnownFiles private List _knownFiles = new List(); private Dictionary _knownArchives = new Dictionary(); public void AddKnown(IEnumerable known, Dictionary archives) { _knownFiles.AddRange(known); foreach (var (key, value) in archives) _knownArchives.TryAdd(key, value); } public async Task BackfillMissing() { var newFiles = _knownArchives.ToDictionary(kv => kv.Key, kv => new VirtualFile { Name = kv.Value, Size = kv.Value.Size, Hash = kv.Key, }); newFiles.Values.Do(f => f.FillFullPath(0)); var parentchild = new Dictionary<(VirtualFile, RelativePath), VirtualFile>(); void BackFillOne(HashRelativePath file) { var parent = newFiles[file.BaseHash]; foreach (var path in file.Paths) { if (parentchild.TryGetValue((parent, path), out var foundParent)) { parent = foundParent; continue; } var nf = new VirtualFile {Name = path, Parent = parent}; nf.FillFullPath(); parent.Children = parent.Children.Add(nf); parentchild.Add((parent, path), nf); parent = nf; } } _knownFiles.Where(f => f.Paths.Length > 0).Do(BackFillOne); var newIndex = await Index.Integrate(newFiles.Values.ToList()); lock (this) Index = newIndex; _knownFiles = new List(); } #endregion } public class DisposableList : List, IDisposable { private Action _unstage; public DisposableList(Action unstage, IEnumerable files) : base(files) { _unstage = unstage; } public void Dispose() { _unstage(); } } public class AsyncDisposableList : List, IAsyncDisposable { private Func _unstage; public AsyncDisposableList(Func unstage, IEnumerable files) : base(files) { _unstage = unstage; } public async ValueTask DisposeAsync() { await _unstage(); } } public class IndexRoot { public static IndexRoot Empty = new IndexRoot(); public IndexRoot(ImmutableList aFiles, Dictionary byFullPath, ImmutableDictionary> byHash, ImmutableDictionary byRoot, ImmutableDictionary> byName) { AllFiles = aFiles; ByFullPath = byFullPath; ByHash = byHash; ByRootPath = byRoot; ByName = byName; } public IndexRoot() { AllFiles = ImmutableList.Empty; ByFullPath = new Dictionary(); ByHash = ImmutableDictionary>.Empty; ByRootPath = ImmutableDictionary.Empty; ByName = ImmutableDictionary>.Empty; } public ImmutableList AllFiles { get; } public Dictionary ByFullPath { get; } public ImmutableDictionary> ByHash { get; } public ImmutableDictionary> ByName { get; set; } public ImmutableDictionary ByRootPath { get; } public async Task Integrate(ICollection files) { Utils.Log($"Integrating {files.Count} files"); var allFiles = AllFiles.Concat(files) .OrderByDescending(f => f.LastModified) .GroupBy(f => f.FullPath).Select(g => g.Last()) .ToImmutableList(); var byFullPath = Task.Run(() => allFiles.SelectMany(f => f.ThisAndAllChildren) .ToDictionary(f => f.FullPath)); var byHash = Task.Run(() => allFiles.SelectMany(f => f.ThisAndAllChildren) .Where(f => f.Hash != Hash.Empty) .ToGroupedImmutableDictionary(f => f.Hash)); var byName = Task.Run(() => allFiles.SelectMany(f => f.ThisAndAllChildren) .ToGroupedImmutableDictionary(f => f.Name)); var byRootPath = Task.Run(() => allFiles.ToImmutableDictionary(f => f.AbsoluteName)); var result = new IndexRoot(allFiles, await byFullPath, await byHash, await byRootPath, await byName); Utils.Log($"Done integrating"); return result; } public VirtualFile FileForArchiveHashPath(HashRelativePath argArchiveHashPath) { var cur = ByHash[argArchiveHashPath.BaseHash].First(f => f.Parent == null); return argArchiveHashPath.Paths.Aggregate(cur, (current, itm) => ByName[itm].First(f => f.Parent == current)); } } public class TemporaryDirectory : IAsyncDisposable { public TemporaryDirectory(AbsolutePath name) { FullName = name; if (!FullName.Exists) FullName.CreateDirectory(); } public AbsolutePath FullName { get; } public async ValueTask DisposeAsync() { if (FullName.Exists) await Utils.DeleteDirectory(FullName); } } }