diff --git a/code-generator/src/main/java/io/streamnative/lightproto/generator/LightProtoField.java b/code-generator/src/main/java/io/streamnative/lightproto/generator/LightProtoField.java index 4417a08..5f0e843 100644 --- a/code-generator/src/main/java/io/streamnative/lightproto/generator/LightProtoField.java +++ b/code-generator/src/main/java/io/streamnative/lightproto/generator/LightProtoField.java @@ -78,6 +78,19 @@ public boolean isRequired() { return field.isRequired(); } + /** + * True when parsing this field invokes LightProtoCodec.readVarInt64(), the one + * reader that can advance past the message limit on truncated input. Messages + * with no such field skip the post-parse truncation check. + */ + public boolean usesVarInt64Parse() { + return isVarInt64Type(field.getProtoType()); + } + + static boolean isVarInt64Type(String protoType) { + return "int64".equals(protoType) || "uint64".equals(protoType) || "sint64".equals(protoType); + } + public void docs(PrintWriter w) { String doc = field.getDoc(); if (doc != null && !doc.isEmpty()) { diff --git a/code-generator/src/main/java/io/streamnative/lightproto/generator/LightProtoGenerator.java b/code-generator/src/main/java/io/streamnative/lightproto/generator/LightProtoGenerator.java index c8f4a2a..2f10958 100644 --- a/code-generator/src/main/java/io/streamnative/lightproto/generator/LightProtoGenerator.java +++ b/code-generator/src/main/java/io/streamnative/lightproto/generator/LightProtoGenerator.java @@ -89,6 +89,33 @@ public static List generate(List descriptors, File ou // Include the coded class once per every generated java package for (String javaPackage : javaPackages) { + // Unchecked ByteBuf accessor companion: lives in the io.netty.buffer + // package (to reach AbstractByteBuf's protected members) and is named + // after the generated package so that multiple generated packages in + // one project don't collide. + String accessClassName = "LightProtoByteBufAccess_" + javaPackage.replace('.', '_'); + try (InputStream is = LightProtoGenerator.class.getResourceAsStream( + "/io/streamnative/lightproto/generator/LightProtoByteBufAccess.java")) { + if (is == null) { + throw new IllegalStateException( + "LightProtoByteBufAccess.java template not found on the classpath"); + } + String accessSource = new String(is.readAllBytes(), java.nio.charset.StandardCharsets.UTF_8) + .replace("LightProtoByteBufAccessTemplate", accessClassName); + Path nettyDir = Paths.get(String.format("%s/io/netty/buffer", outputDirectory)); + Files.createDirectories(nettyDir); + Path accessFile = nettyDir.resolve(accessClassName + ".java"); + try (Writer w = Files.newBufferedWriter(accessFile)) { + w.write(String.format( + "/*%n" + + " * Generated by LightProto v%s%n" + + " * DO NOT MODIFY%n" + + " * DO NOT CHECK IN TO SOURCE CONTROL%n" + + " */%n", VERSION)); + w.write(accessSource); + } + } + try (InputStream is = LightProtoGenerator.class.getResourceAsStream("/io/streamnative/lightproto/generator/LightProtoCodec.java")) { JavaClassSource codecClass = (JavaClassSource) Roaster.parse(is); codecClass.setPackage(javaPackage); @@ -102,7 +129,7 @@ public static List generate(List descriptors, File ou " * DO NOT MODIFY%n" + " * DO NOT CHECK IN TO SOURCE CONTROL%n" + " */%n", VERSION)); - w.write(codecClass.toString()); + w.write(codecClass.toString().replace("LightProtoByteBufAccessTemplate", accessClassName)); } } } diff --git a/code-generator/src/main/java/io/streamnative/lightproto/generator/LightProtoMapField.java b/code-generator/src/main/java/io/streamnative/lightproto/generator/LightProtoMapField.java index e0ea51b..95da500 100644 --- a/code-generator/src/main/java/io/streamnative/lightproto/generator/LightProtoMapField.java +++ b/code-generator/src/main/java/io/streamnative/lightproto/generator/LightProtoMapField.java @@ -30,6 +30,11 @@ public LightProtoMapField(ProtoFieldDescriptor field, int index) { this.valueField = field.getMapValueField(); } + @Override + public boolean usesVarInt64Parse() { + return isVarInt64Type(keyField.getProtoType()) || isVarInt64Type(valueField.getProtoType()); + } + private boolean isStringKey() { return keyField.isStringField(); } diff --git a/code-generator/src/main/java/io/streamnative/lightproto/generator/LightProtoMessage.java b/code-generator/src/main/java/io/streamnative/lightproto/generator/LightProtoMessage.java index 7367b1e..1829821 100644 --- a/code-generator/src/main/java/io/streamnative/lightproto/generator/LightProtoMessage.java +++ b/code-generator/src/main/java/io/streamnative/lightproto/generator/LightProtoMessage.java @@ -173,6 +173,17 @@ private void generateParseFrom(PrintWriter w) { w.format(" LightProtoCodec.skipUnknownField(_tag, _buffer);\n"); w.format(" }\n"); w.format(" }\n"); + // A truncated varint64 is the one read that can consume past the message + // limit without throwing (every other reader is bounds-checked), so the + // check that restores the throw-on-truncated-input contract is emitted + // only for messages that parse 64-bit varints. Messages without them keep + // parseFrom() byte-identical to the fully checked version: even a dead + // extra check measurably perturbs code generation on small messages. + if (fields.stream().anyMatch(LightProtoField::usesVarInt64Parse)) { + w.format(" if (_buffer.readerIndex() > _endIdx) {\n"); + w.format(" throw new IndexOutOfBoundsException(\"Truncated protobuf message\");\n"); + w.format(" }\n"); + } if (hasRequiredFields()) { w.format(" checkRequiredFields();\n"); } diff --git a/code-generator/src/main/resources/io/streamnative/lightproto/generator/LightProtoByteBufAccess.java b/code-generator/src/main/resources/io/streamnative/lightproto/generator/LightProtoByteBufAccess.java new file mode 100644 index 0000000..7c0213e --- /dev/null +++ b/code-generator/src/main/resources/io/streamnative/lightproto/generator/LightProtoByteBufAccess.java @@ -0,0 +1,115 @@ +/** + * Copyright 2026 StreamNative + * + * 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. + */ +package io.netty.buffer; + +/** + * Placed in the {@code io.netty.buffer} package to reach {@link AbstractByteBuf}'s + * protected unchecked accessors ({@code _getByte}) and its {@code readerIndex} + * field directly. The per-byte {@code checkReadableBytes0()} / + * {@code ensureAccessible()} checks that dominate parse profiles are skipped, and + * each varint performs a single {@code readerIndex} store instead of one per byte. + * + *

Reads are not bounds-checked here: on truncated input a read may advance up + * to 9 bytes past the intended message limit (a 10-byte read whose first byte is + * inside the limit). The generated {@code parseFrom()} detects that afterwards + * and throws, preserving the throw-on-truncated contract. + * + *

NOTE: this class creates a split package with netty-buffer under the JPMS + * module path; it is intended for classpath deployments. + */ +public final class LightProtoByteBufAccessTemplate { + + private LightProtoByteBufAccessTemplate() { + } + + /** + * Read a varint of up to 10 bytes as a long. + * + *

The 1-2 byte cases are decoded here so the method stays far below C2's + * FreqInlineSize: the full nested chain compiles to ~371 bytecodes and is + * never inlined, so routing short varints through it costs a call per read — + * a net loss on small values. Longer varints fall into the chain, whose + * multi-byte decode amortizes the call. + */ + public static long readVarInt64Unchecked(AbstractByteBuf buf) { + int i = buf.readerIndex; + byte b0 = buf._getByte(i); + if (b0 >= 0) { + buf.readerIndex = i + 1; + return b0; + } + byte b1 = buf._getByte(i + 1); + if (b1 >= 0) { + buf.readerIndex = i + 2; + return (b0 & 0x7f) | ((long) b1 << 7); + } + return readVarInt64UncheckedChain(buf); + } + + private static long readVarInt64UncheckedChain(AbstractByteBuf buf) { + int i = buf.readerIndex; + long result; + byte tmp = buf._getByte(i++); + if (tmp >= 0) { + buf.readerIndex = i; + return tmp; + } + result = tmp & 0x7fL; + if ((tmp = buf._getByte(i++)) >= 0) { + result |= (long) tmp << 7; + } else { + result |= (tmp & 0x7fL) << 7; + if ((tmp = buf._getByte(i++)) >= 0) { + result |= (long) tmp << 14; + } else { + result |= (tmp & 0x7fL) << 14; + if ((tmp = buf._getByte(i++)) >= 0) { + result |= (long) tmp << 21; + } else { + result |= (tmp & 0x7fL) << 21; + if ((tmp = buf._getByte(i++)) >= 0) { + result |= (long) tmp << 28; + } else { + result |= (tmp & 0x7fL) << 28; + if ((tmp = buf._getByte(i++)) >= 0) { + result |= (long) tmp << 35; + } else { + result |= (tmp & 0x7fL) << 35; + if ((tmp = buf._getByte(i++)) >= 0) { + result |= (long) tmp << 42; + } else { + result |= (tmp & 0x7fL) << 42; + if ((tmp = buf._getByte(i++)) >= 0) { + result |= (long) tmp << 49; + } else { + result |= (tmp & 0x7fL) << 49; + if ((tmp = buf._getByte(i++)) >= 0) { + result |= (long) tmp << 56; + } else { + result |= (tmp & 0x7fL) << 56; + result |= ((long) buf._getByte(i++)) << 63; + } + } + } + } + } + } + } + } + buf.readerIndex = i; + return result; + } +} diff --git a/code-generator/src/main/resources/io/streamnative/lightproto/generator/LightProtoCodec.java b/code-generator/src/main/resources/io/streamnative/lightproto/generator/LightProtoCodec.java index f2840a0..62e01ad 100644 --- a/code-generator/src/main/resources/io/streamnative/lightproto/generator/LightProtoCodec.java +++ b/code-generator/src/main/resources/io/streamnative/lightproto/generator/LightProtoCodec.java @@ -15,6 +15,7 @@ */ package io.streamnative.lightproto.generator; +import io.netty.buffer.AbstractByteBuf; import io.netty.buffer.ByteBuf; import io.netty.buffer.ByteBufUtil; @@ -251,6 +252,20 @@ static int readVarInt(ByteBuf buf) { } static long readVarInt64(ByteBuf buf) { + // Only 64-bit varints use the unchecked path: it skips the per-byte + // bounds/accessibility checks (+31% on varint64-heavy messages). A + // truncated varint64 can overrun the message limit by at most 9 bytes; + // the generated parseFrom() detects that afterwards. The checked + // fallback lives in its own method so this one stays small enough to + // always inline (with the chain inline it is 303 bytecodes and C2 + // refuses it at hot call sites). + if (buf instanceof AbstractByteBuf) { + return io.netty.buffer.LightProtoByteBufAccessTemplate.readVarInt64Unchecked((AbstractByteBuf) buf); + } + return readVarInt64Checked(buf); + } + + private static long readVarInt64Checked(ByteBuf buf) { long result; byte tmp = buf.readByte(); if (tmp >= 0) { diff --git a/tests/src/main/proto/proto3.proto b/tests/src/main/proto/proto3.proto index 30ba20b..a1f7740 100644 --- a/tests/src/main/proto/proto3.proto +++ b/tests/src/main/proto/proto3.proto @@ -45,3 +45,9 @@ message Proto3Message { // Map map string_to_int = 16; } + +// Exercises the varint64 detection through map key/value types: the parser +// must emit the post-parse truncation check for this message. +message MapInt64Message { + map counters = 1; +} diff --git a/tests/src/test/java/io/streamnative/lightproto/tests/TruncatedInputTest.java b/tests/src/test/java/io/streamnative/lightproto/tests/TruncatedInputTest.java new file mode 100644 index 0000000..8967aa1 --- /dev/null +++ b/tests/src/test/java/io/streamnative/lightproto/tests/TruncatedInputTest.java @@ -0,0 +1,132 @@ +/** + * Copyright 2026 StreamNative + * + * 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. + */ +package io.streamnative.lightproto.tests; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; + +import io.netty.buffer.ByteBuf; +import io.netty.buffer.PooledByteBufAllocator; +import io.netty.buffer.Unpooled; +import org.junit.jupiter.api.Test; + +/** + * Truncated input must throw IndexOutOfBoundsException: heap buffers fail + * safely on array bounds, and for pooled direct buffers the unchecked readers + * overrun by at most 9 bytes inside the arena chunk before the post-parse + * limit check detects it. (Exactly-sized unpooled unsafe direct buffers are + * outside this guarantee — see the readVarInt64() notes in LightProtoCodec.) + */ +public class TruncatedInputTest { + + // int_field tag followed by a dangling continuation byte + private static final byte[] TRUNCATED_VARINT = {0x08, (byte) 0x80}; + // long_field (tag 0x10) truncated mid-varint64 + private static final byte[] TRUNCATED_VARINT64 = {0x10, (byte) 0xFF, (byte) 0xFF}; + // double_field (tag 0x21) with only 3 of 8 bytes + private static final byte[] TRUNCATED_FIXED64 = {0x21, 0x00, 0x00, 0x00}; + + private static void assertThrowsOnAllBufferTypes(byte[] data) { + // Heap, exact size + assertThrows(IndexOutOfBoundsException.class, + () -> new Proto3Message().parseFrom(data)); + + // Pooled direct (capacity slack: overrun reads stay in the chunk, the + // post-parse limit check throws) + ByteBuf pooled = PooledByteBufAllocator.DEFAULT.directBuffer(data.length); + try { + pooled.writeBytes(data); + assertThrows(IndexOutOfBoundsException.class, + () -> new Proto3Message().parseFrom(pooled, pooled.readableBytes())); + } finally { + pooled.release(); + } + + } + + @Test + public void testTruncatedVarint() { + assertThrowsOnAllBufferTypes(TRUNCATED_VARINT); + } + + @Test + public void testTruncatedVarint64() { + assertThrowsOnAllBufferTypes(TRUNCATED_VARINT64); + } + + @Test + public void testTruncatedFixed64() { + assertThrowsOnAllBufferTypes(TRUNCATED_FIXED64); + } + + @Test + public void testTruncatedLengthDelimited() { + // string_field (tag 0x32) claims 5 bytes but only 2 are present: + // the checked skipBytes() catches this + byte[] data = {0x32, 0x05, 'a', 'b'}; + assertThrowsOnAllBufferTypes(data); + } + + @Test + public void testTruncatedVarint64InMapValue() { + // MapInt64Message has no top-level 64-bit field: only the map value type + // routes through readVarInt64(), so this exercises the check emission + // driven by map key/value types. counters entry: key "k", value varint64 + // truncated at the end of the buffer. + byte[] data = {0x0A, 0x06, 0x0A, 0x01, 'k', 0x10, (byte) 0xFF, (byte) 0xFF}; + + assertThrows(IndexOutOfBoundsException.class, + () -> new MapInt64Message().parseFrom(data)); + + ByteBuf pooled = PooledByteBufAllocator.DEFAULT.directBuffer(data.length); + try { + pooled.writeBytes(data); + assertThrows(IndexOutOfBoundsException.class, + () -> new MapInt64Message().parseFrom(pooled, pooled.readableBytes())); + } finally { + pooled.release(); + } + + // Sanity: the same entry with a complete varint64 parses + byte[] valid = {0x0A, 0x06, 0x0A, 0x01, 'k', 0x10, (byte) 0xB9, 0x60}; + MapInt64Message parsed = new MapInt64Message(); + parsed.parseFrom(valid); + assertEquals(1, parsed.getCountersCount()); + } + + @Test + public void testValidMessageOnTightCapacityBuffer() { + // A valid message on a zero-slack buffer never overruns + Proto3Message src = new Proto3Message(); + src.setIntField(42); + src.setLongField(123456789L); + src.setStringField("tight-capacity"); + byte[] data = src.toByteArray(); + + ByteBuf tight = Unpooled.directBuffer(data.length, data.length); + try { + tight.writeBytes(data); + Proto3Message parsed = new Proto3Message(); + parsed.parseFrom(tight, tight.readableBytes()); + assertEquals(42, parsed.getIntField()); + assertEquals(123456789L, parsed.getLongField()); + assertEquals("tight-capacity", parsed.getStringField()); + assertEquals(0, tight.readableBytes()); + } finally { + tight.release(); + } + } +}