Add clone-based asset editing and SerializedFile/UnityFS writeback

UnityAsset.NET.BundleFiles/BundleFile.cs

@@ -96,9 +96,10 @@ public static BlocksAndDirectoryInfo ParseDataInfo(IReader reader, Header header
         }
 
         public BundleFile(Header header, BlocksAndDirectoryInfo dataInfo, List<FileWrapper> files, string? key = null,
-            BlockCacheContext? blockCache = null)
+            BlockCacheContext? blockCache = null, UnityCN? unityCnInfo = null)
         {
             UnityCnKey = key;
+            UnityCnInfo = unityCnInfo;
             Header = header;
             DataInfo = dataInfo;
             Files = files;
@@ -170,6 +171,33 @@ public void Serialize(IWriter writer, CompressionType infoPacker = CompressionTy
 
             blockWriter.Finish();
             using var blockStream = blockWriter.GetDataStream();
+            if (unityCnKey != null)
+            {
+                if (UnityCnInfo == null || UnityCnInfo.Key != unityCnKey)
+                    throw new InvalidOperationException("UnityCN writing requires matching encryption metadata and key.");
+                for (var index = 0; index < blockWriter.BlockInfos.Count; index++)
+                {
+                    var block = blockWriter.BlockInfos[index];
+                    if ((block.Flags & StorageBlockFlags.CompressionTypeMask) == (StorageBlockFlags)CompressionType.None)
+                    {
+                        blockStream.Position += block.CompressedSize;
+                        continue;
+                    }
+                    var position = blockStream.Position;
+                    var bytes = new byte[checked((int)block.CompressedSize)];
+                    var read = 0;
+                    while (read < bytes.Length)
+                    {
+                        var count = blockStream.Read(bytes, read, bytes.Length - read);
+                        if (count == 0) throw new EndOfStreamException();
+                        read += count;
+                    }
+                    UnityCnInfo.EncryptBlock(bytes, bytes.Length, index);
+                    blockStream.Position = position;
+                    blockStream.Write(bytes, 0, bytes.Length);
+                }
+                blockStream.Position = 0;
+            }
 
             var dataInfo = new BlocksAndDirectoryInfo(DataInfo.UncompressedDataHash, blockWriter.BlockInfos.ToArray(),
                 directoryInfo.ToArray());
@@ -221,7 +249,6 @@ public void Serialize(IWriter writer, CompressionType infoPacker = CompressionTy
             if (infoAtEnd)
             {
                 header.Size += compressedDataInfoStream.Length;
-                if (needAlignAfterInfo) header.Size = Align(header.Size, 16);
             }
 
             header.Serialize(writer);
@@ -239,7 +266,6 @@ public void Serialize(IWriter writer, CompressionType infoPacker = CompressionTy
             if (infoAtEnd)
             {
                 writer.WriteStream(compressedDataInfoStream);
-                if (needAlignAfterInfo) writer.Align(16);
             }
         }
 
@@ -276,4 +302,4 @@ public static bool UnityCNVersionJudge(UnityRevision version)
                    (version >= "2022.3" && version <= "2022.3.1"); //2022.3.1 and earlier
         }
     }
-}
+}

UnityAsset.NET.BundleFiles/IO/Writer/BlockStreamWriter.cs

@@ -259,15 +259,15 @@ private void FlushBlock()
         {
             var compressed = Compression.CompressToStream(_buffer.AsSpan(0, _bufferPos), CompressionType);
 
-            if (compressed.Length < ChunkSize)
+            if (compressed.Length < _bufferPos)
             {
                 compressed.CopyTo(Stream);
                 BlockInfos.Add(new StorageBlockInfo((uint)_bufferPos, (uint)compressed.Length, (StorageBlockFlags)CompressionType));
             }
             else
             {
-                Stream.Write(_buffer, 0, ChunkSize);
-                BlockInfos.Add(new StorageBlockInfo((uint)ChunkSize, (uint)ChunkSize, (StorageBlockFlags)CompressionType.None));
+                Stream.Write(_buffer, 0, _bufferPos);
+                BlockInfos.Add(new StorageBlockInfo((uint)_bufferPos, (uint)_bufferPos, (StorageBlockFlags)CompressionType.None));
             }
 
             CheckMemoryLimit();
@@ -434,4 +434,4 @@ public override void Dispose()
             _bufferStream.Dispose();
         }
     }
-}
+}

UnityAsset.NET.TypeTreeHelper/Compiler/RoslynTypeBuilder.cs

@@ -39,16 +39,28 @@ public ClassDeclarationSyntax BuildClass(ClassTypeInfo typeInfo)
 
         classDeclaration = classDeclaration.AddMembers(classNameProperty);
 
