# Copyright (c) 2025-2026 Buf Technologies, Inc. # # Licensed 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. from __future__ import annotations import struct from dataclasses import dataclass from typing import final from ._wire_type import WireType DEPTH_LIMIT = 100 @dataclass(frozen=True, slots=True, init=False) class Tag: """Tag of a protobuf field, consisting of a field number and a wire type.""" number: int wire_type: WireType raw: int def __init__(self, key: int) -> None: number = key >> 3 wire_type = WireType(key & 0x07) if number == 0: msg = "invalid tag with field number 0" raise ValueError(msg) object.__setattr__(self, "number", number) object.__setattr__(self, "wire_type", wire_type) object.__setattr__(self, "raw", key) @final class BinaryReader: """A reader for deserializing Protocol Buffer wire format.""" def __init__(self, data: memoryview) -> None: """Initialize a BinaryReader. Args: data: The underlying memoryview to read from. """ self._data = data self._offset = 0 @property def offset(self) -> int: """The current read position within the buffer.""" return self._offset def read(self, size: int) -> memoryview: """Read a given number of bytes from the buffer.""" if self._offset + size > len(self._data): msg = "unexpected end of buffer" raise EOFError(msg) result = self._data[self._offset : self._offset + size] self._offset += size return result def seek(self, offset: int) -> None: """Move the buffer position to the given offset.""" if offset < 0 or offset > len(self._data): msg = "invalid seek offset" raise ValueError(msg) self._offset = offset def varint(self, limit: int = 10, last_byte_limit: int = 1) -> int: """Read a varint from the buffer.""" result = 0 shift = 0 data = self._data offset = self._offset end = len(data) # protobuf guarantees varints are at most 10 bytes for i in range(limit): if offset >= end: msg = "unexpected end of buffer while reading varint" raise EOFError(msg) byte = data[offset] offset += 1 result |= (byte & 0x7F) << shift shift += 7 if byte < 0x80: if i == limit - 1 and byte > last_byte_limit: msg = "invalid varint" raise ValueError(msg) self._offset = offset return result msg = "invalid varint" raise ValueError(msg) def tag(self) -> Tag: """Read a varint and interpret it as a protobuf field tag, consisting of a field number and a wire type.""" # Protobuf tags are limited to 5-byte varints key = self.varint(5, 0x0F) return Tag(key) def bool_(self) -> bool: """Read a varint and interpret it as a boolean.""" return bool(self.varint()) def int32(self) -> int: """Read a varint and interpret it as a signed 32-bit integer.""" value = self.varint() value %= 1 << 32 return value if value < (1 << 31) else value - (1 << 32) def int64(self) -> int: """Read a varint and interpret it as a signed 64-bit integer.""" value = self.varint() return value if value < (1 << 63) else value - (1 << 64) def uint32(self) -> int: """Read a varint and interpret it as an unsigned 32-bit integer.""" value = self.varint() return value % (1 << 32) def uint64(self) -> int: """Read a varint and interpret it as an unsigned 64-bit integer.""" return self.varint() def sint32(self) -> int: """Read a varint and interpret it as a signed 32-bit integer with zigzag encoding.""" value = self.varint() value = value % (1 << 32) return ((value + 1) >> 1) * (-1 if value & 1 else 1) def sint64(self) -> int: """Read a varint and interpret it as a signed 64-bit integer with zigzag encoding.""" value = self.varint() return ((value + 1) >> 1) * (-1 if value & 1 else 1) def float_(self) -> float: """Read 4 bytes and interpret them as a 32-bit floating point number.""" return struct.unpack(" float: """Read 8 bytes and interpret them as a 64-bit floating point number.""" return struct.unpack(" int: """Read 4 bytes and interpret them as an unsigned 32-bit integer.""" return struct.unpack(" int: """Read 4 bytes and interpret them as a signed 32-bit integer.""" return struct.unpack(" int: """Read 8 bytes and interpret them as an unsigned 64-bit integer.""" return struct.unpack(" int: """Read 8 bytes and interpret them as a signed 64-bit integer.""" return struct.unpack(" memoryview: """Skip a field value and return a memoryview on the skipped bytes. Args: wire_type: The wire type of the field to skip. depth: The current recursion depth of the message tree. field_number: The field number. Used for delimited encoding. Returns: A [`memoryview`][] over the skipped bytes. """ if depth > DEPTH_LIMIT: msg = f"exceeded maximum recursion depth {DEPTH_LIMIT} while skipping field {field_number}" raise RecursionError(msg) start = self._offset match wire_type: case WireType.VARINT: self.varint() case WireType.BIT64: self.read(8) case WireType.LENGTH_DELIMITED: length = self.varint() self.read(length) case WireType.BIT32: self.read(4) case WireType.SGROUP: while True: inner_tag = self.tag() if inner_tag.wire_type == WireType.EGROUP: if inner_tag.number != field_number: msg = f"mismatched group tag: expected {field_number}, got {inner_tag.number}" raise ValueError(msg) break nested_depth = ( depth + 1 if inner_tag.wire_type == WireType.SGROUP else depth ) self.skip( inner_tag.wire_type, depth=nested_depth, field_number=inner_tag.number, ) case WireType.EGROUP: msg = "unexpected end group tag outside of group" raise ValueError(msg) return self._data[start : self._offset] def bytes_(self) -> bytes: """Read a length-delimited byte sequence.""" length = self.varint() return bytes(self.read(length)) def string(self) -> str: """Read a length-delimited byte sequence and decode it as UTF-8.""" length = self.varint() view = self.read(length) return str(view, "utf-8")