Skip to content

Commit

Permalink
Address review comments
Browse files Browse the repository at this point in the history
  • Loading branch information
lfittl committed Nov 26, 2021
1 parent fab7498 commit 51433c5
Show file tree
Hide file tree
Showing 3 changed files with 56 additions and 30 deletions.
4 changes: 2 additions & 2 deletions ruby/ext/google/protobuf_c/message.c
Expand Up @@ -943,7 +943,7 @@ static VALUE Message_decode(int argc, VALUE* argv, VALUE klass) {
VALUE depth = rb_hash_lookup(hash_args, ID2SYM(rb_intern("max_recursion_depth")));

if (depth != Qnil && TYPE(depth) == T_FIXNUM) {
options = FIX2INT(depth) << 16;
options |= UPB_DECODE_MAXDEPTH(FIX2INT(depth));
}
}

Expand Down Expand Up @@ -1058,7 +1058,7 @@ static VALUE Message_encode(int argc, VALUE* argv, VALUE klass) {
VALUE depth = rb_hash_lookup(hash_args, ID2SYM(rb_intern("max_recursion_depth")));

if (depth != Qnil && TYPE(depth) == T_FIXNUM) {
options = FIX2INT(depth) << 16;
options |= UPB_DECODE_MAXDEPTH(FIX2INT(depth));
}
}

Expand Down
4 changes: 2 additions & 2 deletions ruby/src/main/java/com/google/protobuf/jruby/RubyMap.java
Expand Up @@ -389,7 +389,7 @@ protected IRubyObject deepCopy(ThreadContext context) {
return newMap;
}