-        var members = new List<MemberDeclarationSyntax>();
+        var members = new List<MemberDeclarationSyntax>
+        {
+            SyntaxFactory.ParseMemberDeclaration("private bool __assetEditable = false;")!,
+            SyntaxFactory.ParseMemberDeclaration($"public {typeInfo.GeneratedClassName} Clone() => UnityAsset.NET.Types.AssetCloner.Clone(this);")!
+        };
 
         foreach (var fieldInfo in typeInfo.Fields)
         {
+            members.Add(SyntaxFactory.ParseMemberDeclaration(
+                $"private {fieldInfo.DeclaredTypeSyntax} __asset_{fieldInfo.Name.TrimStart('@')} = default!;")!);
             var propertyDeclaration = SyntaxFactory.PropertyDeclaration(
                     fieldInfo.DeclaredTypeSyntax,
                     fieldInfo.Name)
                 .AddModifiers(SyntaxFactory.Token(SyntaxKind.PublicKeyword))
                 .AddAccessorListAccessors(
-                    SyntaxFactory.AccessorDeclaration(SyntaxKind.GetAccessorDeclaration).WithSemicolonToken(SyntaxFactory.Token(SyntaxKind.SemicolonToken))
+                    SyntaxFactory.AccessorDeclaration(SyntaxKind.GetAccessorDeclaration)
+                        .WithBody(SyntaxFactory.Block(SyntaxFactory.ParseStatement(
+                            $"return __assetEditable ? __asset_{fieldInfo.Name.TrimStart('@')} : UnityAsset.NET.Types.AssetCloner.CopyField(__asset_{fieldInfo.Name.TrimStart('@')});"))),
+                    SyntaxFactory.AccessorDeclaration(SyntaxKind.SetAccessorDeclaration)
+                        .WithBody(SyntaxFactory.Block(
+                            SyntaxFactory.ParseStatement("if (!__assetEditable) throw new InvalidOperationException(\"Clone the asset before editing it.\");"),
+                            SyntaxFactory.ParseStatement($"__asset_{fieldInfo.Name.TrimStart('@')} = value;")))
                 );
             members.Add(propertyDeclaration);
         }
@@ -863,7 +875,7 @@ private ConstructorDeclarationSyntax BuildConstructor(string className, List<Uni
             assignments.Add(SyntaxFactory.ExpressionStatement(
                 SyntaxFactory.AssignmentExpression(
                     SyntaxKind.SimpleAssignmentExpression,
-                    SyntaxFactory.IdentifierName(fieldInfo.Name),
+                    SyntaxFactory.IdentifierName($"__asset_{fieldInfo.Name.TrimStart('@')}"),
                         CreateReaderExpression(fieldInfo.TypeInfo, expectedType: fieldInfo.DeclaredTypeSyntax)
                 )
             ));
@@ -907,4 +919,4 @@ private ConstructorDeclarationSyntax BuildConstructor(string className, List<Uni
 
         return constructor;
     }
-}
+}

UnityAsset.NET.UnityTypeGen/InterfaceGenerator.cs

@@ -504,10 +504,16 @@ private CompilationUnitSyntax GenerateClassInterface(UnityTypeSchema schema, boo
             else
             {
                 var abstractGetter = SyntaxFactory.AccessorDeclaration(SyntaxKind.GetAccessorDeclaration)
-                    .WithSemicolonToken(SyntaxFactory.Token(SyntaxKind.SemicolonToken));
+                    .WithBody(SyntaxFactory.Block(SyntaxFactory.ParseStatement(
+                        "throw new System.NotSupportedException(\"This member has no implementation.\");")));
                 propertyDeclaration = propertyDeclaration.AddAccessorListAccessors(abstractGetter);
             }
 
+            // Versions lacking the member keep a rejecting setter; generated concrete members implement it normally.
+            propertyDeclaration = propertyDeclaration.AddAccessorListAccessors(
+                SyntaxFactory.AccessorDeclaration(SyntaxKind.SetAccessorDeclaration)
+                    .WithBody(SyntaxFactory.Block(SyntaxFactory.ParseStatement(
+                        "throw new System.NotSupportedException(\"This member is not editable in this version.\");"))));
             members.Add(propertyDeclaration);
         }
 
@@ -532,4 +538,4 @@ private CompilationUnitSyntax GenerateClassInterface(UnityTypeSchema schema, boo
             .AddMembers(namespaceDeclaration.AddMembers(interfaceDeclaration))
             .WithLeadingTrivia(leadingTrivia);
     }
-}
+}

UnityAsset.NET/Asset.cs

@@ -15,6 +15,7 @@ public class Asset : IEquatable<Asset>
     public bool IsNamedAsset;
 
     internal void AttachContext(AssetContext context) => Volatile.Write(ref _context, context);
