Add support for getting byte sizes to dex backed references

NOTICE

@@ -30,8 +30,8 @@ THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 
 
 Unless otherwise stated in the code/commit message, any changes with the
-committer of bgruv@google.com is copyrighted by Google Inc. and released
+committer of bgruv@google.com or wkal@google.com is copyrighted by 
-under the following license:
+Google Inc. and released under the following license:
 
 *******************************************************************************
 Copyright 2011, Google Inc.

dexlib2/src/main/java/org/jf/dexlib2/dexbacked/BaseDexReader.java

@@ -88,10 +88,41 @@ public class BaseDexReader<T extends BaseDexBuffer> {
         return result;
     }
 
+    public int peekSleb128Size() {
+        int end = dexBuf.baseOffset + offset;
+        int currentByteValue;
+        int result;
+        byte[] buf = dexBuf.buf;
+
+        result = buf[end++] & 0xff;
+        if (result > 0x7f) {
+            currentByteValue = buf[end++] & 0xff;
+            if (currentByteValue > 0x7f) {
+                currentByteValue = buf[end++] & 0xff;
+                if (currentByteValue > 0x7f) {
+                    currentByteValue = buf[end++] & 0xff;
+                    if (currentByteValue > 0x7f) {
+                        currentByteValue = buf[end++] & 0xff;
+                        if (currentByteValue > 0x7f) {
+                            throw new ExceptionWithContext(
+                                "Invalid sleb128 integer encountered at offset 0x%x", offset);
+                        }
+                    }
+                }
+            }
+        }
+
+        return end - (dexBuf.baseOffset + offset);
+    }
+
     public int readSmallUleb128() {
         return readUleb128(false);
     }
 
