Add option to dispose base streams

hactoolnet/Program.cs

@@ -29,7 +29,7 @@ namespace hactoolnet
 
             using (var nax0 = Nax0.CreateFromPath(keyset, args[2], args[3]))
             {
-                var nca = new Nca(keyset, nax0.Stream);
+                var nca = new Nca(keyset, nax0.Stream, true);
 
                 using (var output = new FileStream(args[4], FileMode.Create))
                 using (var progress = new ProgressBar())
@@ -54,12 +54,15 @@ namespace hactoolnet
             }
 
             keyset.SetSdSeed(args[1].ToBytes());
-            var sdfs = new SdFs(keyset, args[2]);
+            using (var sdfs = new SdFs(keyset, args[2]))
-            var ncas = sdfs.ReadAllNca();
-
-            foreach (Nca nca in ncas)
             {
-                Console.WriteLine($"{nca.Header.TitleId:X16} {nca.Header.ContentType.ToString().PadRight(10, ' ')} {nca.Name}");
+                sdfs.OpenAllNcas();
+
+                foreach (Nca nca in sdfs.Ncas)
+                {
+                    Console.WriteLine(
+                        $"{nca.Header.TitleId:X16} {nca.Header.ContentType.ToString().PadRight(10, ' ')} {nca.Name}");
+                }
             }
         }
 
@@ -67,43 +70,48 @@ namespace hactoolnet
         {
             var keyset = ExternalKeys.ReadKeyFile(args[0]);
             keyset.SetSdSeed(args[1].ToBytes());
-            var sdfs = new SdFs(keyset, args[2]);
+            List<Nca> ncas;
-            var ncas = sdfs.ReadAllNca().ToArray();
+            using (var sdfs = new SdFs(keyset, args[2]))
-
-            var metadata = new List<Cnmt>();
-
-            using (var progress = new ProgressBar())
             {
-                foreach (var nca in ncas.Where(x => x.Header.ContentType == ContentType.Meta))
+                sdfs.OpenAllNcas();
+                ncas = sdfs.Ncas;
+
+                var metadata = new List<Cnmt>();
+
+                using (var progress = new ProgressBar())
                 {
-                    foreach (var section in nca.Sections.Where(x => x.Header.FsType == SectionFsType.Pfs0))
+                    foreach (var nca in ncas.Where(x => x.Header.ContentType == ContentType.Meta))
                     {
-                        var sect = nca.OpenSection(section.SectionNum);
+                        foreach (var section in nca.Sections.Where(x => x.Header.FsType == SectionFsType.Pfs0))
-                        var pfs0 = sect.Pfs0;
-                        pfs0.Open(sect.Stream);
-
-                        foreach (var entry in pfs0.Entries)
                         {
-                            var path = Path.Combine("meta", entry.Name);
+                            var sect = nca.OpenSection(section.SectionNum);
-                            Directory.CreateDirectory(Path.GetDirectoryName(path));
+                            var pfs0 = new Pfs0(sect);
-                            var file = pfs0.GetFile(entry.Index);
 
-                            metadata.Add(new Cnmt(new MemoryStream(file)));
+                            foreach (var entry in pfs0.Entries)
-                            File.WriteAllBytes(path, file);
+                            {
+                                var path = Path.Combine("meta", entry.Name);
+                                Directory.CreateDirectory(Path.GetDirectoryName(path));
+                                var file = pfs0.GetFile(entry.Index);
+
+                                metadata.Add(new Cnmt(new MemoryStream(file)));
+                                File.WriteAllBytes(path, file);
+                            }
                         }
                     }
-                }
 
-                foreach (var meta in metadata.OrderBy(x => x.TitleId))
+                    foreach (var meta in metadata.OrderBy(x => x.TitleId))
-                {
-                   // progress.LogMessage($"{meta.TitleId:X16} v{meta.TitleVersion} {meta.Type}");
-
-                    foreach (var content in meta.ContentEntries)
                     {
-                        // Add an actual hexdump function
+                        progress.LogMessage($"{meta.TitleId:X16} v{meta.TitleVersion} {meta.Type}");
-                       // progress.LogMessage($"    {BitConverter.ToString(content.NcaId).Replace("-", "")}.nca {content.Type}");
+
+                        foreach (var content in meta.ContentEntries)
+                        {
+                            // Add an actual hexdump function
+                            progress.LogMessage(
+                                $"    {BitConverter.ToString(content.NcaId).Replace("-", "")}.nca {content.Type}");
+                        }
+
+                        progress.LogMessage("");
                     }
-                    //progress.LogMessage("");
                 }
             }
         }