+    internal AssetContext? Context => Volatile.Read(ref _context);
 
     internal void DetachContext() => Volatile.Write(ref _context, null);
 
@@ -40,6 +41,7 @@ private IUnityAsset GetValue()
             "(AssetManager), or that session has been cleared or disposed.");
 
         var value = context.Factory.Invoke(Info.Type, DataReader);
+        UnityAsset.NET.Types.AssetCloner.RecordOrigin(value, Info.Type);
 
         // Only after a successful materialisation: the session uses this to account for the blocks it just read.
         context.OnParsed?.Invoke(this);
@@ -148,4 +150,4 @@ private string ReadName(IUnityAsset value)
             : throw new InvalidOperationException(
                 $"Asset {Type}/{PathId} is named but its generated type " +
                 $"{value.GetType().Name} does not implement INamedObject.");
-}
+}

UnityAsset.NET/Editing/AssetValidation.cs

@@ -0,0 +1,53 @@
+using System.Collections;
+using System.Reflection;
+using UnityAsset.NET.Files.SerializedFiles;
+using UnityAsset.NET.Types.PreDefined;
+using UnityAsset.NET.Types.PreDefined.Interfaces;
+using UnityAsset.NET.TypeTreeHelper;
+
+namespace UnityAsset.NET.Editing;
+
+internal static class AssetValidation
+{
+    public static void References(IUnityAsset value, SerializedFile file, IReadOnlySet<long> pathIds)
+    {
+        var seen = new HashSet<object>(ReferenceEqualityComparer.Instance);
+        void Validate(int fileId, long pathId)
+        {
+            if (fileId < 0 || fileId > file.Metadata.Externals.Length)
+                throw new InvalidOperationException($"Reference uses undeclared external FileID {fileId}.");
+            if (fileId == 0 && pathId != 0 && !pathIds.Contains(pathId))
+                throw new InvalidOperationException($"Reference targets missing local PathID {pathId}.");
+        }
+        void Visit(object? item)
+        {
+            if (item is null || item is string or TypeTreeRepr || item.GetType().IsPrimitive || item.GetType().IsEnum) return;
+            if (!seen.Add(item)) return;
+            if (item is IPPtr pointer)
+            {
+                Validate(pointer.m_FileID, pointer.m_PathID);
+                return;
+            }
+            if (item is NodeData node && node.Type.StartsWith("PPtr<", StringComparison.Ordinal))
+            {
+                var fields = node.As<Dictionary<string, NodeData>>();
+                Validate(fields["m_FileID"].As<int>(), fields["m_PathID"].As<long>());
+                return;
+            }
+            if (item is IDictionary dictionary)
+            {
+                foreach (DictionaryEntry pair in dictionary) { Visit(pair.Key); Visit(pair.Value); }
+                return;
+            }
+            if (item is IEnumerable sequence)
+            {
+                foreach (var element in sequence) Visit(element);
+                return;
+            }
+            for (var type = item.GetType(); type != null; type = type.BaseType)
+                foreach (var field in type.GetFields(BindingFlags.Instance | BindingFlags.Public | BindingFlags.NonPublic | BindingFlags.DeclaredOnly))
+                    Visit(field.GetValue(item));
+        }
+        Visit(value);
+    }
+}

UnityAsset.NET/Editing/AtomicOutput.cs

@@ -0,0 +1,50 @@
+namespace UnityAsset.NET.Editing;
+
+internal static class AtomicOutput
+{
+    public static void Write(Stream destination, byte[] bytes)
+    {
+        ArgumentNullException.ThrowIfNull(destination);
+        if (!destination.CanWrite || !destination.CanSeek)
+            throw new NotSupportedException("Transactional stream output requires a writable, seekable stream. Use the path overload for atomic file replacement.");
+        var position = destination.Position;
+        var length = destination.Length;
+        var remaining = Math.Max(0, length - position);
+        if (remaining > 0 && !destination.CanRead)
+            throw new NotSupportedException("Overwriting stream content requires read access for rollback.");
+        var previous = new byte[checked((int)remaining)];
+        try
+        {
+            if (previous.Length > 0) destination.ReadExactly(previous);
+        }
+        finally { destination.Position = position; }
+        try { destination.Write(bytes); }
+        catch
+        {
+            destination.Position = position;
+            if (previous.Length > 0) destination.Write(previous);
+            destination.SetLength(length);
+            destination.Position = position;
+            throw;
+        }
+    }
+
+    public static void Write(string path, byte[] bytes)
+    {
+        var fullPath = Path.GetFullPath(path);
+        var temporary = Path.Combine(Path.GetDirectoryName(fullPath)!, $".{Path.GetFileName(path)}.{Guid.NewGuid():N}.tmp");
+        try
+        {
+            using (var stream = new FileStream(temporary, FileMode.CreateNew, FileAccess.Write, FileShare.None))
+            {
+                stream.Write(bytes);
+                stream.Flush(flushToDisk: true);
+            }
+            File.Move(temporary, fullPath, overwrite: true);
+        }
+        finally
+        {
+            if (File.Exists(temporary)) File.Delete(temporary);
+        }
+    }
+}

