Java Protobuf 反序列化内存边界控制策略与挑战

import com.google.protobuf.CodedInputStream;
import com.google.protobuf.InvalidProtocolBufferException;
import com.google.protobuf.Message;
import com.google.protobuf.DynamicMessage;
import com.google.protobuf.Descriptors.Descriptor;

import java.io.IOException;
import java.io.InputStream;
import java.nio.ByteBuffer;

public class ProtobufDeserializationLimiter {

    /**
     * 使用CodedInputStream限制序列化消息的最大字节数。
     *
     * @param dataStream 包含序列化Protobuf消息的输入流
     * @param descriptor 消息的描述符
     * @param maxSerializedBytes 允许的最大序列化字节数
     * @return 反序列化后的消息对象
     * @throws IOException 如果I/O操作失败或消息超出大小限制
     */
    public static Message parseMessageWithSerializedLimit(
            InputStream dataStream, Descriptor descriptor, int maxSerializedBytes) throws IOException {

        CodedInputStream codedInputStream = CodedInputStream.newInstance(dataStream);
        // 设置最大读取字节数限制
        codedInputStream.setSizeLimit(maxSerializedBytes);

        try {
            // 使用DynamicMessage进行反序列化，因为描述符可能是动态加载的
            return DynamicMessage.parseFrom(descriptor, codedInputStream);
        } catch (InvalidProtocolBufferException e) {
            // 当消息超过setSizeLimit设定的限制时，会抛出此异常
            if (e.getMessage() != null && e.getMessage().contains("size limit was exceeded")) {
                throw new IOException("Serialized message size exceeded the allowed limit of " + maxSerializedBytes + " bytes.", e);
            }
            throw e; // 其他Protobuf解析错误
        }
    }

    public static void main(String[] args) {
        // 假设我们有一个简单的Protobuf定义
        // message MyMessage {
        //   string name = 1;
        //   int32 id = 2;
        // }
        // 实际应用中，descriptor会通过FileDescriptorSet动态获取
        // 这里只是一个模拟的描述符获取过程
        Descriptor myMessageDescriptor = getExampleDescriptor(); // 模拟获取描述符

        // 模拟一个合法的短消息 (e.g., "name: 'test', id: 1")
        byte[] smallMessageBytes = ByteBuffer.allocate(10)
                .put((byte) (1 << 3 | 2)) // field 1, wire type 2 (length-delimited string)
                .put((byte) 4) // length of "test"
                .put("test".getBytes())
                .put((byte) (2 << 3 | 0)) // field 2, wire type 0 (varint)
                .put((byte) 1) // value 1
                .array();

        // 模拟一个过长的消息 (实际中可能是一个恶意构造的大消息)
        byte[] largeMessageBytes = new byte[2000]; // 超过1KB限制
        // 填充一些数据以模拟Protobuf消息
        largeMessageBytes[0] = (byte) (1 << 3 | 2); // field 1, wire type 2
        largeMessageBytes[1] = (byte) 127; // length prefix for a long string
        for (int i = 2; i < 130; i++) { // Fill part of the string
            largeMessageBytes[i] = 'a';
        }
        // 剩余部分保持0，或填充其他数据

        int maxAllowedBytes = 1024; // 1KB限制

        try {
            // 尝试解析合法消息
            InputStream smallStream = new java.io.ByteArrayInputStream(smallMessageBytes);
            Message msg1 = parseMessageWithSerializedLimit(smallStream, myMessageDescriptor, maxAllowedBytes);
            System.out.println("Successfully parsed small message: " + msg1.toString());

            // 尝试解析过长消息
            InputStream largeStream = new java.io.ByteArrayInputStream(largeMessageBytes);
            parseMessageWithSerializedLimit(largeStream, myMessageDescriptor, maxAllowedBytes);
            System.out.println("Successfully parsed large message (this should not happen)");

        } catch (IOException e) {
            System.err.println("Error parsing message: " + e.getMessage());
        }
    }

    // 模拟获取描述符的方法 (在实际应用中，这会从FileDescriptorSet中解析)
    private static Descriptor getExampleDescriptor() {
        // 这是一个非常简化的模拟，实际需要使用DescriptorProtos和DescriptorPool
        // 这里仅为示例提供一个虚拟的描述符
        try {
            // 使用Protobuf的反射机制来获取一个简单的描述符
            // 假设你有一个proto文件定义了 MyMessage
            // syntax = "proto3";
            // package com.example;
            // message MyMessage {
            //   string name = 1;
            //   int32 id = 2;
            // }
            // 你需要编译这个proto文件，然后使用生成的Java类来获取描述符
            // 例如：return com.example.MyMessage.getDescriptor();
            // 由于这里没有实际的.proto文件和生成的类，我们返回一个null或抛出异常
            // 实际应用中，你需要确保这里能获取到正确的描述符
            // 为了让示例编译通过，我们创建一个假的描述符，这在实际中不可取
            // 这是一个复杂的步骤，通常涉及 FileDescriptorSet
            System.out.println("Warning: Using a placeholder descriptor. In real applications, load from FileDescriptorSet.");
            return com.google.protobuf.DescriptorProtos.DescriptorProto.newBuilder()
                    .setName("MyMessage")
                    .addField(com.google.protobuf.DescriptorProtos.FieldDescriptorProto.newBuilder()
                            .setName("name")
                            .setNumber(1)
                            .setType(com.google.protobuf.DescriptorProtos.FieldDescriptorProto.Type.TYPE_STRING)
                            .build())
                    .addField(com.google.protobuf.DescriptorProtos.FieldDescriptorProto.newBuilder()
                            .setName("id")
                            .setNumber(2)
                            .setType(com.google.protobuf.DescriptorProtos.FieldDescriptorProto.Type.TYPE_INT32)
                            .build())
                    .build()
                    .getDescriptorForType(); // 这是一个简化的获取方式，可能不完全正确，但用于演示
        } catch (Exception e) {
            e.printStackTrace();
            return null;
        }
    }
}

message MyMessage {
  repeated string names = 1;
  repeated int32 ids = 2;
}

// 假设这是接收到的原始字节数组
byte[] receivedProtobufBytes = getReceivedBytes();

// 如果仅需转发，直接将字节数组发送到数据存储
dataStoreService.storeRawProtobuf(receivedProtobufBytes);

// 避免：
// Message parsedMessage = DynamicMessage.parseFrom(descriptor, receivedProtobufBytes);
// dataStoreService.storeParsedMessage(parsedMessage);