protected List<DynamicMessage> build(ThreadContext context, RubyDescriptor descriptor, int depth) {
protected List<DynamicMessage> build(ThreadContext context, RubyDescriptor descriptor, int depth, int maxRecursionDepth) {
List<DynamicMessage> list = new ArrayList<DynamicMessage>();
RubyClass rubyClass = (RubyClass) descriptor.msgclass(context);
FieldDescriptor keyField = descriptor.getField("key");
Expand All @@ -398,7 +398,7 @@ protected List<DynamicMessage> build(ThreadContext context, RubyDescriptor descr
RubyMessage mapMessage = (RubyMessage) rubyClass.newInstance(context, Block.NULL_BLOCK);
mapMessage.setField(context, keyField, key);
mapMessage.setField(context, valueField, table.get(key));
list.add(mapMessage.build(context, depth + 1));
list.add(mapMessage.build(context, depth + 1, maxRecursionDepth));
}
return list;
}
Expand Down
78 changes: 52 additions & 26 deletions ruby/src/main/java/com/google/protobuf/jruby/RubyMessage.java
Expand Up @@ -39,6 +39,7 @@
import com.google.protobuf.Descriptors.FileDescriptor;
import com.google.protobuf.Descriptors.OneofDescriptor;
import com.google.protobuf.ByteString;
import com.google.protobuf.CodedInputStream;
import com.google.protobuf.DynamicMessage;
import com.google.protobuf.InvalidProtocolBufferException;
import com.google.protobuf.Message;
Expand Down Expand Up @@ -461,35 +462,63 @@ public static IRubyObject getDescriptor(ThreadContext context, IRubyObject recv)

/*
* call-seq:
* MessageClass.encode(msg) => bytes
* MessageClass.encode(msg, options = {}) => bytes
*
* Encodes the given message object to its serialized form in protocol buffers
* wire format.
* @param options [Hash] options for the encoder
* max_recursion_depth: set to maximum encoding depth for message (default is 64)
*/
@JRubyMethod(meta = true)
public static IRubyObject encode(ThreadContext context, IRubyObject recv, IRubyObject value) {
if (recv != value.getMetaClass()) {
throw context.runtime.newArgumentError("Tried to encode a " + value.getMetaClass() + " message with " + recv);
@JRubyMethod(required = 1, optional = 1, meta = true)
public static IRubyObject encode(ThreadContext context, IRubyObject recv, IRubyObject[] args) {
if (recv != args[0].getMetaClass()) {
throw context.runtime.newArgumentError("Tried to encode a " + args[0].getMetaClass() + " message with " + recv);
}
RubyMessage message = (RubyMessage) value;
return context.runtime.newString(new ByteList(message.build(context).toByteArray()));
RubyMessage message = (RubyMessage) args[0];
int maxRecursionDepthInt = SINK_MAXIMUM_NESTING;

if (args.length > 1) {
RubyHash options = (RubyHash) args[1];
IRubyObject maxRecursionDepth = options.fastARef(context.runtime.newSymbol("max_recursion_depth"));

if (maxRecursionDepth != null) {
maxRecursionDepthInt = ((RubyNumeric) maxRecursionDepth).getIntValue();
}
}
return context.runtime.newString(new ByteList(message.build(context, 0, maxRecursionDepthInt).toByteArray()));
}

/*
* call-seq:
* MessageClass.decode(data) => message
* MessageClass.decode(data, options = {}) => message
*
* Decodes the given data (as a string containing bytes in protocol buffers wire
* format) under the interpretation given by this message class's definition
* and returns a message object with the corresponding field values.
* @param options [Hash] options for the decoder
* max_recursion_depth: set to maximum decoding depth for message (default is 100)
*/
@JRubyMethod(meta = true)
public static IRubyObject decode(ThreadContext context, IRubyObject recv, IRubyObject data) {
@JRubyMethod(required = 1, optional = 1, meta = true)
public static IRubyObject decode(ThreadContext context, IRubyObject recv, IRubyObject[] args) {
IRubyObject data = args[0];
byte[] bin = data.convertToString().getBytes();
CodedInputStream input = CodedInputStream.newInstance(bin);
RubyMessage ret = (RubyMessage) ((RubyClass) recv).newInstance(context, Block.NULL_BLOCK);

if (args.length == 2) {
if (!(args[1] instanceof RubyHash)) {
throw context.runtime.newArgumentError("Expected hash arguments.");
}

IRubyObject maxRecursionDepth = ((RubyHash) args[1]).fastARef(context.runtime.newSymbol("max_recursion_depth"));
if (maxRecursionDepth != null) {
input.setRecursionLimit(((RubyNumeric) maxRecursionDepth).getIntValue());
}
}

try {
ret.builder.mergeFrom(bin);
} catch (InvalidProtocolBufferException e) {
ret.builder.mergeFrom(input);
} catch (Exception e) {
throw RaiseException.from(context.runtime, (RubyClass) context.runtime.getClassFromPath("Google::Protobuf::ParseError"), e.getMessage());
}

Expand Down Expand Up @@ -541,7 +570,7 @@ public static IRubyObject encodeJson(ThreadContext context, IRubyObject recv, IR
printer = printer.usingTypeRegistry(JsonFormat.TypeRegistry.newBuilder().add(message.descriptor).build());

try {
result = printer.print(message.build(context));
result = printer.print(message.build(context, 0, SINK_MAXIMUM_NESTING));
} catch (InvalidProtocolBufferException e) {
throw runtime.newRuntimeError(e.getMessage());
} catch (IllegalArgumentException e) {
Expand Down Expand Up @@ -635,12 +664,8 @@ public IRubyObject toHash(ThreadContext context) {
return ret;
}

protected DynamicMessage build(ThreadContext context) {
return build(context, 0);
}

protected DynamicMessage build(ThreadContext context, int depth) {
if (depth > SINK_MAXIMUM_NESTING) {
protected DynamicMessage build(ThreadContext context, int depth, int maxRecursionDepth) {
if (depth >= maxRecursionDepth) {
throw context.runtime.newRuntimeError("Maximum recursion depth exceeded during encoding.");
}

Expand All @@ -651,7 +676,7 @@ protected DynamicMessage build(ThreadContext context, int depth) {
if (value instanceof RubyMap) {
builder.clearField(fieldDescriptor);
RubyDescriptor mapDescriptor = (RubyDescriptor) getDescriptorForField(context, fieldDescriptor);
for (DynamicMessage kv : ((RubyMap) value).build(context, mapDescriptor, depth)) {
for (DynamicMessage kv : ((RubyMap) value).build(context, mapDescriptor, depth, maxRecursionDepth)) {
builder.addRepeatedField(fieldDescriptor, kv);
}

Expand All @@ -660,7 +685,7 @@ protected DynamicMessage build(ThreadContext context, int depth) {

builder.clearField(fieldDescriptor);
for (int i = 0; i < repeatedField.size(); i++) {
Object item = convert(context, fieldDescriptor, repeatedField.get(i), depth,
Object item = convert(context, fieldDescriptor, repeatedField.get(i), depth, maxRecursionDepth,
/*isDefaultValueForBytes*/ false);
builder.addRepeatedField(fieldDescriptor, item);
}
Expand All @@ -682,7 +707,7 @@ protected DynamicMessage build(ThreadContext context, int depth) {
fieldDescriptor.getFullName().equals("google.protobuf.FieldDescriptorProto.default_value")) {
isDefaultStringForBytes = true;
}
builder.setField(fieldDescriptor, convert(context, fieldDescriptor, value, depth, isDefaultStringForBytes));
builder.setField(fieldDescriptor, convert(context, fieldDescriptor, value, depth, maxRecursionDepth, isDefaultStringForBytes));
}
}

Expand All @@ -702,7 +727,7 @@ protected DynamicMessage build(ThreadContext context, int depth) {
builder.clearField(fieldDescriptor);
RubyDescriptor mapDescriptor = (RubyDescriptor) getDescriptorForField(context,
fieldDescriptor);
for (DynamicMessage kv : ((RubyMap) value).build(context, mapDescriptor, depth)) {
for (DynamicMessage kv : ((RubyMap) value).build(context, mapDescriptor, depth, maxRecursionDepth)) {
builder.addRepeatedField(fieldDescriptor, kv);
}
}
Expand Down Expand Up @@ -814,7 +839,8 @@ private FieldDescriptor findField(ThreadContext context, IRubyObject fieldName,
// convert a ruby object to protobuf type, skip type check since it is checked on the way in
private Object convert(ThreadContext context,
FieldDescriptor fieldDescriptor,
IRubyObject value, int depth, boolean isDefaultStringForBytes) {
IRubyObject value, int depth, int maxRecursionDepth,
boolean isDefaultStringForBytes) {
Object val = null;
switch (fieldDescriptor.getType()) {
case INT32:
Expand Down Expand Up @@ -855,7 +881,7 @@ private Object convert(ThreadContext context,
}
break;
case MESSAGE:
val = ((RubyMessage) value).build(context, depth + 1);
val = ((RubyMessage) value).build(context, depth + 1, maxRecursionDepth);
break;
case ENUM:
EnumDescriptor enumDescriptor = fieldDescriptor.getEnumType();
Expand Down Expand Up @@ -1214,7 +1240,7 @@ private void validateMessageType(ThreadContext context, FieldDescriptor fieldDes
private static final String CONST_SUFFIX = "_const";
private static final String HAS_PREFIX = "has_";
private static final String QUESTION_MARK = "?";
private static final int SINK_MAXIMUM_NESTING = 63;
private static final int SINK_MAXIMUM_NESTING = 64;

private Descriptor descriptor;
private DynamicMessage.Builder builder;
Expand Down

0 comments on commit 51433c5

Please sign in to comment.