spring/spring-framework
1
0
Fork 0
mirror of https://github.com/spring-projects/spring-framework.git synced 2026-09-20 04:09:04 +00:00
Code Issues Projects Releases Wiki Activity

Allow Protobuf codec extensions

Browse Source 
Closes gh-35403
This commit is contained in:
rstoyanchev
2025-10-23 12:16:52 +01:00
parent 6ebb2071f5
commit 9331e1e86c
3 changed files with 154 additions and 79 deletions
Download Patch File Download Diff File Expand all files Collapse all files
spring-web/src/main/java/org/springframework/http/codec/protobuf/ProtobufDecoder.java
+107 -59
View File
@@ -129,13 +129,22 @@ public class ProtobufDecoder extends ProtobufCodecSupport implements Decoder<Mes
public Flux<Message> decode(Publisher<DataBuffer> inputStream, ResolvableType elementType,
@Nullable MimeType mimeType, @Nullable Map<String, Object> hints) {
MessageDecoderFunction decoderFunction = new MessageDecoderFunction(elementType, this.maxMessageSize);
MessageDecoderFunction decoderFunction =
new MessageDecoderFunction(elementType, this.maxMessageSize, initMessageSizeReader());
return Flux.from(inputStream)
.flatMapIterable(decoderFunction)
.doOnTerminate(decoderFunction::discard);
}
/**
* Return a reader for message size information encoded in the input stream.
* @since 7.0
*/
protected MessageSizeReader initMessageSizeReader() {
return new DefaultMessageSizeReader();
}
@Override
public Mono<Message> decodeToMono(Publisher<DataBuffer> inputStream, ResolvableType elementType,
@Nullable MimeType mimeType, @Nullable Map<String, Object> hints) {
@@ -150,9 +159,7 @@ public class ProtobufDecoder extends ProtobufCodecSupport implements Decoder<Mes
try {
Message.Builder builder = getMessageBuilder(targetType.toClass());
ByteBuffer byteBuffer = ByteBuffer.allocate(dataBuffer.readableByteCount());
dataBuffer.toByteBuffer(byteBuffer);
builder.mergeFrom(CodedInputStream.newInstance(byteBuffer), this.extensionRegistry);
merge(dataBuffer, builder);
return builder.build();
}
catch (IOException ex) {
@@ -166,6 +173,17 @@ public class ProtobufDecoder extends ProtobufCodecSupport implements Decoder<Mes
}
}
/**
* Use merge methods on {@link Message.Builder} to read a single message
* from the given {@code DataBuffer}.
* @since 7.0
*/
protected void merge(DataBuffer dataBuffer, Message.Builder builder) throws IOException {
ByteBuffer byteBuffer = ByteBuffer.allocate(dataBuffer.readableByteCount());
dataBuffer.toByteBuffer(byteBuffer);
builder.mergeFrom(CodedInputStream.newInstance(byteBuffer), this.extensionRegistry);
}
/**
* Create a new {@code Message.Builder} instance for the given class.
@@ -196,15 +214,14 @@ public class ProtobufDecoder extends ProtobufCodecSupport implements Decoder<Mes
private int messageBytesToRead;
private int offset;
private final MessageSizeReader messageSizeReader;
public MessageDecoderFunction(ResolvableType elementType, int maxMessageSize) {
public MessageDecoderFunction(ResolvableType elementType, int maxMessageSize, MessageSizeReader messageSizeReader) {
this.elementType = elementType;
this.maxMessageSize = maxMessageSize;
this.messageSizeReader = messageSizeReader;
}
@Override
public Iterable<? extends Message> apply(DataBuffer input) {
try {
@@ -214,9 +231,11 @@ public class ProtobufDecoder extends ProtobufCodecSupport implements Decoder<Mes
do {
if (this.output == null) {
if (!readMessageSize(input)) {
Integer messageSize = this.messageSizeReader.readMessageSize(input);
if (messageSize == null) {
return messages;
}
this.messageBytesToRead = messageSize;
if (this.maxMessageSize > 0 && this.messageBytesToRead > this.maxMessageSize) {
throw new DataBufferLimitException(
"The number of bytes to read for message " +
@@ -262,56 +281,6 @@ public class ProtobufDecoder extends ProtobufCodecSupport implements Decoder<Mes
}
}
/**
* Parse message size as a varint from the input stream, updating {@code messageBytesToRead} and
* {@code offset} fields if needed to allow processing of upcoming chunks.
* Inspired from {@link CodedInputStream#readRawVarint32(int, java.io.InputStream)}
* @return {@code true} when the message size is parsed successfully, {@code false} when the message size is
* truncated
* @see <a href="https://developers.google.com/protocol-buffers/docs/encoding#varints">Base 128 Varints</a>
*/
private boolean readMessageSize(DataBuffer input) {
if (this.offset == 0) {
if (input.readableByteCount() == 0) {
return false;
}
int firstByte = input.read();
if ((firstByte & 0x80) == 0) {
this.messageBytesToRead = firstByte;
return true;
}
this.messageBytesToRead = firstByte & 0x7f;
this.offset = 7;
}
if (this.offset < 32) {
for (; this.offset < 32; this.offset += 7) {
if (input.readableByteCount() == 0) {
return false;
}
final int b = input.read();
this.messageBytesToRead |= (b & 0x7f) << this.offset;
if ((b & 0x80) == 0) {
this.offset = 0;
return true;
}
}
}
// Keep reading up to 64 bits.
for (; this.offset < 64; this.offset += 7) {
if (input.readableByteCount() == 0) {
return false;
}
final int b = input.read();
if ((b & 0x80) == 0) {
this.offset = 0;
return true;
}
}
this.offset = 0;
throw new DecodingException("Cannot parse message size: malformed varint");
}
public void discard() {
if (this.output != null) {
DataBufferUtils.release(this.output);
@@ -319,4 +288,83 @@ public class ProtobufDecoder extends ProtobufCodecSupport implements Decoder<Mes
}
}
/**
* Component to read the size of a message. Implementations must be
* stateful and expect size information is potentially split
* across input chunks.
* @since 7.0
*/
protected interface MessageSizeReader {
/**
* Read the message size from the given buffer. This method may be
* called multiple times before the message size is fully read.
* @return return the message size, or {@code null} if the data in the
* input buffer was insufficient
*/
@Nullable Integer readMessageSize(DataBuffer input);
}
/**
* Default reader for Protobuf messages.
* <p>Parses the message size as a varint from the input stream.
* Inspired by {@link CodedInputStream#readRawVarint32(int, java.io.InputStream)},
* @see <a href="https://developers.google.com/protocol-buffers/docs/encoding#varints">Base 128 Varints</a>
*/
private static class DefaultMessageSizeReader implements MessageSizeReader {
private int offset;
private int messageSize;
@Override
public @Nullable Integer readMessageSize(DataBuffer input) {
if (this.offset == 0) {
if (input.readableByteCount() == 0) {
return null;
}
int firstByte = input.read();
if ((firstByte & 0x80) == 0) {
this.messageSize = firstByte;
return getAndReset();
}
this.messageSize = firstByte & 0x7f;
this.offset = 7;
}
if (this.offset < 32) {
for (; this.offset < 32; this.offset += 7) {
if (input.readableByteCount() == 0) {
return null;
}
final int b = input.read();
this.messageSize |= (b & 0x7f) << this.offset;
if ((b & 0x80) == 0) {
return getAndReset();
}
}
}
// Keep reading up to 64 bits.
for (; this.offset < 64; this.offset += 7) {
if (input.readableByteCount() == 0) {
return null;
}
final int b = input.read();
if ((b & 0x80) == 0) {
return getAndReset();
}
}
getAndReset();
throw new DecodingException("Cannot parse message size: malformed varint");
}
private @Nullable Integer getAndReset() {
Integer result = (this.messageSize != 0 ? this.messageSize : null);
this.offset = 0;
this.messageSize = 0;
return result;
}
}
}
spring-web/src/main/java/org/springframework/http/codec/protobuf/ProtobufEncoder.java
+18 -8
View File
@@ -107,15 +107,10 @@ public class ProtobufEncoder extends ProtobufCodecSupport implements HttpMessage
}
private DataBuffer encodeValue(Message message, DataBufferFactory bufferFactory, boolean delimited) {
FastByteArrayOutputStream bos = new FastByteArrayOutputStream();
FastByteArrayOutputStream outputStream = new FastByteArrayOutputStream();
try {
if (delimited) {
message.writeDelimitedTo((OutputStream) bos);
}
else {
message.writeTo((OutputStream) bos);
}
byte[] bytes = bos.toByteArrayUnsafe();
writeMessage(message, delimited, outputStream);
byte[] bytes = outputStream.toByteArrayUnsafe();
return bufferFactory.wrap(bytes);
}
catch (IOException ex) {
@@ -123,4 +118,19 @@ public class ProtobufEncoder extends ProtobufCodecSupport implements HttpMessage
}
}
/**
* Use write methods on {@link Message} to write to the given {@code OutputStream}.
* @since 7.0
*/
protected void writeMessage(
Message message, boolean delimited, OutputStream outputStream) throws IOException {
if (delimited) {
message.writeDelimitedTo(outputStream);
}
else {
message.writeTo(outputStream);
}
}
}
spring-web/src/main/java/org/springframework/http/codec/protobuf/ProtobufHttpMessageWriter.java
+29 -12
View File
@@ -77,30 +77,47 @@ public class ProtobufHttpMessageWriter extends EncoderHttpMessageWriter<Message>
@SuppressWarnings("unchecked")
@Override
public Mono<Void> write(Publisher<? extends Message> inputStream, ResolvableType elementType,
@Nullable MediaType mediaType, ReactiveHttpOutputMessage message, Map<String, Object> hints) {
@Nullable MediaType mediaType, ReactiveHttpOutputMessage outputMessage, Map<String, Object> hints) {
try {
Message.Builder builder = getMessageBuilder(elementType.toClass());
Descriptors.Descriptor descriptor = builder.getDescriptorForType();
message.getHeaders().add(X_PROTOBUF_SCHEMA_HEADER, descriptor.getFile().getName());
message.getHeaders().add(X_PROTOBUF_MESSAGE_HEADER, descriptor.getFullName());
outputMessage.getHeaders().add(X_PROTOBUF_SCHEMA_HEADER, descriptor.getFile().getName());
outputMessage.getHeaders().add(X_PROTOBUF_MESSAGE_HEADER, descriptor.getFullName());
if (inputStream instanceof Flux) {
if (mediaType == null) {
message.getHeaders().setContentType(((HttpMessageEncoder<?>)getEncoder()).getStreamingMediaTypes().get(0));
}
else if (!ProtobufEncoder.DELIMITED_VALUE.equals(mediaType.getParameters().get(ProtobufEncoder.DELIMITED_KEY))) {
Map<String, String> parameters = new HashMap<>(mediaType.getParameters());
parameters.put(ProtobufEncoder.DELIMITED_KEY, ProtobufEncoder.DELIMITED_VALUE);
message.getHeaders().setContentType(new MediaType(mediaType.getType(), mediaType.getSubtype(), parameters));
}
outputMessage.getHeaders().setContentType(getStreamingContentType(mediaType));
}
return super.write(inputStream, elementType, mediaType, message, hints);
extendHeaders(outputMessage, hints);
return super.write(inputStream, elementType, mediaType, outputMessage, hints);
}
catch (Exception ex) {
return Mono.error(new EncodingException("Could not write Protobuf message: " + ex.getMessage(), ex));
}
}
/**
* Return the {@code MediaType} to use when the input Publisher is multivalued.
* @since 7.0
*/
protected MediaType getStreamingContentType(@Nullable MediaType mediaType) {
if (mediaType == null) {
return ((HttpMessageEncoder<?>) getEncoder()).getStreamingMediaTypes().get(0);
}
Map<String, String> params = new HashMap<>(mediaType.getParameters());
if (!ProtobufEncoder.DELIMITED_VALUE.equals(params.get(ProtobufEncoder.DELIMITED_KEY))) {
params.put(ProtobufEncoder.DELIMITED_KEY, ProtobufEncoder.DELIMITED_VALUE);
mediaType = new MediaType(mediaType, params);
}
return mediaType;
}
/**
* Make further updates to headers.
* @since 7.0
*/
protected void extendHeaders(ReactiveHttpOutputMessage message, Map<String, Object> hints) {
}
/**
* Create a new {@code Message.Builder} instance for the given class.
* <p>This method uses a ConcurrentHashMap for caching method lookups.