AssetRipper_AssetRipper/Source/AssetRipper.Export.Modules.Shader/DXShaderProgramRestorer.cs

using AssetRipper.Export.Modules.Shaders.Resources;
using AssetRipper.Export.Modules.Shaders.ShaderBlob;
using AssetRipper.IO.Endian;
using AssetRipper.Primitives;

namespace AssetRipper.Export.Modules.Shaders;

public static class DXShaderProgramRestorer
{
	public static byte[] RestoreProgramData(BinaryReader reader, UnityVersion version, ref ShaderSubProgram shaderSubProgram)
	{
		using MemoryStream dest = new MemoryStream();
		using BinaryWriter writer = new BinaryWriter(dest);
		uint baseOffset = (uint)reader.BaseStream.Position;
		byte[] magicBytes = reader.ReadBytes(4);
		byte[] checksum = reader.ReadBytes(16);
		uint unknown0 = reader.ReadUInt32();
		uint totalSize = reader.ReadUInt32();
		uint chunkCount = reader.ReadUInt32();
		List<uint> chunkOffsets = new List<uint>();
		for (int i = 0; i < chunkCount; i++)
		{
			chunkOffsets.Add(reader.ReadUInt32());
		}
		uint bodyOffset = (uint)reader.BaseStream.Position;
		// Check if shader already has resource chunk
		foreach (uint chunkOffset in chunkOffsets)
		{
			reader.BaseStream.Position = chunkOffset + baseOffset;
			uint fourCc = reader.ReadUInt32();
			if (fourCc == RDEFFourCC)
			{
				reader.BaseStream.Position = baseOffset;
				byte[] original = reader.ReadBytes((int)reader.BaseStream.Length);
				return original;
			}
		}
		reader.BaseStream.Position = bodyOffset;
		byte[] resourceChunkData = GetResourceChunk(version, ref shaderSubProgram);
		//Adjust for new chunk
		totalSize += (uint)resourceChunkData.Length;
		for (int i = 0; i < chunkCount; i++)
		{
			chunkOffsets[i] += (uint)resourceChunkData.Length + 4;
		}
		chunkOffsets.Insert(0, bodyOffset - baseOffset + 4);
		chunkCount += 1;
		totalSize += (uint)resourceChunkData.Length;

		writer.Write(magicBytes);
		writer.Write(checksum);
		writer.Write(unknown0);
		writer.Write(totalSize);
		writer.Write(chunkCount);
		foreach (uint chunkOffset in chunkOffsets)
		{
			writer.Write(chunkOffset);
		}
		writer.Write(resourceChunkData);
		byte[] rest = reader.ReadBytes((int)reader.BaseStream.Length - (int)reader.BaseStream.Position);
		writer.Write(rest);
		return dest.ToArray();
	}

	private static byte[] GetResourceChunk(UnityVersion version, ref ShaderSubProgram shaderSubprogram)
	{
		using MemoryStream memoryStream = new MemoryStream();
		using EndianWriter writer = new EndianWriter(memoryStream, EndianType.LittleEndian);
		ResourceChunk resourceChunk = new ResourceChunk(version, ref shaderSubprogram);
		resourceChunk.Write(writer);
		//uint size = resourceChunk.Size;
		//if (memoryStream.Length != resourceChunk.Size) throw new Exception("Expected size does not match actual size");
		return memoryStream.ToArray();
	}

	/// <summary>
	/// 'RDEF' ascii
	/// </summary>
	public const uint RDEFFourCC = 0x46454452;
}