Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -39,10 +39,15 @@
public class WireFormatInfo implements Command, MarshallAware {

public static final byte DATA_STRUCTURE_TYPE = CommandTypes.WIREFORMAT_INFO;
private static final int MAX_PROPERTY_SIZE = 1024 * 4;
private static final byte MAGIC[] = new byte[] {'A', 'c', 't', 'i', 'v', 'e', 'M', 'Q'};

protected byte magic[] = MAGIC;
// Max number of properties allowed in the map is 64
static final int MAX_PROPERTY_SIZE = 64;
// Used to validate property values that allocate buffers, limit to 512 bytes
static final int MAX_PROPERTY_BUFFER_SIZE = 512;
// Do not allow any nested collections in properties
static final int MAX_PROPERTY_DEPTH = 0;
private static final byte[] MAGIC = new byte[] {'A', 'c', 't', 'i', 'v', 'e', 'M', 'Q'};

protected byte[] magic = MAGIC;
protected int version;
protected ByteSequence marshalledProperties;

Expand Down Expand Up @@ -136,7 +141,7 @@ public Object getProperty(String name) throws IOException {
if (marshalledProperties == null) {
return null;
}
properties = unmarsallProperties(marshalledProperties);
properties = unmarshalProperties(marshalledProperties);
}
return properties.get(name);
}
Expand All @@ -147,7 +152,7 @@ public Map<String, Object> getProperties() throws IOException {
if (marshalledProperties == null) {
return Collections.EMPTY_MAP;
}
properties = unmarsallProperties(marshalledProperties);
properties = unmarshalProperties(marshalledProperties);
}
return Collections.unmodifiableMap(properties);
}
Expand All @@ -167,14 +172,15 @@ protected void lazyCreateProperties() throws IOException {
if (marshalledProperties == null) {
properties = new HashMap<String, Object>();
} else {
properties = unmarsallProperties(marshalledProperties);
properties = unmarshalProperties(marshalledProperties);
marshalledProperties = null;
}
}
}