+    public int peekSmallUleb128Size() {
+        return peekUleb128Size(false);
+    }
+
     private int readUleb128(boolean allowLarge) {
         int end = dexBuf.baseOffset + offset;
         int currentByteValue;
@@ -133,6 +164,43 @@ public class BaseDexReader<T extends BaseDexBuffer> {
         return result;
     }
 
+    private int peekUleb128Size(boolean allowLarge) {
+        int end = dexBuf.baseOffset + offset;
+        int currentByteValue;
+        int result;
+        byte[] buf = dexBuf.buf;
+
+        result = buf[end++] & 0xff;
+        if (result > 0x7f) {
+            currentByteValue = buf[end++] & 0xff;
+            if (currentByteValue > 0x7f) {
+                currentByteValue = buf[end++] & 0xff;
+                if (currentByteValue > 0x7f) {
+                    currentByteValue = buf[end++] & 0xff;
+                    if (currentByteValue > 0x7f) {
+                        currentByteValue = buf[end++];
+
+                        // MSB shouldn't be set on last byte
+                        if (currentByteValue < 0) {
+                            throw new ExceptionWithContext(
+                                "Invalid uleb128 integer encountered at offset 0x%x", offset);
+                        } else if ((currentByteValue & 0xf) > 0x07) {
+                            if (!allowLarge) {
+                                // for non-large uleb128s, we assume most significant bit of the result will not be
+                                // set, so that it can fit into a signed integer without wrapping
+                                throw new ExceptionWithContext(
+                                    "Encountered valid uleb128 that is out of range at offset 0x%x", offset);
+                            }
+                        }
+                    }
+                }
+            }
+        }
+
+        return end - (dexBuf.baseOffset + offset);
+    }
+
+
     /**
      * Reads a "large" uleb128. That is, one that may legitimately be greater than a signed int.
      *
@@ -183,6 +251,35 @@ public class BaseDexReader<T extends BaseDexBuffer> {
         return result;
     }
 
+    public int peekBigUleb128Size() {
+        int end = dexBuf.baseOffset + offset;
+        int currentByteValue;
+        int result;
+        byte[] buf = dexBuf.buf;
+
+        result = buf[end++] & 0xff;
+        if (result > 0x7f) {
+            currentByteValue = buf[end++] & 0xff;
+            if (currentByteValue > 0x7f) {
+                currentByteValue = buf[end++] & 0xff;
+                if (currentByteValue > 0x7f) {
+                    currentByteValue = buf[end++] & 0xff;
+                    if (currentByteValue > 0x7f) {
+                        currentByteValue = buf[end++];
+
+                        // MSB shouldn't be set on last byte
+                        if (currentByteValue < 0) {
+                            throw new ExceptionWithContext(
+                                "Invalid uleb128 integer encountered at offset 0x%x", offset);
+                        }
+                    }
+                }
+            }
+        }
+
+        return end - (dexBuf.baseOffset + offset);
+    }
+
     public void skipUleb128() {
         int end = dexBuf.baseOffset + offset;
         byte currentByteValue;
@@ -516,4 +613,11 @@ public class BaseDexReader<T extends BaseDexBuffer> {
         offset += ret[0];
         return value;
     }
+
+    public int peekStringLength(int utf16Length) {
+        int[] ret = new int[1];
+        Utf8Utils.utf8BytesWithUtf16LengthToString(
+            dexBuf.buf, dexBuf.baseOffset + offset, utf16Length, ret);
+        return ret[0];
+    }
 }

dexlib2/src/main/java/org/jf/dexlib2/dexbacked/DexBackedClassDef.java

@@ -36,6 +36,7 @@ import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.Iterables;
 import org.jf.dexlib2.base.reference.BaseTypeReference;
 import org.jf.dexlib2.dexbacked.raw.ClassDefItem;
+import org.jf.dexlib2.dexbacked.raw.TypeIdItem;
 import org.jf.dexlib2.dexbacked.util.AnnotationsDirectory;
 import org.jf.dexlib2.dexbacked.util.FixedSizeSet;
 import org.jf.dexlib2.dexbacked.util.StaticInitialValueIterator;
@@ -434,4 +435,66 @@ public class DexBackedClassDef extends BaseTypeReference implements ClassDef {
         virtualMethodsOffset = reader.getOffset();
         return virtualMethodsOffset;
     }
+
+    /**
+     * Calculate and return the private size of a class definition.
+     *
+     * Calculated as: class_def_item size + type_id size + interfaces type_list +
+     * annotations_directory_item overhead + class_data_item + static values overhead +
+     * methods size + fields size
+     *
+     * @return size in bytes
+     */
+    public int getSize() {
+        int size = 8 * 4; //class_def_item has 8 uint fields in dex files
+        size += TypeIdItem.ITEM_SIZE; //type_ids size
+
+        //add interface list size if any
+        int interfacesLength = getInterfaces().size();
+        if (interfacesLength > 0) {
+            //add size of the type_list
+            size += 4; //uint for size
+            size += interfacesLength * 2; //ushort per type_item
+        }
+
+        //annotations directory size if it exists
+        AnnotationsDirectory directory = getAnnotationsDirectory();
+        if (!AnnotationsDirectory.EMPTY.equals(directory)) {
+            size += 4 * 4; //4 uints in annotations_directory_item
+            Set<? extends DexBackedAnnotation> classAnnotations = directory.getClassAnnotations();
+            if (!classAnnotations.isEmpty()) {
+                size += 4; //uint for size
+                size += classAnnotations.size() * 4; //uint per annotation_off
+                //TODO: should we add annotation_item size? what if it's shared?
+            }
+        }
+
+        //static values and/or metadata
+        int staticInitialValuesOffset =
+            dexFile.readSmallUint(classDefOffset + ClassDefItem.STATIC_VALUES_OFFSET);
+        if (staticInitialValuesOffset != 0) {
+            DexReader reader = dexFile.readerAt(staticInitialValuesOffset);
+            size += reader.peekSmallUleb128Size(); //encoded_array size field
+        }
+
+        //class_data_item
+        int classDataOffset = dexFile.readSmallUint(classDefOffset + ClassDefItem.CLASS_DATA_OFFSET);
+        if (classDataOffset > 0) {
+            DexReader reader = dexFile.readerAt(classDataOffset);
+            reader.readSmallUleb128(); //staticFieldCount
+            reader.readSmallUleb128(); //instanceFieldCount
+            reader.readSmallUleb128(); //directMethodCount
+            reader.readSmallUleb128(); //virtualMethodCount
+            size += reader.getOffset() - classDataOffset;
+        }
+
+        for (DexBackedField dexBackedField : getFields()) {
+            size += dexBackedField.getSize();
+        }
+
+        for (DexBackedMethod dexBackedMethod : getMethods()) {
+            size += dexBackedMethod.getSize();
+        }
+        return size;
+    }
 }