libhac/Nax0.cs

@@ -26,7 +26,7 @@ namespace libhac
 
             stream.Position = 0x4000;
             var xts = XtsAes128.Create(Keys[0], Keys[1]);
-            Stream = new RandomAccessSectorStream(new XtsSectorStream(stream, xts, 0x4000, 0x4000));
+            Stream = new RandomAccessSectorStream(new XtsSectorStream(stream, xts, 0x4000, 0x4000), keepOpen);
         }
 
         private void ReadHeader(Stream stream)

libhac/Nca.cs

@@ -5,7 +5,7 @@ using libhac.XTSSharp;
 
 namespace libhac
 {
-    public class Nca
+    public class Nca : IDisposable
     {
         public NcaHeader Header { get; private set; }
         public string Name { get; set; }
@@ -13,13 +13,16 @@ namespace libhac
         public int CryptoType { get; private set; }
         public byte[][] DecryptedKeys { get; } = Util.CreateJaggedArray<byte[][]>(4, 0x10);
         public Stream Stream { get; private set; }
+        private bool KeepOpen { get; }
 
         public List<NcaSection> Sections = new List<NcaSection>();
 
-        public Nca(Keyset keyset, Stream stream)
+        public Nca(Keyset keyset, Stream stream, bool keepOpen)
         {
+            stream.Position = 0;
+            KeepOpen = keepOpen;
             Stream = stream;
-            ReadHeader(keyset, stream);
+            DecryptHeader(keyset, stream);
 
             CryptoType = Math.Max(Header.CryptoType, Header.CryptoType2);
             if (CryptoType > 0) CryptoType--;
@@ -33,25 +36,24 @@ namespace libhac
 
             for (int i = 0; i < 4; i++)
             {
-                var section = ParseSection(keyset, stream, i);
+                var section = ParseSection(i);
                 if (section != null) Sections.Add(section);
             }
         }
 
-        public NcaSection OpenSection(int index)
+        public Stream OpenSection(int index)
         {
             if (index >= Sections.Count) throw new ArgumentOutOfRangeException(nameof(index));
             var sect = Sections[index];
 
-            sect.Stream = null;
             long offset = sect.Offset;
             long size = sect.Size;
 
             switch (sect.Header.FsType)
             {
                 case SectionFsType.Pfs0:
-                    offset = sect.Offset + sect.Pfs0.Superblock.Pfs0Offset;
+                    offset = sect.Offset + sect.Pfs0.Pfs0Offset;
-                    size = sect.Pfs0.Superblock.Pfs0Size;
+                    size = sect.Pfs0.Pfs0Size;
                     break;
                 case SectionFsType.Romfs:
                     break;
@@ -68,27 +70,28 @@ namespace libhac
                 case SectionCryptType.XTS:
                     break;
                 case SectionCryptType.CTR:
-                    sect.Stream = new RandomAccessSectorStream(new AesCtrStream(Stream, DecryptedKeys[2], offset, size, offset));
+                    return new RandomAccessSectorStream(new AesCtrStream(Stream, DecryptedKeys[2], offset, size, offset), false);
-                    break;
                 case SectionCryptType.BKTR:
                     break;
                 default:
                     throw new ArgumentOutOfRangeException();
             }
 
-            return sect;
+            return Stream;
         }
 
-        private void ReadHeader(Keyset keyset, Stream stream)
+        private void DecryptHeader(Keyset keyset, Stream stream)
         {
-            stream.Position = 0;
+            byte[] headerBytes = new byte[0xC00];
             var xts = XtsAes128.Create(keyset.header_key);
-            var headerDec = new RandomAccessSectorStream(new XtsSectorStream(stream, xts, 0x200));
+            using (var headerDec = new RandomAccessSectorStream(new XtsSectorStream(stream, xts, 0x200)))
-            var reader = new BinaryReader(headerDec);
+            {
+                headerDec.Read(headerBytes, 0, headerBytes.Length);
+            }
+
+            var reader = new BinaryReader(new MemoryStream(headerBytes));
 
             Header = NcaHeader.Read(reader);
-
-            headerDec.Close();
         }
 
         private void DecryptKeyArea(Keyset keyset)
@@ -100,7 +103,7 @@ namespace libhac
             }
         }
 