private Map<String, Object> unmarsallProperties(ByteSequence marshalledProperties) throws IOException {
return MarshallingSupport.unmarshalPrimitiveMap(new DataInputStream(new ByteArrayInputStream(marshalledProperties)), MAX_PROPERTY_SIZE);
private Map<String, Object> unmarshalProperties(ByteSequence marshalledProperties) throws IOException {
return MarshallingSupport.unmarshalPrimitiveMap(new DataInputStream(new ByteArrayInputStream(marshalledProperties)),
MAX_PROPERTY_SIZE, MAX_PROPERTY_BUFFER_SIZE, MAX_PROPERTY_DEPTH);
}

@Override
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -67,15 +67,11 @@ public static void marshalPrimitiveMap(Map<String, Object> map, DataOutputStream
}

public static Map<String, Object> unmarshalPrimitiveMap(DataInputStream in) throws IOException {
return unmarshalPrimitiveMap(in, Integer.MAX_VALUE);
return unmarshalPrimitiveMap(in, Integer.MAX_VALUE, Integer.MAX_VALUE, Integer.MAX_VALUE);
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If possible consider capping the max size from these methods to in.available() to that if the encoded size is bogus you will fail early vs large allocation attempt. Unsure if this is possible if these are read via a socket input stream which is expected to block on reads waiting more bytes

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Same as other 2 comments to close the loop, I'll investigate as a follow on and look at other places using the util.

}

public static Map<String, Object> unmarshalPrimitiveMap(DataInputStream in, boolean force) throws IOException {
return unmarshalPrimitiveMap(in, Integer.MAX_VALUE, force);
}

public static Map<String, Object> unmarshalPrimitiveMap(DataInputStream in, int maxPropertySize) throws IOException {
return unmarshalPrimitiveMap(in, maxPropertySize, false);
public static Map<String, Object> unmarshalPrimitiveMap(DataInputStream in, int maxPropertySize, int maxBufferSize, int maxDepth) throws IOException {
return unmarshalPrimitiveMap(in, false, maxPropertySize, maxBufferSize, maxDepth, 0);
}

/**
Expand All @@ -84,18 +80,23 @@ public static Map<String, Object> unmarshalPrimitiveMap(DataInputStream in, int
* @throws IOException
* @throws IOException
*/
public static Map<String, Object> unmarshalPrimitiveMap(DataInputStream in, int maxPropertySize, boolean force) throws IOException {
private static Map<String, Object> unmarshalPrimitiveMap(DataInputStream in, boolean force, int maxPropertySize, int maxBufferSize,
int maxDepth, int currentDepth) throws IOException {
// increment after validation, so future calls get the incremented depth
validateDepth(maxDepth, currentDepth++);

int size = in.readInt();
if (size > maxPropertySize) {
throw new IOException("Primitive map is larger than the allowed size: " + size);
}
if (size < 0) {
return null;
} else {
Map<String, Object> rc = new HashMap<String, Object>(size);
// Size here was already validated above
Map<String, Object> rc = new HashMap<>(size);
for (int i = 0; i < size; i++) {
String name = in.readUTF();
rc.put(name, unmarshalPrimitive(in, force));
rc.put(name, unmarshalPrimitive(in, force, maxPropertySize, maxBufferSize, maxDepth, currentDepth));
}
return rc;
}
Expand All @@ -108,15 +109,20 @@ public static void marshalPrimitiveList(List<Object> list, DataOutputStream out)
}
}

public static List<Object> unmarshalPrimitiveList(DataInputStream in) throws IOException {
return unmarshalPrimitiveList(in, false);
public static List<Object> unmarshalPrimitiveList(DataInputStream in, int maxPropertySize, int maxBufferSize,
int maxDepth) throws IOException {
return unmarshalPrimitiveList(in, false, maxPropertySize, maxBufferSize, maxDepth, 0);
}

public static List<Object> unmarshalPrimitiveList(DataInputStream in, boolean force) throws IOException {
int size = in.readInt();
List<Object> answer = new ArrayList<Object>(size);
private static List<Object> unmarshalPrimitiveList(DataInputStream in, boolean force, int maxPropertySize, int maxBufferSize,
int maxDepth, int currentDepth) throws IOException {
// increment after validation, so future calls get the incremented depth
validateDepth(maxDepth, currentDepth++);

int size = validateBufferSize(maxBufferSize, in.readInt());
List<Object> answer = new ArrayList<>(size);
while (size-- > 0) {
answer.add(unmarshalPrimitive(in, force));
answer.add(unmarshalPrimitive(in, force, maxPropertySize, maxBufferSize, maxDepth, currentDepth));
}
return answer;
}
Expand Down Expand Up @@ -158,10 +164,11 @@ public static void marshalPrimitive(DataOutputStream out, Object value) throws I
}

public static Object unmarshalPrimitive(DataInputStream in) throws IOException {
return unmarshalPrimitive(in, false);
return unmarshalPrimitive(in, false, Integer.MAX_VALUE, Integer.MAX_VALUE, Integer.MAX_VALUE, 0);
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Another case where I'd consider if the passed size limit could be the stream.available()

Copy link
Copy Markdown
Contributor Author

@cshannon cshannon Jun 5, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

}

public static Object unmarshalPrimitive(DataInputStream in, boolean force) throws IOException {
private static Object unmarshalPrimitive(DataInputStream in, boolean force, int maxPropertySize, int maxBufferSize,
int maxDepth, int currentDepth) throws IOException {
Object value = null;
byte type = in.readByte();
switch (type) {
Expand Down Expand Up @@ -190,29 +197,29 @@ public static Object unmarshalPrimitive(DataInputStream in, boolean force) throw
value = in.readDouble();
break;
case BYTE_ARRAY_TYPE:
value = new byte[in.readInt()];
value = new byte[validateBufferSize(maxBufferSize, in.readInt())];
in.readFully((byte[])value);
break;
case STRING_TYPE:
if (force) {
value = in.readUTF();
} else {
value = readUTF(in, in.readUnsignedShort());
value = readUTF(in, maxBufferSize, in.readUnsignedShort());
}
break;
case BIG_STRING_TYPE: {
if (force) {
value = readUTF8(in);
value = readUTF8(in, maxBufferSize);
} else {
value = readUTF(in, in.readInt());
value = readUTF(in, maxBufferSize, in.readInt());
}
break;
}
case MAP_TYPE:
value = unmarshalPrimitiveMap(in, true);
value = unmarshalPrimitiveMap(in, true, maxPropertySize, maxBufferSize, maxDepth, currentDepth);
break;
case LIST_TYPE:
value = unmarshalPrimitiveList(in, true);
value = unmarshalPrimitiveList(in, true, maxPropertySize, maxBufferSize, maxDepth, currentDepth);
break;
case NULL:
value = null;
Expand All @@ -223,8 +230,9 @@ public static Object unmarshalPrimitive(DataInputStream in, boolean force) throw
return value;
}

public static UTF8Buffer readUTF(DataInputStream in, int length) throws IOException {
byte data[] = new byte[length];
public static UTF8Buffer readUTF(DataInputStream in, int maxLength, int length) throws IOException {
validateBufferSize(maxLength, length);
byte[] data = new byte[length];
in.readFully(data);
return new UTF8Buffer(data);
}
Expand Down Expand Up @@ -350,7 +358,11 @@ public static int writeUTFBytesToBuffer(String str, long count,
}

public static String readUTF8(DataInput dataIn) throws IOException {
int utflen = dataIn.readInt();
return readUTF8(dataIn, Integer.MAX_VALUE);
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This one is a bit dangerous so should check on who is calling it. I might actually pass a value of dataIn.available() unless there is a case where this can be used on blocking reads like a read from a socket input stream.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

That is a good point, I agree this could be looked into for improvement as well. This is a utility class and method so I would need to look where else it could be used.

This particular method wasn't touched and is used at all by WireFormatInfo so I don't think it needs to be changed now. It could be used in the future by various use cases, also there's a chance (and likely) the stream could be a compressed stream which would impact the available() call as well. So I can look into it as a separate follow on.

}

public static String readUTF8(DataInput dataIn, int maxBufferSize) throws IOException {
int utflen = validateBufferSize(maxBufferSize, dataIn.readInt());
if (utflen > -1) {
byte bytearr[] = new byte[utflen];
char chararr[] = new char[utflen];
Expand Down Expand Up @@ -419,4 +431,17 @@ public static String truncate64(String text) {
}
return text;
}

private static int validateBufferSize(int maxSize, int size) throws IOException {
if (size > maxSize) {
throw new IOException("Max buffer size: " + maxSize + " exceeded, size: " + size);
}
return size;
}

private static void validateDepth(int maxDepth, int currentDepth) throws IOException {
if (currentDepth > maxDepth) {
throw new IOException("Max unmarshaling depth: " + maxDepth + " exceeded, depth: " + currentDepth);
}
}
}
Loading
Loading