UnityAsset.NET/Editing/BundleEditor.cs

@@ -0,0 +1,85 @@
+using UnityAsset.NET.BundleFiles;
+using UnityAsset.NET.Extensions;
+using UnityAsset.NET.Files.SerializedFiles;
+using UnityAsset.NET.FileSystem;
+using UnityAsset.NET.IO;
+using UnityAsset.NET.IO.Reader;
+using UnityAsset.NET.IO.Writer;
+
+namespace UnityAsset.NET.Editing;
+
+public sealed class BundleEditor
+{
+    private readonly BundleFile _original;
+    private readonly Dictionary<string, SerializedFileEditor> _files;
+    private readonly Lazy<(byte[] Bytes, BundleFile File)> _result;
+
+    public BundleEditor(BundleFile original) : this(original, new(StringComparer.Ordinal)) { }
+    private BundleEditor(BundleFile original, Dictionary<string, SerializedFileEditor> files)
+    {
+        _original = original ?? throw new ArgumentNullException(nameof(original));
+        _files = files;
+        _result = new(Build);
+    }
+
+    public BundleEditor WithFile(string path, SerializedFileEditor file)
+    {
+        ArgumentNullException.ThrowIfNull(file);
+        if (!_original.Files.Any(f => f.Info.Path == path)) throw new ArgumentException($"Unknown bundle path: {path}");
+        var files = new Dictionary<string, SerializedFileEditor>(_files, StringComparer.Ordinal) { [path] = file };
+        return new BundleEditor(_original, files);
+    }
+
+    public BundleFile Materialize() => _result.Value.File;
+    public void Write(Stream destination) => AtomicOutput.Write(destination, _result.Value.Bytes);
+    public void Write(string path) => AtomicOutput.Write(path, _result.Value.Bytes);
+
+    private (byte[], BundleFile) Build()
+    {
+        var source = _original.SourceVirtualFile ?? throw new InvalidOperationException("The original bundle must have source bytes.");
+        var reader = new CustomFileReaderProvider(source).CreateReader();
+        var originalBytes = reader.ReadBytes(checked((int)reader.Length));
+        var unchanged = true;
+        var files = new List<FileWrapper>();
+        foreach (var entry in _original.Files)
+        {
+            byte[] bytes;
+            var raw = entry.File is SerializedFile sf ? sf.ReaderProvider.CreateReader()
+                : entry.File is IReaderProvider provider ? provider.CreateReader() : null;
+            if (raw == null) throw new NotSupportedException($"No source bytes for {entry.Info.Path}.");
+            var original = raw.ReadBytes(checked((int)raw.Length));
+            if (_files.TryGetValue(entry.Info.Path, out var editor))
+            {
+                bytes = editor.Bytes;
+                unchanged &= original.AsSpan().SequenceEqual(bytes);
+            }
+            else bytes = original;
+            files.Add(new FileWrapper(new MemoryFileProvider(bytes), entry.Info));
+        }
+        byte[] result;
+        if (unchanged) result = originalBytes;
+        else
+        {
+            var bundle = new BundleFile(_original.Header!, _original.DataInfo!, files,
+                _original.UnityCnKey, unityCnInfo: _original.UnityCnInfo);
+            using var output = new MemoryStream();
+            using (var writer = new CustomStreamWriter(output, leaveOpen: true))
+            {
+                bundle.Serialize(writer, CompressionType.Lz4HC, CompressionType.Lz4HC, _original.UnityCnKey);
+                writer.Finish();
+            }
+            result = output.ToArray();
+        }
+        var materialized = new BundleFile(new MemoryFileInfo(source.Path, result, FileType.BundleFile), _original.UnityCnKey);
+        materialized.ParseFilesWithTypeConversion();
+        foreach (var entry in materialized.Files)
+        {
+            if (entry.File is not SerializedFile sf) continue;
+            var template = _files.TryGetValue(entry.Info.Path, out var editor) ? editor.Materialize()
+                : _original.Files.First(e => e.Info.Path == entry.Info.Path).File as SerializedFile;
+            foreach (var asset in sf.Assets)
+                if (template?.PathToAsset.GetValueOrDefault(asset.PathId)?.Context is { } context) asset.AttachContext(context);
+        }
+        return (result, materialized);
+    }
+}