-        private NcaSection ParseSection(Keyset keyset, Stream stream, int index)
+        private NcaSection ParseSection(int index)
         {
             var entry = Header.SectionEntries[index];
             var header = Header.FsHeaders[index];
@@ -116,11 +119,19 @@ namespace libhac
 
             if (sect.Type == SectionType.Pfs0)
             {
-                sect.Pfs0 = new Pfs0 { Superblock = header.Pfs0 };
+                sect.Pfs0 = header.Pfs0;
             }
 
             return sect;
         }
+
+        public void Dispose()
+        {
+            if (!KeepOpen)
+            {
+                Stream?.Dispose();
+            }
+        }
     }
 
     public class NcaSection
@@ -132,6 +143,6 @@ namespace libhac
         public long Offset { get; set; }
         public long Size { get; set; }
 
-        public Pfs0 Pfs0 { get; set; }
+        public Pfs0Superblock Pfs0 { get; set; }
     }
 }

libhac/Pfs0.cs

@@ -5,20 +5,19 @@ namespace libhac
 {
     public class Pfs0
     {
-        public Pfs0Superblock Superblock { get; set; }
         public Pfs0Header Header { get; set; }
         public int HeaderSize { get; set; }
         public Pfs0FileEntry[] Entries { get; set; }
         private Stream Stream { get; set; }
 
-        public void Open(Stream file)
+        public Pfs0(Stream stream)
         {
             byte[] headerBytes;
-            using (var reader = new BinaryReader(file, Encoding.Default, true))
+            using (var reader = new BinaryReader(stream, Encoding.Default, true))
             {
                 Header = new Pfs0Header(reader);
                 HeaderSize = (int)(16 + 24 * Header.NumFiles + Header.StringTableSize);
-                file.Position = 0;
+                stream.Position = 0;
                 headerBytes = reader.ReadBytes(HeaderSize);
             }
 
@@ -40,7 +39,7 @@ namespace libhac
                 }
             }
 
-            Stream = file;
+            Stream = stream;
         }
 
         public byte[] GetFile(int index)

libhac/SdFs.cs

@@ -5,12 +5,14 @@ using System.Linq;
 
 namespace libhac
 {
-    public class SdFs
+    public class SdFs : IDisposable
     {
         public Keyset Keyset { get; }
         public string RootDir { get; }
         public string ContentsDir { get; }
         public string[] Files { get; }
+        public List<Nca> Ncas { get; } = new List<Nca>();
+        private List<Nax0> Nax0s { get; } = new List<Nax0>();
 
         public SdFs(Keyset keyset, string sdPath)
         {
@@ -24,7 +26,7 @@ namespace libhac
             Files = Directory.GetFiles(ContentsDir, "00", SearchOption.AllDirectories).Select(Path.GetDirectoryName).ToArray();
         }
 
-        public IEnumerable<Nca> ReadAllNca()
+        public void OpenAllNcas()
         {
             foreach (var file in Files)
             {
@@ -33,7 +35,8 @@ namespace libhac
                 {
                     var sdPath = "/" + Util.GetRelativePath(file, ContentsDir).Replace('\\', '/');
                     var nax0 = Nax0.CreateFromPath(Keyset, file, sdPath);
-                    nca = new Nca(Keyset, nax0.Stream);
+                    Nax0s.Add(nax0);
+                    nca = new Nca(Keyset, nax0.Stream, false);
                     nca.Name = Path.GetFileName(file);
                 }
                 catch (Exception ex)
@@ -41,9 +44,28 @@ namespace libhac
                     Console.WriteLine($"{ex.Message} {file}");
                 }
 
-                if (nca != null) yield return nca;
+                if (nca != null) Ncas.Add(nca);
             }
+        }
 
+        private void DisposeNcas()
+        {
+            foreach (var nca in Ncas)
+            {
+                nca.Dispose();
+            }
+            Ncas.Clear();
+
+            foreach (var nax0 in Nax0s)
+            {
+                nax0.Dispose();
+            }
+            Nax0s.Clear();
+        }
+
+        public void Dispose()
+        {
+            DisposeNcas();
         }
     }
 }

libhac/XTSSharp/RandomAccessSectorStream.cs

