jianguotian commented on code in PR #79: URL: https://github.com/apache/paimon-mosaic/pull/79#discussion_r3844207487
########## tools/native_binary.py: ########## @@ -0,0 +1,1724 @@ +#!/usr/bin/env python3 + +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Validate native-library format, architecture, structure, and exports.""" + +from __future__ import annotations + +import struct +from dataclasses import dataclass + + +TARGET_ARCHITECTURE = { + "x86_64-unknown-linux-gnu": ("ELF", "x86_64"), + "aarch64-unknown-linux-gnu": ("ELF", "aarch64"), + "aarch64-apple-darwin": ("Mach-O", "aarch64"), + "x86_64-pc-windows-msvc": ("PE", "x86_64"), +} + +MACHINE_ARCHITECTURE = { + 62: "x86_64", + 183: "aarch64", +} +PE_MACHINE_ARCHITECTURE = { + 0x8664: "x86_64", + 0xAA64: "aarch64", +} +MACHO_CPU_ARCHITECTURE = { + 0x01000007: "x86_64", + 0x0100000C: "aarch64", +} + +# A Mosaic native library exports on the order of a hundred symbols; the bound +# exists to keep hostile input from driving quadratic hash-membership work. +MAX_DYNAMIC_SYMBOLS = 100_000 + +MOSAIC_SYMBOL_FAMILIES = { + "JNI": { + "Java_org_apache_paimon_mosaic_NativeLib_nativeReaderExportSchema", + "Java_org_apache_paimon_mosaic_NativeLib_nativeReaderFree", + "Java_org_apache_paimon_mosaic_NativeLib_nativeReaderNumRowGroups", + "Java_org_apache_paimon_mosaic_NativeLib_nativeReaderOpen", + "Java_org_apache_paimon_mosaic_NativeLib_nativeReaderOpenRowGroup", + "Java_org_apache_paimon_mosaic_NativeLib_nativeReaderRowGroupNumRows", + "Java_org_apache_paimon_mosaic_NativeLib_nativeReaderRowGroupStatMaxs", + "Java_org_apache_paimon_mosaic_NativeLib_nativeReaderRowGroupStatMins", + "Java_org_apache_paimon_mosaic_NativeLib_nativeReaderRowGroupStatNames", + "Java_org_apache_paimon_mosaic_NativeLib_nativeReaderRowGroupStatNullCounts", + "Java_org_apache_paimon_mosaic_NativeLib_nativeReaderSetProjection", + "Java_org_apache_paimon_mosaic_NativeLib_nativeRowGroupReaderFree", + "Java_org_apache_paimon_mosaic_NativeLib_nativeRowGroupReaderNumRows", + "Java_org_apache_paimon_mosaic_NativeLib_nativeRowGroupReaderReadColumns", + "Java_org_apache_paimon_mosaic_NativeLib_nativeWriterClose", + "Java_org_apache_paimon_mosaic_NativeLib_nativeWriterEstimatedSize", + "Java_org_apache_paimon_mosaic_NativeLib_nativeWriterFree", + "Java_org_apache_paimon_mosaic_NativeLib_nativeWriterNumRowGroups", + "Java_org_apache_paimon_mosaic_NativeLib_nativeWriterOpen", + "Java_org_apache_paimon_mosaic_NativeLib_nativeWriterRowGroupStatMaxs", + "Java_org_apache_paimon_mosaic_NativeLib_nativeWriterRowGroupStatMins", + "Java_org_apache_paimon_mosaic_NativeLib_nativeWriterRowGroupStatNames", + "Java_org_apache_paimon_mosaic_NativeLib_nativeWriterRowGroupStatNullCounts", + "Java_org_apache_paimon_mosaic_NativeLib_nativeWriterWriteBatch", + }, + "FFI": { + "mosaic_last_error", + "mosaic_reader_export_schema", + "mosaic_reader_free", + "mosaic_reader_num_row_groups", + "mosaic_reader_open", + "mosaic_reader_open_row_group", + "mosaic_reader_row_group_num_rows", + "mosaic_reader_row_group_num_stats", + "mosaic_reader_row_group_stats", + "mosaic_reader_set_projection", + "mosaic_record_batch_export", + "mosaic_record_batch_free", + "mosaic_record_batch_num_columns", + "mosaic_record_batch_num_rows", + "mosaic_row_group_reader_free", + "mosaic_row_group_reader_num_rows", + "mosaic_row_group_reader_read_columns", + "mosaic_writer_close", + "mosaic_writer_estimated_file_size", + "mosaic_writer_free", + "mosaic_writer_num_row_groups", + "mosaic_writer_open", + "mosaic_writer_options_default", + "mosaic_writer_row_group_num_stats", + "mosaic_writer_row_group_stats", + "mosaic_writer_write_batch", + }, +} + + +@dataclass(frozen=True) +class NativeBinary: + binary_format: str + architectures: frozenset[str] + exported_symbols: frozenset[str] + + +@dataclass(frozen=True) +class ElfSection: + section_type: int + flags: int + address: int + offset: int + size: int + link: int + entry_size: int + + +def require_range(data: bytes, offset: int, size: int, description: str) -> None: + if ( + offset < 0 + or size < 0 + or offset > len(data) + or size > len(data) - offset + ): + raise ValueError(f"{description} is out of bounds") + + +def is_power_of_two(value: int) -> bool: + return value > 0 and value & (value - 1) == 0 + + +def c_string_bytes( + data: bytes, offset: int, limit: int, description: str +) -> bytes: + if offset < 0 or offset >= limit or limit > len(data): + raise ValueError(f"{description} is out of bounds") + terminator = data.find(b"\0", offset, limit) + if terminator < 0: + raise ValueError(f"{description} is not null-terminated") + return data[offset:terminator] + + +def ascii_symbol(raw_name: bytes) -> str | None: + try: + return raw_name.decode("ascii") + except UnicodeDecodeError: + return None + + +def elf_virtual_range( + data: bytes, + address: int, + size: int, + load_segments: list[tuple[int, int, int, int]], + description: str, +) -> int: + mapped_offsets = set() + memory_mapped = False + for file_offset, virtual_address, file_size, memory_size in load_segments: + if address < virtual_address: + continue + delta = address - virtual_address + if delta > memory_size or size > memory_size - delta: + continue + memory_mapped = True + if delta <= file_size and size <= file_size - delta: + mapped_offsets.add(file_offset + delta) + + if not mapped_offsets: + if memory_mapped: + raise ValueError(f"{description} is not file-backed") + raise ValueError(f"{description} is not mapped by an ELF PT_LOAD segment") + if len(mapped_offsets) != 1: + raise ValueError(f"{description} has ambiguous ELF PT_LOAD mappings") + + offset = mapped_offsets.pop() + require_range(data, offset, size, description) + return offset + + +def elf_loader_section( + data: bytes, + sections: list[ElfSection], + load_segments: list[tuple[int, int, int, int]], + address: int, + section_type: int, + minimum_size: int, + linked_section: int | None, + description: str, + section_name: str, +) -> tuple[int, ElfSection]: + offset = elf_virtual_range( + data, address, minimum_size, load_segments, description + ) + matches = [ + (index, section) + for index, section in enumerate(sections) + if section.section_type == section_type + and section.address == address + and section.offset == offset + ] + if not matches: + raise ValueError( + f"{description} does not reference an {section_name} section" + ) + if len(matches) != 1: + raise ValueError(f"{description} references multiple {section_name} sections") + + index, section = matches[0] + if not section.flags & 0x2: + raise ValueError(f"{description} section is not allocated") + if section.size < minimum_size: + raise ValueError(f"{description} is truncated") + if linked_section is not None and section.link != linked_section: + raise ValueError(f"{description} section does not link to DT_SYMTAB") + if ( + elf_virtual_range( + data, address, section.size, load_segments, description + ) + != section.offset + ): + raise ValueError(f"{description} has inconsistent file mapping") + return index, section + + +def elf_sysv_hash(name: bytes) -> int: + value = 0 + for byte in name: + value = (value << 4) + byte + high = value & 0xF0000000 + if high: + value ^= high >> 24 + value &= ~high + return value & 0xFFFFFFFF + + +def elf_gnu_hash(name: bytes) -> int: + value = 5381 + for byte in name: + value = (value * 33 + byte) & 0xFFFFFFFF + return value + + +@dataclass(frozen=True) +class ElfSysvHash: + buckets: tuple[int, ...] + chains: tuple[int, ...] + + @property + def symbol_count(self) -> int: + return len(self.chains) + + def contains(self, symbol_index: int, name: bytes) -> bool: + index = self.buckets[elf_sysv_hash(name) % len(self.buckets)] + while index: + if index == symbol_index: Review Comment: Fixed. ELF hash ownership is built once for O(1) lookup, and PE/Mach-O symbol counts are capped. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
