From 181a0f6b6abe5328d41130f55133de139cc1b2a7 Mon Sep 17 00:00:00 2001
From: Timothy Baldridge <tbaldridge@gmail.com>
Date: Wed, 4 Mar 2020 22:27:15 -0700
Subject: [PATCH] Bug fixes for virtual memory/stream leaks

---
 Compression.BSA/BA2Builder.cs              |  4 ++++
 Compression.BSA/BSABuilder.cs              |  3 +++
 Compression.BSA/TES3Builder.cs             |  1 +
 Wabbajack.Common/Util/DiskSlabAllocator.cs |  2 +-
 Wabbajack.Lib/MO2Installer.cs              | 10 +++++-----
 5 files changed, 14 insertions(+), 6 deletions(-)

diff --git a/Compression.BSA/BA2Builder.cs b/Compression.BSA/BA2Builder.cs
index 813a1163..7b77e54a 100644
--- a/Compression.BSA/BA2Builder.cs
+++ b/Compression.BSA/BA2Builder.cs
@@ -172,6 +172,7 @@ namespace Compression.BSA
                 using var ms = new MemoryStream();
                 using (var ds = new DeflaterOutputStream(ms))
                 {
+                    ds.IsStreamOwner = false;
                     src.CopyToLimit(ds, (int)chunk.FullSz);
                 }
 
@@ -204,6 +205,7 @@ namespace Compression.BSA
             bw.Write((ulong)pos);
             bw.BaseStream.Position = pos;
             _dataSlab.CopyToLimit(bw.BaseStream, (int)_dataSlab.Length);
+            _dataSlab.Dispose();
         }
     }
 
@@ -230,6 +232,7 @@ namespace Compression.BSA
             {
                 using (var ds = new DeflaterOutputStream(ms))
                 {
+                    ds.IsStreamOwner = false;
                     builder._dataSrc.CopyTo(ds);
                 }
 
@@ -268,6 +271,7 @@ namespace Compression.BSA
             wtr.BaseStream.Position = pos;
             _dataSrc.Position = 0;
             _dataSrc.CopyToLimit(wtr.BaseStream, (int)_dataSrc.Length);
+            _dataSrc.Dispose();
         }
     }
 }
diff --git a/Compression.BSA/BSABuilder.cs b/Compression.BSA/BSABuilder.cs
index 6569be09..628a21df 100644
--- a/Compression.BSA/BSABuilder.cs
+++ b/Compression.BSA/BSABuilder.cs
@@ -298,6 +298,7 @@ namespace Compression.BSA
                     var r = new MemoryStream();
                     using (var w = new DeflaterOutputStream(r))
                     {
+                        w.IsStreamOwner = false;
                         _srcData.CopyTo(w);
                     }
                     _srcData = _bsa._slab.Allocate(r.Length);
@@ -341,11 +342,13 @@ namespace Compression.BSA
                 wtr.Write((uint) _originalSize);
                 _srcData.Position = 0;
                 _srcData.CopyToLimit(wtr.BaseStream, (int)_srcData.Length);
+                _srcData.Dispose();
             }
             else
             {
                 _srcData.Position = 0;
                 _srcData.CopyToLimit(wtr.BaseStream, (int)_srcData.Length);
+                _srcData.Dispose();
             }
         }
     }
diff --git a/Compression.BSA/TES3Builder.cs b/Compression.BSA/TES3Builder.cs
index 2fccc6a1..3f85bd49 100644
--- a/Compression.BSA/TES3Builder.cs
+++ b/Compression.BSA/TES3Builder.cs
@@ -66,6 +66,7 @@ namespace Compression.BSA
             {
                 bw.BaseStream.Position = _state.DataOffset + state.Offset;
                 data.CopyTo(bw.BaseStream);
+                data.Dispose();
             }
         }
 
diff --git a/Wabbajack.Common/Util/DiskSlabAllocator.cs b/Wabbajack.Common/Util/DiskSlabAllocator.cs
index 937e38b1..effdee3b 100644
--- a/Wabbajack.Common/Util/DiskSlabAllocator.cs
+++ b/Wabbajack.Common/Util/DiskSlabAllocator.cs
@@ -18,7 +18,7 @@ namespace Wabbajack.Common
         public DiskSlabAllocator()
         {
             _name = Guid.NewGuid().ToString();
-            _mmap = MemoryMappedFile.CreateNew(_name, (long)1 << 34);
+            _mmap = MemoryMappedFile.CreateNew(null, (long)1 << 34);
         }
 
         public Stream Allocate(long size)
diff --git a/Wabbajack.Lib/MO2Installer.cs b/Wabbajack.Lib/MO2Installer.cs
index 55f97b8d..6a7b23f2 100644
--- a/Wabbajack.Lib/MO2Installer.cs
+++ b/Wabbajack.Lib/MO2Installer.cs
@@ -239,17 +239,17 @@ namespace Wabbajack.Lib
 
                 using (var a = bsa.State.MakeBuilder())
                 {
-                    await bsa.FileStates.PMap(Queue, state =>
+                    var streams = await bsa.FileStates.PMap(Queue, state =>
                     {
                         Status($"Adding {state.Path} to BSA");
-                        using (var fs = File.OpenRead(Path.Combine(sourceDir, state.Path)))
-                        {
-                            a.AddFile(state, fs);
-                        }
+                        var fs = File.OpenRead(Path.Combine(sourceDir, state.Path));
+                        a.AddFile(state, fs);
+                        return fs;
                     });
 
                     Info($"Writing {bsa.To}");
                     a.Build(Path.Combine(OutputFolder, bsa.To));
+                    streams.Do(s => s.Dispose());
                 }
             }