added function for size limit of raw, array and map data to Unpacker interface

muga · muga · commit f5bcdace2681 · 2011-11-04T02:18:15.000+09:00
diff --git a/src/main/java/org/msgpack/unpacker/AbstractUnpacker.java b/src/main/java/org/msgpack/unpacker/AbstractUnpacker.java
@@ -28,6 +28,13 @@
 public abstract class AbstractUnpacker implements Unpacker {
     protected MessagePack msgpack;
 
+    // TODO #MN and #SF about default values
+    protected long rawSizeLimit = 65536;
+
+    protected int arraySizeLimit = 65536;
+
+    protected int mapSizeLimit = 65536;
+
     protected AbstractUnpacker(MessagePack msgpack) {
 	this.msgpack = msgpack;
     }
@@ -109,4 +116,22 @@ public long getReadByteCount() {
     public void resetReadByteCount() {
         throw new UnsupportedOperationException("Not implemented yet");
     }
+
+    public void setRawSizeLimit(long size) {
+        if (size > 0 && size > 65536) {
+            rawSizeLimit = size;
+        }
+    }
+
+    public void setArraySizeLimit(int size) {
+        if (size > 0 && size > 65536) {
+            arraySizeLimit = size;
+        }
+    }
+
+    public void setMapSizeLimit(int size) {
+        if (size > 0 && size > 65536) {
+            mapSizeLimit = size;
+        }
+    }
 }
diff --git a/src/main/java/org/msgpack/unpacker/MessagePackUnpacker.java b/src/main/java/org/msgpack/unpacker/MessagePackUnpacker.java
@@ -224,8 +224,10 @@ private boolean readOneWithoutStackLarge(Accept a, final int b) throws IOExcepti
         case 0xdb:  // raw 32
             {
                 int count = in.getInt();
-                if (count < 0) {
-                    throw new IOException("Raw size too large");  // TODO error MessageSizeException
+                if (count < 0 || count > rawSizeLimit) {
+                    String reason = String.format("Size of raw (%d) over limit at %d",
+                            new Object[] { count, rawSizeLimit });
+                    throw new SizeLimitException(reason);
                 }
                 if (count == 0) {
                     a.acceptEmptyRaw();
@@ -255,8 +257,10 @@ private boolean readOneWithoutStackLarge(Accept a, final int b) throws IOExcepti
         case 0xdd:  // array 32
             {
                 int count = in.getInt();
-                if (count < 0) {
-                    throw new IOException("Array size too large");  // TODO error MessageSizeException
+                if (count < 0 || count > arraySizeLimit) {
+                    String reason = String.format("Size of array (%d) over limit at %d",
+                            new Object[] { count, arraySizeLimit });
+                    throw new SizeLimitException(reason);
                 }
                 a.acceptArray(count);
                 stack.reduceCount();
@@ -278,8 +282,10 @@ private boolean readOneWithoutStackLarge(Accept a, final int b) throws IOExcepti
         case 0xdf:  // map 32
             {
                 int count = in.getInt();
-                if (count < 0) {
-                    throw new IOException("Map size too large");  // TODO error MessageSizeException
+                if (count < 0 && count > mapSizeLimit) {
+                    String reason = String.format("Size of map (%d) over limit at %d",
+                            new Object[] { count, mapSizeLimit });
+                    throw new SizeLimitException(reason);
                 }
                 a.acceptMap(count);
                 stack.reduceCount();
diff --git a/src/main/java/org/msgpack/unpacker/SizeLimitException.java b/src/main/java/org/msgpack/unpacker/SizeLimitException.java
@@ -0,0 +1,40 @@
+//
+// MessagePack for Java
+//
+// Copyright (C) 2011 Muga Nishizawa
+//
+//    Licensed under the Apache License, Version 2.0 (the "License");
+//    you may not use this file except in compliance with the License.
+//    You may obtain a copy of the License at
+//
+//        http://www.apache.org/licenses/LICENSE-2.0
+//
+//    Unless required by applicable law or agreed to in writing, software
+//    distributed under the License is distributed on an "AS IS" BASIS,
+//    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+//    See the License for the specific language governing permissions and
+//    limitations under the License.
+//
+package org.msgpack.unpacker;
+
+import org.msgpack.MessageTypeException;
+
+
+@SuppressWarnings("serial")
+public class SizeLimitException extends MessageTypeException {
+    public SizeLimitException() {
+        super();
+    }
+
+    public SizeLimitException(String message) {
+        super(message);
+    }
+
+    public SizeLimitException(String message, Throwable cause) {
+        super(message, cause);
+    }
+
+    public SizeLimitException(Throwable cause) {
+        super(cause);
+    }
+}
diff --git a/src/main/java/org/msgpack/unpacker/Unpacker.java b/src/main/java/org/msgpack/unpacker/Unpacker.java
@@ -90,5 +90,12 @@ public interface Unpacker extends Iterable<Value>, Closeable {
     public long getReadByteCount();
 
     public void resetReadByteCount();
+
+
+    public void setRawSizeLimit(long size);
+
+    public void setArraySizeLimit(int size);
+
+    public void setMapSizeLimit(int size);
 }
 
