diff --git a/UnityPy/export/ShaderConverter.py b/UnityPy/export/ShaderConverter.py index dd3ac1c5..9b7f017a 100644 --- a/UnityPy/export/ShaderConverter.py +++ b/UnityPy/export/ShaderConverter.py @@ -61,28 +61,32 @@ def ConvertSerializedShader(m_Shader: Shader) -> str: platformNumber = len(m_Shader.platforms) compressed_blob = bytes(m_Shader.compressedBlob) - def get_entry(array: Union[List[T], List[List[T]]], index: int) -> T: + def get_entries(array: Union[List[T], List[List[T]]], index: int) -> List[T]: item = array[index] - if isinstance(item, List): - return item[0] - return item + return item if isinstance(item, list) else [item] for i in range(platformNumber): if i >= len(m_Shader.compressedLengths) or i >= len(m_Shader.decompressedLengths): # m_Shader.platforms shouldn't be longer than m_shader.[de]compressedLengths, but it is break - compressedSize = get_entry(m_Shader.compressedLengths, i) - decompressedSize = get_entry(m_Shader.decompressedLengths, i) - offset = get_entry(m_Shader.offsets, i) + compressedSizes = get_entries(m_Shader.compressedLengths, i) + decompressedSizes = get_entries(m_Shader.decompressedLengths, i) + offsets = get_entries(m_Shader.offsets, i) + if not offsets or not (len(offsets) == len(compressedSizes) == len(decompressedSizes)): + raise ValueError("Shader platform has inconsistent segment arrays") - compressedBytes = compressed_blob[offset : offset + compressedSize] - decompressedBytes = CompressionHelper.decompress_lz4(compressedBytes, decompressedSize) + segmentReaders = [] + for offset, compressedSize, decompressedSize in zip(offsets, compressedSizes, decompressedSizes): + compressedBytes = compressed_blob[offset : offset + compressedSize] + decompressedBytes = CompressionHelper.decompress_lz4(compressedBytes, decompressedSize) + segmentReaders.append(EndianBinaryReader(decompressedBytes, endian="<")) shaderPrograms.append( ShaderProgram( - EndianBinaryReader(decompressedBytes, endian="<"), + segmentReaders[0], m_Shader.object_reader.version, + segmentReaders, ) ) @@ -543,7 +547,14 @@ def GetPlatformString(platform: int): class ShaderProgram: m_SubPrograms: List[ShaderSubProgram] - def __init__(self, reader: EndianBinaryReader, version: Tuple[int, int, int, int]): + def __init__( + self, + reader: EndianBinaryReader, + version: Tuple[int, int, int, int], + segment_readers: Optional[List[EndianBinaryReader]] = None, + ): + if segment_readers is None: + segment_readers = [reader] subProgramCapacity = reader.read_int() self.m_SubPrograms = [None] * subProgramCapacity @@ -555,8 +566,14 @@ def __init__(self, reader: EndianBinaryReader, version: Tuple[int, int, int, int for i in range(subProgramCapacity): reader.Position = 4 + i * entrySize offset = reader.read_int() - reader.Position = offset - self.m_SubPrograms[i] = ShaderSubProgram(reader) + reader.read_int() # length + segment = reader.read_int() if entrySize == 12 else 0 + if segment < 0 or segment >= len(segment_readers): + raise ValueError(f"Shader subprogram references missing segment {segment}") + # Entry offsets are relative to the selected decompressed segment. + segment_reader = segment_readers[segment] + segment_reader.Position = offset + self.m_SubPrograms[i] = ShaderSubProgram(segment_reader) def Export(self, shader: str) -> str: shader = re.sub( diff --git a/tests/test_shader_segments.py b/tests/test_shader_segments.py new file mode 100644 index 00000000..a537ad62 --- /dev/null +++ b/tests/test_shader_segments.py @@ -0,0 +1,234 @@ +"""Synthetic shader records exercise decompression, parsing and text export.""" + +import io +import struct +from types import SimpleNamespace + +import lz4.block +import pytest + +from UnityPy.enums import PassType, ShaderCompilerPlatform, ShaderGpuProgramType +from UnityPy.export.ShaderConverter import ShaderProgram, export_shader +from UnityPy.streams import EndianBinaryReader + + +def aligned_string(value): + data = value.encode("utf8") + return struct.pack("= 201608170: + data += struct.pack("= (2019, 3): + data += struct.pack("= (2019, 3) else 8) + entries = [(header_size, len(record), 0)] * copies + return entry_table(entries, version) + record + + +def shader(segments, version, program_count, nested=True): + offsets, compressed_sizes, decompressed_sizes = [], [], [] + payload = b"prefix" + for segment in segments: + compressed = lz4.block.compress(segment, store_size=False) + offsets.append(len(payload)) + compressed_sizes.append(len(compressed)) + decompressed_sizes.append(len(segment)) + payload += compressed + b"gap" + + tags = SimpleNamespace(tags=[]) + state = SimpleNamespace( + m_Name="", + m_LOD=0, + m_Tags=tags, + alphaToMask=SimpleNamespace(val=0), + zClip=SimpleNamespace(val=1), + zTest=SimpleNamespace(val=4), + zWrite=SimpleNamespace(val=1), + culling=SimpleNamespace(val=2), + offsetFactor=SimpleNamespace(val=0), + offsetUnits=SimpleNamespace(val=0), + lighting=False, + gpuProgramID=0, + ) + subprograms = [ + SimpleNamespace( + m_BlobIndex=index, + m_GpuProgramType=ShaderGpuProgramType.kShaderGpuProgramGLES3, + m_ShaderHardwareTier=0, + ) + for index in range(program_count) + ] + empty_program = SimpleNamespace(m_SubPrograms=[]) + shader_pass = SimpleNamespace( + m_Type=PassType.kPassTypeNormal, + m_State=state, + progVertex=SimpleNamespace(m_SubPrograms=subprograms), + progFragment=empty_program, + progGeometry=empty_program, + progHull=empty_program, + progDomain=empty_program, + ) + parsed_form = SimpleNamespace( + m_Name="SyntheticShader", + m_PropInfo=SimpleNamespace(m_Props=[]), + m_SubShaders=[SimpleNamespace(m_LOD=0, m_Tags=tags, m_Passes=[shader_pass])], + m_FallbackName="", + m_CustomEditorName="", + ) + return SimpleNamespace( + m_SubProgramBlob=None, + compressedBlob=payload, + platforms=[ShaderCompilerPlatform.kShaderCompPlatformGLES3Plus], + offsets=[offsets] if nested else offsets, + compressedLengths=[compressed_sizes] if nested else compressed_sizes, + decompressedLengths=[decompressed_sizes] if nested else decompressed_sizes, + object_reader=SimpleNamespace(version=version), + m_ParsedForm=parsed_form, + ) + + +@pytest.mark.parametrize( + "version,record_version", + [ + ((5, 3, 0, 0), 201509030), + ((2018, 4, 0, 0), 201802150), + ((2019, 2, 0, 0), 201806140), + ((2019, 3, 0, 0), 201806140), + ((2022, 3, 0, 0), 202012090), + ], +) +def test_single_reader_constructor_preserves_entry_layouts(version, record_version): + code = b"void single_segment() {}" + data = single_segment(program_record(code, record_version), version) + program = ShaderProgram(EndianBinaryReader(data, endian="<"), version) + subprogram = program.m_SubPrograms[0] + assert subprogram.m_ProgramCode == code + assert subprogram.m_Keywords == ["SYNTHETIC_GLOBAL"] + assert subprogram.m_LocalKeywords == (["SYNTHETIC_LOCAL"] if record_version == 201806140 else None) + assert code.decode() in program.Export("GpuProgramIndex 0") + + +@pytest.mark.parametrize( + "version,record_version,nested", + [ + ((2018, 4, 0, 0), 201802150, False), + ((2019, 3, 0, 0), 201806140, True), + ((2022, 3, 0, 0), 202012090, True), + ], +) +def test_single_segment_shader_export(version, record_version, nested): + code = b"void single_export() {}" + segment = single_segment(program_record(code, record_version), version) + text = export_shader(shader([segment], version, 1, nested=nested)) + assert 'Shader "SyntheticShader"' in text + assert text.count(code.decode()) == 1 + + +@pytest.mark.parametrize("version,record_version", [((2019, 3, 0, 0), 201806140), ((2022, 3, 0, 0), 202012090)]) +def test_shader_export_reads_segment_local_offsets_in_index_order(version, record_version): + codes = [ + b"void third_segment() {}", + b"void table_segment() {}", + b"void second_segment() {}", + b"void second_segment_later() {}", + ] + records = [program_record(code, record_version) for code in codes] + header_size = 4 + 4 * 12 + table = entry_table( + [ + (0, len(records[0]), 2), + (header_size, len(records[1]), 0), + (0, len(records[2]), 1), + (len(records[2]), len(records[3]), 1), + ], + version, + ) + segments = [table + records[1], records[2] + records[3], records[0]] + text = export_shader(shader(segments, version, 4)) + assert all(text.count(code.decode()) == 1 for code in codes) + positions = [text.index(code.decode()) for code in codes] + assert positions == sorted(positions) + + +def test_shader_export_keeps_platform_segment_lists_separate(): + version = (2022, 3, 0, 0) + gles_code, glcore_code = b"void gles_program() {}", b"void glcore_program() {}" + gles_record = program_record(gles_code, 202012090) + glcore_record = program_record(glcore_code, 202012090, ShaderGpuProgramType.kShaderGpuProgramGLCore32) + first = shader([entry_table([(0, len(gles_record), 1)], version), gles_record], version, 1) + second = shader([entry_table([(0, len(glcore_record), 1)], version), glcore_record], version, 1) + first.offsets.append([offset + len(first.compressedBlob) for offset in second.offsets[0]]) + first.compressedBlob += second.compressedBlob + first.compressedLengths.extend(second.compressedLengths) + first.decompressedLengths.extend(second.decompressedLengths) + first.platforms.append(ShaderCompilerPlatform.kShaderCompPlatformOpenGLCore) + first.m_ParsedForm.m_SubShaders[0].m_Passes[0].progVertex.m_SubPrograms.append( + SimpleNamespace( + m_BlobIndex=0, + m_GpuProgramType=ShaderGpuProgramType.kShaderGpuProgramGLCore32, + m_ShaderHardwareTier=0, + ) + ) + text = export_shader(first) + assert text.count(gles_code.decode()) == 1 + assert text.count(glcore_code.decode()) == 1 + + +def test_entries_may_share_a_program_record(): + version = (2019, 3, 0, 0) + code = b"void shared_program() {}" + segment = single_segment(program_record(code, 201806140), version, copies=2) + program = ShaderProgram(EndianBinaryReader(segment, endian="<"), version) + assert [subprogram.m_ProgramCode for subprogram in program.m_SubPrograms] == [code, code] + + +@pytest.mark.parametrize("stream_backed", [False, True]) +def test_shared_record_in_nonzero_segment_with_reader_backends(stream_backed): + version = (2022, 3, 0, 0) + code = b"void shared_nonzero_record() {}" + record = program_record(code, 202012090) + table = entry_table([(4, len(record), 2), (4, len(record), 2)], version) + segments = [table, b"unused", b"pad!" + record] + readers = [EndianBinaryReader(io.BytesIO(part) if stream_backed else part, endian="<") for part in segments] + program = ShaderProgram(readers[0], version, readers) + + assert [item.m_ProgramCode for item in program.m_SubPrograms] == [code, code] + assert [item.m_Keywords for item in program.m_SubPrograms] == [["SYNTHETIC_GLOBAL"], ["SYNTHETIC_GLOBAL"]] + text = program.Export("GpuProgramIndex 0\nGpuProgramIndex 1") + assert text.count(code.decode()) == 2 + + +@pytest.mark.parametrize("segment", [-1, 1]) +def test_unavailable_segment_is_reported(segment): + version = (2019, 3, 0, 0) + table = entry_table([(0, 4, segment)], version) + with pytest.raises(ValueError, match=f"missing segment {segment}"): + ShaderProgram(EndianBinaryReader(table, endian="<"), version) + + +@pytest.mark.parametrize("field", ["offsets", "compressedLengths", "decompressedLengths"]) +def test_mismatched_segment_arrays_are_not_silently_truncated(field): + version = (2019, 3, 0, 0) + record = program_record(b"void incomplete_metadata() {}", 201806140) + source = shader([entry_table([(0, len(record), 1)], version), record], version, 1) + getattr(source, field)[0].pop() + with pytest.raises(ValueError, match="inconsistent segment arrays"): + export_shader(source)