dexlib2/src/main/java/org/jf/dexlib2/dexbacked/DexBackedField.java

@@ -33,8 +33,10 @@ package org.jf.dexlib2.dexbacked;
 
 import org.jf.dexlib2.base.reference.BaseFieldReference;
 import org.jf.dexlib2.dexbacked.raw.FieldIdItem;
+import org.jf.dexlib2.dexbacked.reference.DexBackedFieldReference;
 import org.jf.dexlib2.dexbacked.util.AnnotationsDirectory;
 import org.jf.dexlib2.dexbacked.util.StaticInitialValueIterator;
+import org.jf.dexlib2.dexbacked.value.DexBackedEncodedValue;
 import org.jf.dexlib2.iface.ClassDef;
 import org.jf.dexlib2.iface.Field;
 import org.jf.dexlib2.iface.value.EncodedValue;
@@ -52,6 +54,8 @@ public class DexBackedField extends BaseFieldReference implements Field {
     public final int annotationSetOffset;
 
     public final int fieldIndex;
+    private final int startOffset;
+    private final int initialValueOffset;
 
     private int fieldIdItemOffset;
 
@@ -65,11 +69,13 @@ public class DexBackedField extends BaseFieldReference implements Field {
 
         // large values may be used for the index delta, which cause the cumulative index to overflow upon
         // addition, effectively allowing out of order entries.
+        startOffset = reader.getOffset();
         int fieldIndexDiff = reader.readLargeUleb128();
         this.fieldIndex = fieldIndexDiff + previousFieldIndex;
         this.accessFlags = reader.readSmallUleb128();
 
         this.annotationSetOffset = annotationIterator.seekTo(fieldIndex);
+        initialValueOffset = staticInitialValueIterator.getReaderOffset();
         this.initialValue = staticInitialValueIterator.getNextOrNull();
     }
 
@@ -82,11 +88,13 @@ public class DexBackedField extends BaseFieldReference implements Field {
 
         // large values may be used for the index delta, which cause the cumulative index to overflow upon
         // addition, effectively allowing out of order entries.
+        startOffset = reader.getOffset();
         int fieldIndexDiff = reader.readLargeUleb128();
         this.fieldIndex = fieldIndexDiff + previousFieldIndex;
         this.accessFlags = reader.readSmallUleb128();
 
         this.annotationSetOffset = annotationIterator.seekTo(fieldIndex);
+        initialValueOffset = 0;
         this.initialValue = null;
     }
 
@@ -131,4 +139,38 @@ public class DexBackedField extends BaseFieldReference implements Field {
         }
         return fieldIdItemOffset;
     }
