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);
}

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);
}

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 = Double.valueOf(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);
}

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