@@ -37,7 +37,7 @@ namespace libhac.XTSSharp
         private readonly byte[] _buffer;
         private readonly int _bufferSize;
         private readonly SectorStream _s;
-        private readonly bool _isStreamOwned;
+        private readonly bool _keepOpen;
         private bool _bufferDirty;
         private bool _bufferLoaded;
         private int _bufferPos;
@@ -48,7 +48,7 @@ namespace libhac.XTSSharp
         /// </summary>
         /// <param name="s">Base stream</param>
         public RandomAccessSectorStream(SectorStream s)
-            : this(s, false)
+            : this(s, true)
         {
         }
 
@@ -56,11 +56,11 @@ namespace libhac.XTSSharp
         /// Creates a new stream
         /// </summary>
         /// <param name="s">Base stream</param>
-        /// <param name="isStreamOwned">Does this stream own the base stream? i.e. should it be automatically disposed?</param>
+        /// <param name="keepOpen">Should this stream leave the base stream open when disposed?</param>
-        public RandomAccessSectorStream(SectorStream s, bool isStreamOwned)
+        public RandomAccessSectorStream(SectorStream s, bool keepOpen)
         {
             _s = s;
-            _isStreamOwned = isStreamOwned;
+            _keepOpen = keepOpen;
             _buffer = new byte[s.SectorSize];
             _bufferSize = s.SectorSize;
         }
@@ -152,7 +152,7 @@ namespace libhac.XTSSharp
 
             base.Dispose(disposing);
 
-            if (_isStreamOwned)
+            if (!_keepOpen)
                 _s.Dispose();
         }
 

libhac/XTSSharp/SectorStream.cs

@@ -37,6 +37,7 @@ namespace libhac.XTSSharp
         private readonly Stream _baseStream;
         private readonly long _offset;
         private ulong _currentSector;
+        private bool _keepOpen;
 
         /// <summary>
         /// Creates a new stream
@@ -55,12 +56,28 @@ namespace libhac.XTSSharp
         /// <param name="sectorSize">The size of the sectors to read/write</param>
         /// <param name="offset">Offset to start counting sectors</param>
         public SectorStream(Stream baseStream, int sectorSize, long offset)
+            : this(baseStream, sectorSize, offset, false)
         {
             SectorSize = sectorSize;
             _baseStream = baseStream;
             _offset = offset;
         }
 
+        /// <summary>
+        /// Creates a new stream
+        /// </summary>
+        /// <param name="baseStream">The base stream to read/write from</param>
+        /// <param name="sectorSize">The size of the sectors to read/write</param>
+        /// <param name="offset">Offset to start counting sectors</param>
+        /// <param name="keepOpen">Should this stream leave the base stream open when disposed?</param>
+        public SectorStream(Stream baseStream, int sectorSize, long offset, bool keepOpen)
+        {
+            SectorSize = sectorSize;
+            _baseStream = baseStream;
+            _offset = offset;
+            _keepOpen = keepOpen;
+        }
+
         /// <summary>
         /// The size of the sectors
         /// </summary>
@@ -229,5 +246,15 @@ namespace libhac.XTSSharp
             _baseStream.Write(buffer, offset, count);
             _currentSector++;
         }
+
+        protected override void Dispose(bool disposing)
+        {
+            if (!_keepOpen)
+            {
+                _baseStream.Dispose();
+            }
+
+            base.Dispose(disposing);
+        }
     }
 }

libhac/XTSSharp/XtsStream.cs

@@ -50,7 +50,7 @@ namespace libhac.XTSSharp
         /// <param name="xts">Xts implementation to use</param>
         /// <param name="sectorSize">Sector size</param>
         public XtsStream(Stream baseStream, Xts xts, int sectorSize)
-            : base(new XtsSectorStream(baseStream, xts, sectorSize), true)
+            : base(new XtsSectorStream(baseStream, xts, sectorSize), false)
         {
         }
 
@@ -63,7 +63,7 @@ namespace libhac.XTSSharp
         /// <param name="sectorSize">Sector size</param>
         /// <param name="offset">Offset to start counting sectors</param>
         public XtsStream(Stream baseStream, Xts xts, int sectorSize, long offset)
-            : base(new XtsSectorStream(baseStream, xts, sectorSize, offset), true)
+            : base(new XtsSectorStream(baseStream, xts, sectorSize, offset), false)
         {
         }
     }