+
+    /**
+     * Calculate and return the private size of a field definition.
+     *
+     * Calculated as: field_idx_diff + access_flags + annotations overhead +
+     * initial value size + field reference size
+     *
+     * @return size in bytes
+     */
+    public int getSize() {
+        int size = 0;
+        DexReader reader = dexFile.readerAt(startOffset);
+        reader.readLargeUleb128(); //field_idx_diff
+        reader.readSmallUleb128(); //access_flags
+        size += reader.getOffset() - startOffset;
+
+        Set<? extends DexBackedAnnotation> annotations = getAnnotations();
+        if (!annotations.isEmpty()) {
+            size += 2 * 4; //2 * uint overhead from field_annotation
+        }
+
+        if (initialValueOffset > 0) {
+            reader.setOffset(initialValueOffset);
+            if (initialValue != null) {
+                DexBackedEncodedValue.skipFrom(reader);
+                size += reader.getOffset() - initialValueOffset;
+            }
+        }
+
+        DexBackedFieldReference fieldRef = new DexBackedFieldReference(dexFile, fieldIndex);
+        size += fieldRef.getSize();
+
+        return size;
+    }
 }

dexlib2/src/main/java/org/jf/dexlib2/dexbacked/DexBackedMethod.java

@@ -37,11 +37,13 @@ import org.jf.dexlib2.base.reference.BaseMethodReference;
 import org.jf.dexlib2.dexbacked.raw.MethodIdItem;
 import org.jf.dexlib2.dexbacked.raw.ProtoIdItem;
 import org.jf.dexlib2.dexbacked.raw.TypeListItem;
+import org.jf.dexlib2.dexbacked.reference.DexBackedMethodReference;
 import org.jf.dexlib2.dexbacked.util.AnnotationsDirectory;
 import org.jf.dexlib2.dexbacked.util.FixedSizeList;
 import org.jf.dexlib2.dexbacked.util.ParameterIterator;
 import org.jf.dexlib2.iface.Annotation;
 import org.jf.dexlib2.iface.Method;
+import org.jf.dexlib2.iface.MethodImplementation;
 import org.jf.dexlib2.iface.MethodParameter;
 import org.jf.util.AbstractForwardSequentialList;
 
@@ -62,6 +64,7 @@ public class DexBackedMethod extends BaseMethodReference implements Method {
     private final int methodAnnotationSetOffset;
 
     public final int methodIndex;
+    private final int startOffset;
 
     private int methodIdItemOffset;
     private int protoIdItemOffset;
@@ -72,6 +75,7 @@ public class DexBackedMethod extends BaseMethodReference implements Method {
                            int previousMethodIndex) {
         this.dexFile = reader.dexBuf;
         this.classDef = classDef;
+        startOffset = reader.getOffset();
 
         // large values may be used for the index delta, which cause the cumulative index to overflow upon
         // addition, effectively allowing out of order entries.
@@ -91,6 +95,7 @@ public class DexBackedMethod extends BaseMethodReference implements Method {
                            @Nonnull AnnotationsDirectory.AnnotationIterator paramaterAnnotationIterator) {
         this.dexFile = reader.dexBuf;
         this.classDef = classDef;
+        startOffset = reader.getOffset();
 
         // large values may be used for the index delta, which cause the cumulative index to overflow upon
         // addition, effectively allowing out of order entries.
@@ -224,4 +229,32 @@ public class DexBackedMethod extends BaseMethodReference implements Method {
             reader.skipUleb128();
         }
     }
+
+    /**
+     * Calculate and return the private size of a method definition.
+     *
+     * Calculated as: method_idx_diff + access_flags + code_off +
+     * implementation size + reference size
+     *
+     * @return size in bytes
+     */
+    public int getSize() {
+        int size = 0;
+
+        DexReader reader = dexFile.readerAt(startOffset);
+        reader.readLargeUleb128(); //method_idx_diff
+        reader.readSmallUleb128(); //access_flags
+        reader.readSmallUleb128(); //code_off
+        size += reader.getOffset() - startOffset;
+
+        DexBackedMethodImplementation impl = getImplementation();
+        if (impl != null) {
+            size += impl.getSize();
+        }
+
+        DexBackedMethodReference methodRef = new DexBackedMethodReference(dexFile, methodIndex);
+        size += methodRef.getSize();
+
+        return size;
+    }
 }

dexlib2/src/main/java/org/jf/dexlib2/dexbacked/DexBackedMethodImplementation.java

@@ -36,6 +36,7 @@ import org.jf.dexlib2.dexbacked.instruction.DexBackedInstruction;
 import org.jf.dexlib2.dexbacked.raw.CodeItem;
 import org.jf.dexlib2.dexbacked.util.DebugInfo;
 import org.jf.dexlib2.dexbacked.util.FixedSizeList;
+import org.jf.dexlib2.dexbacked.util.VariableSizeListIterator;
 import org.jf.dexlib2.dexbacked.util.VariableSizeLookaheadIterator;
 import org.jf.dexlib2.iface.MethodImplementation;
 import org.jf.dexlib2.iface.debug.DebugItem;
@@ -148,4 +149,31 @@ public class DexBackedMethodImplementation implements MethodImplementation {
     public Iterator<String> getParameterNames(@Nullable DexReader dexReader) {
         return getDebugInfo().getParameterNames(dexReader);
     }
+
+    /**
+     * Calculate and return the private size of a method implementation.
+     *
+     * Calculated as: debug info size + instructions size + try-catch size
+     *
+     * @return size in bytes
+     */
+    public int getSize() {
+        int debugSize = getDebugInfo().getSize();
+
+        //set code_item ending offset to the end of instructions list (insns_size * ushort)
+        int lastOffset = dexFile.readSmallUint(codeOffset + CodeItem.INSTRUCTION_COUNT_OFFSET) * 2;
+
+        //read any exception handlers and move code_item offset to the end
+        for (DexBackedTryBlock tryBlock: getTryBlocks()) {
+            Iterator<? extends DexBackedExceptionHandler> tryHandlerIter =
+                tryBlock.getExceptionHandlers().iterator();
+            while (tryHandlerIter.hasNext()) {
+                tryHandlerIter.next();
+            }
+            lastOffset = ((VariableSizeListIterator)tryHandlerIter).getReaderOffset();
+        }
+
+        //method impl size = debug block size + code_item size
+        return debugSize + (lastOffset - codeOffset);
+    }
 }

dexlib2/src/main/java/org/jf/dexlib2/dexbacked/reference/DexBackedFieldReference.java

@@ -63,4 +63,15 @@ public class DexBackedFieldReference extends BaseFieldReference {
     public String getType() {
         return dexFile.getType(dexFile.readUshort(fieldIdItemOffset + FieldIdItem.TYPE_OFFSET));
     }
+
+    /**
+     * Calculate and return the private size of a field reference.
+     *
+     * Calculated as: class_idx + type_idx + name_idx
+     *
+     * @return size in bytes
+     */
+    public int getSize() {
+        return FieldIdItem.ITEM_SIZE;
+    }
 }

dexlib2/src/main/java/org/jf/dexlib2/dexbacked/reference/DexBackedMethodProtoReference.java

@@ -74,4 +74,20 @@ public class DexBackedMethodProtoReference extends BaseMethodProtoReference {
     public String getReturnType() {
         return dexFile.getType(dexFile.readSmallUint(protoIdItemOffset + ProtoIdItem.RETURN_TYPE_OFFSET));
     }
+
+    /**
+     * Calculate and return the private size of a method proto.
+     *
+     * Calculated as: shorty_idx + return_type_idx + parameters_off + type_list size
+     *
+     * @return size in bytes
+     */
+    public int getSize() {
+        int size = ProtoIdItem.ITEM_SIZE; //3 * uint
+        List<String> parameters = getParameterTypes();
+        if (!parameters.isEmpty()) {
+            size += 4 + parameters.size() * 2; //uint + size * ushort for type_idxs
+        }
+        return size;
+    }
 }

dexlib2/src/main/java/org/jf/dexlib2/dexbacked/reference/DexBackedMethodReference.java

@@ -98,4 +98,19 @@ public class DexBackedMethodReference extends BaseMethodReference {
         }
         return protoIdItemOffset;
     }
+
+    /**
+     * Calculate and return the private size of a method reference.
+     *
+     * Calculated as: class_idx + proto_idx + name_idx + prototype size
+     *
+     * @return size in bytes
+     */
+    public int getSize() {
+        int size = MethodIdItem.ITEM_SIZE; //ushort + ushort + uint for indices
+        DexBackedMethodProtoReference protoRef = new DexBackedMethodProtoReference(dexFile,
+            dexFile.readUshort(methodIdItemOffset + MethodIdItem.PROTO_OFFSET));
+        size += protoRef.getSize();
+        return size;
+    }
 }

dexlib2/src/main/java/org/jf/dexlib2/dexbacked/reference/DexBackedStringReference.java

@@ -33,6 +33,8 @@ package org.jf.dexlib2.dexbacked.reference;
 
 import org.jf.dexlib2.base.reference.BaseStringReference;
 import org.jf.dexlib2.dexbacked.DexBackedDexFile;
+import org.jf.dexlib2.dexbacked.DexReader;
+import org.jf.dexlib2.dexbacked.raw.StringIdItem;
 
 import javax.annotation.Nonnull;
 
@@ -50,4 +52,25 @@ public class DexBackedStringReference extends BaseStringReference {
     public String getString() {
         return dexFile.getString(stringIndex);
     }
+
+
+    /**
+     * Calculate and return the private size of a string reference.
+     *
+     * Calculated as: string_data_off + string_data_item size
+     *
+     * @return size in bytes
+     */
+    public int getSize() {
+        int size = StringIdItem.ITEM_SIZE; //uint for string_data_off
+        //add the string data length:
+        int stringOffset = dexFile.getStringIdItemOffset(stringIndex);
+        int stringDataOffset = dexFile.readSmallUint(stringOffset);
+        DexReader reader = dexFile.readerAt(stringDataOffset);
+        size += reader.peekSmallUleb128Size();
+        int utf16Length = reader.readSmallUleb128();
+        //and string data itself:
+        size += reader.peekStringLength(utf16Length);
+        return size;
+    }
 }

dexlib2/src/main/java/org/jf/dexlib2/dexbacked/reference/DexBackedTypeReference.java

@@ -33,6 +33,7 @@ package org.jf.dexlib2.dexbacked.reference;
 
 import org.jf.dexlib2.base.reference.BaseTypeReference;
 import org.jf.dexlib2.dexbacked.DexBackedDexFile;
+import org.jf.dexlib2.dexbacked.raw.TypeIdItem;
 
 import javax.annotation.Nonnull;
 
@@ -49,4 +50,16 @@ public class DexBackedTypeReference extends BaseTypeReference {
     @Nonnull public String getType() {
         return dexFile.getType(typeIndex);
     }
+
+
+    /**
+     * Calculate and return the private size of a type reference.
+     *
+     * Calculated as: descriptor_idx
+     *
+     * @return size in bytes
+     */
+    public int getSize() {
+        return TypeIdItem.ITEM_SIZE; //uint for descriptor_idx
+    }
 }

dexlib2/src/main/java/org/jf/dexlib2/dexbacked/util/DebugInfo.java

@@ -59,6 +59,13 @@ public abstract class DebugInfo implements Iterable<DebugItem> {
      */
     @Nonnull public abstract Iterator<String> getParameterNames(@Nullable DexReader reader);
 
+    /**
+     * Calculate and return the private size of debuginfo.
+     *
+     * @return size in bytes
+     */
+    public abstract int getSize();
+
     public static DebugInfo newOrEmpty(@Nonnull DexBackedDexFile dexFile, int debugInfoOffset,
                                        @Nonnull DexBackedMethodImplementation methodImpl) {
         if (debugInfoOffset == 0) {
@@ -78,6 +85,11 @@ public abstract class DebugInfo implements Iterable<DebugItem> {
         @Nonnull @Override public Iterator<String> getParameterNames(@Nullable DexReader reader) {
             return ImmutableSet.<String>of().iterator();
         }
+
+        @Override
+        public int getSize() {
+            return 0;
+        }
     }
 
     private static class DebugInfoImpl extends DebugInfo {
@@ -279,5 +291,14 @@ public abstract class DebugInfo implements Iterable<DebugItem> {
                 }
             };
         }
+
+        @Override
+        public int getSize() {
+            Iterator<DebugItem> iter = iterator();
+            while(iter.hasNext()) {
+                iter.next();
+            }
+            return ((VariableSizeLookaheadIterator) iter).getReaderOffset() - debugInfoOffset;
+        }
     }
 }

dexlib2/src/main/java/org/jf/dexlib2/dexbacked/util/StaticInitialValueIterator.java

@@ -43,10 +43,12 @@ public abstract class StaticInitialValueIterator {
     public static final StaticInitialValueIterator EMPTY = new StaticInitialValueIterator() {
         @Nullable @Override public EncodedValue getNextOrNull() { return null; }
         @Override public void skipNext() {}
+        @Override public int getReaderOffset() { return 0; }
     };
 
     @Nullable public abstract EncodedValue getNextOrNull();
     public abstract void skipNext();
+    public abstract int getReaderOffset();
 
     @Nonnull
     public static StaticInitialValueIterator newOrEmpty(@Nonnull DexBackedDexFile dexFile, int offset) {
@@ -81,5 +83,9 @@ public abstract class StaticInitialValueIterator {
                 DexBackedEncodedValue.skipFrom(reader);
             }
         }
+
+        public int getReaderOffset() {
+            return reader.getOffset();
+        }
     }
 }

dexlib2/src/main/java/org/jf/dexlib2/dexbacked/util/VariableSizeLookaheadIterator.java

@@ -59,4 +59,8 @@ public abstract class VariableSizeLookaheadIterator<T> extends AbstractIterator<
     protected T computeNext() {
         return readNextItem(reader);
     }
+
+    public final int getReaderOffset() {
+        return reader.getOffset();
+    }
 }

dexlib2/src/test/java/org/jf/dexlib2/dexbacked/BaseDexReaderLeb128Test.java

@@ -254,11 +254,21 @@ public class BaseDexReaderLeb128Test {
         reader = dexBuf.readerAt(0);
         reader.skipUleb128();
         Assert.assertEquals(expectedLength, reader.getOffset());
+
+        reader = dexBuf.readerAt(0);
+        Assert.assertEquals(expectedLength, reader.peekSmallUleb128Size());
     }
 
     private void performFailureTest(byte[] buf) {
         BaseDexBuffer dexBuf = new BaseDexBuffer(buf);
         BaseDexReader reader = dexBuf.readerAt(0);
+        try {
+            reader.peekSmallUleb128Size();
+            Assert.fail();
+        } catch (ExceptionWithContext ex) {
+            // expected exception
+        }
+
         try {
             reader.readSmallUleb128();
             Assert.fail();

dexlib2/src/test/java/org/jf/dexlib2/dexbacked/BaseDexReaderSleb128Test.java

@@ -257,11 +257,20 @@ public class BaseDexReaderSleb128Test {
         BaseDexReader reader = dexBuf.readerAt(0);
         Assert.assertEquals(expectedValue, reader.readSleb128());
         Assert.assertEquals(expectedLength, reader.getOffset());
+
+        reader = dexBuf.readerAt(0);
+        Assert.assertEquals(expectedLength, reader.peekSleb128Size());
     }
 
     private void performFailureTest(byte[] buf) {
         BaseDexBuffer dexBuf = new BaseDexBuffer(buf);
         BaseDexReader reader = dexBuf.readerAt(0);
+        try {
+            reader.peekSleb128Size();
+            Assert.fail();
+        } catch (ExceptionWithContext ex) {
+            // expected exception
+        }
         try {
             reader.readSleb128();
             Assert.fail();