Add support to dexlib2 for writing hidden api restrictions

dexlib2/src/main/java/org/jf/dexlib2/HiddenApiRestriction.java

@@ -107,4 +107,31 @@ public enum HiddenApiRestriction {
         }
         return joiner.toString();
     }
+
+    public static int combineFlags(Iterable<HiddenApiRestriction> flags) {
+        boolean gotHiddenApiFlag = false;
+        boolean gotDomainSpecificApiFlag = false;
+
+        int value = 0;
+
+        for (HiddenApiRestriction flag : flags) {
+            if (flag.isDomainSpecificApiFlag) {
+                if (gotDomainSpecificApiFlag) {
+                    throw new IllegalArgumentException(
+                            "Cannot combine multiple flags for domain-specific api restrictions");
+                }
+                gotDomainSpecificApiFlag = true;
+                value += flag.value;
+            } else {
+                if (gotHiddenApiFlag) {
+                    throw new IllegalArgumentException(
+                            "Cannot combine multiple flags for hidden api restrictions");
+                }
+                gotHiddenApiFlag = true;
+                value += flag.value;
+            }
+        }
+
+        return value;
+    }
 }

dexlib2/src/main/java/org/jf/dexlib2/writer/ClassSection.java

@@ -31,6 +31,7 @@
 
 package org.jf.dexlib2.writer;
 
+import org.jf.dexlib2.HiddenApiRestriction;
 import org.jf.dexlib2.builder.MutableMethodImplementation;
 import org.jf.dexlib2.iface.ExceptionHandler;
 import org.jf.dexlib2.iface.TryBlock;
@@ -43,6 +44,7 @@ import java.io.IOException;
 import java.util.Collection;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
 
 public interface ClassSection<StringKey extends CharSequence, TypeKey extends CharSequence, TypeListKey, ClassKey,
         FieldKey, MethodKey, AnnotationSetKey, EncodedArrayKey> extends IndexSection<ClassKey> {
@@ -67,6 +69,9 @@ public interface ClassSection<StringKey extends CharSequence, TypeKey extends Ch
     int getFieldAccessFlags(@Nonnull FieldKey key);
     int getMethodAccessFlags(@Nonnull MethodKey key);
 
+    @Nonnull Set<HiddenApiRestriction> getFieldHiddenApiRestrictions(@Nonnull FieldKey key);
+    @Nonnull Set<HiddenApiRestriction> getMethodHiddenApiRestrictions(@Nonnull MethodKey key);
+
     @Nullable AnnotationSetKey getClassAnnotations(@Nonnull ClassKey key);
     @Nullable AnnotationSetKey getFieldAnnotations(@Nonnull FieldKey key);
     @Nullable AnnotationSetKey getMethodAnnotations(@Nonnull MethodKey key);

dexlib2/src/main/java/org/jf/dexlib2/writer/DexWriter.java

@@ -130,8 +130,11 @@ public abstract class DexWriter<
     protected int annotationDirectorySectionOffset = NO_OFFSET;
     protected int debugSectionOffset = NO_OFFSET;
     protected int codeSectionOffset = NO_OFFSET;
+    protected int hiddenApiRestrictionsOffset = NO_OFFSET;
     protected int mapSectionOffset = NO_OFFSET;
 
+    protected boolean hasHiddenApiRestrictions = false;
+
     protected int numAnnotationSetRefItems = 0;
     protected int numAnnotationDirectoryItems = 0;
     protected int numDebugInfoItems = 0;
@@ -304,6 +307,7 @@ public abstract class DexWriter<
             DexDataWriter headerWriter = outputAt(dest, 0);
             DexDataWriter indexWriter = outputAt(dest, HeaderItem.ITEM_SIZE);
             DexDataWriter offsetWriter = outputAt(dest, dataSectionOffset);
+
             try {
                 writeStrings(indexWriter, offsetWriter);
                 writeTypes(indexWriter);
@@ -339,7 +343,7 @@ public abstract class DexWriter<
                 writeAnnotationSetRefs(offsetWriter);
                 writeAnnotationDirectories(offsetWriter);
                 writeDebugAndCodeItems(offsetWriter, tempFactory.makeDeferredOutputStream());
-                writeClasses(indexWriter, offsetWriter);
+                writeClasses(dest, indexWriter, offsetWriter);
 
                 writeMapItem(offsetWriter);
                 writeHeader(headerWriter, dataSectionOffset, offsetWriter.getPosition());
@@ -481,7 +485,8 @@ public abstract class DexWriter<
         }
     }
 
-    private void writeClasses(@Nonnull DexDataWriter indexWriter, @Nonnull DexDataWriter offsetWriter) throws IOException {
+    private void writeClasses(@Nonnull DexDataStore dataStore, @Nonnull DexDataWriter indexWriter,
+                              @Nonnull DexDataWriter offsetWriter) throws IOException {
         classIndexSectionOffset = indexWriter.getPosition();
         classDataSectionOffset = offsetWriter.getPosition();
 
@@ -492,6 +497,124 @@ public abstract class DexWriter<
         for (Map.Entry<? extends ClassKey, Integer> key: classEntries) {
             index = writeClass(indexWriter, offsetWriter, index, key);
         }
+
+        if (!shouldWriteHiddenApiRestrictions()) {
+            return;
+        }
+
+        hiddenApiRestrictionsOffset = offsetWriter.getPosition();
+
+        RestrictionsWriter restrictionsWriter = new RestrictionsWriter(dataStore, offsetWriter, classEntries.size());
+
+        try {
+            for (Map.Entry<? extends ClassKey, Integer> key : classEntries) {
+
+                for (FieldKey fieldKey : classSection.getSortedStaticFields(key.getKey())) {
+                    restrictionsWriter.writeRestriction(classSection.getFieldHiddenApiRestrictions(fieldKey));
+                }
+
+                for (FieldKey fieldKey : classSection.getSortedInstanceFields(key.getKey())) {
+                    restrictionsWriter.writeRestriction(classSection.getFieldHiddenApiRestrictions(fieldKey));
+                }
+
+                for (MethodKey methodKey : classSection.getSortedDirectMethods(key.getKey())) {
+                    restrictionsWriter.writeRestriction(classSection.getMethodHiddenApiRestrictions(methodKey));
+                }
+
+                for (MethodKey methodKey : classSection.getSortedVirtualMethods(key.getKey())) {
+                    restrictionsWriter.writeRestriction(classSection.getMethodHiddenApiRestrictions(methodKey));
+                }
+
+                restrictionsWriter.finishClass();
+            }
+        } finally {
+            restrictionsWriter.close();
+        }
+    }
+
+    private boolean shouldWriteHiddenApiRestrictions() {
+        return hasHiddenApiRestrictions && opcodes.api >= 29;
+    }
+
+    private static class RestrictionsWriter {
+        private final int startOffset;
+
+        private final DexDataStore dataStore;
+        private final DexDataWriter offsetsWriter;
+        private final DexDataWriter restrictionsWriter;
+
+        private boolean writeRestrictionsForClass = false;
+        private int pendingBlankEntries = 0;
+
+        public RestrictionsWriter(DexDataStore dataStore, DexDataWriter offsetWriter, int numClasses)
+                throws IOException {
+            this.startOffset = offsetWriter.getPosition();
+            this.dataStore = dataStore;
+            this.restrictionsWriter = offsetWriter;
+
+            int offsetsSize = numClasses * HiddenApiClassDataItem.OFFSET_ITEM_SIZE;
+
+            // We don't know the size yet, so skip over it
+            restrictionsWriter.writeInt(0);
+
+            this.offsetsWriter = outputAt(dataStore, restrictionsWriter.getPosition());
+
+            // Skip over the offsets
+            for (int i=0; i<offsetsSize; i++) {
+                restrictionsWriter.write(0);
+            }
+            restrictionsWriter.flush();
+        }
+
+        public void finishClass() throws IOException {
+            if (!writeRestrictionsForClass) {
+                // Normally the offset gets written when the first non-blank restriction gets written. If only blank
+                // restrictions were added, nothing actually gets written, and we write out a blank offset here.
+                offsetsWriter.writeInt(0);
+            }
+
+            writeRestrictionsForClass = false;
+            pendingBlankEntries = 0;
+        }
+
+        private void addBlankEntry() throws IOException {
+            if (writeRestrictionsForClass) {
+                restrictionsWriter.writeUleb128(HiddenApiRestriction.WHITELIST.getValue());
+            } else {
+                pendingBlankEntries++;
+            }
+        }
+
+        public void writeRestriction(@Nonnull Set<HiddenApiRestriction> hiddenApiRestrictions) throws IOException {
+            if (hiddenApiRestrictions.isEmpty()) {
+                addBlankEntry();
+                return;
+            }
+
+            if (!writeRestrictionsForClass) {
+                writeRestrictionsForClass = true;
+                offsetsWriter.writeInt(restrictionsWriter.getPosition() - startOffset);
+
+                for (int i = 0; i < pendingBlankEntries; i++) {
+                    restrictionsWriter.writeUleb128(HiddenApiRestriction.WHITELIST.getValue());
+                }
+                pendingBlankEntries = 0;
+            }
+            restrictionsWriter.writeUleb128(HiddenApiRestriction.combineFlags(hiddenApiRestrictions));
+        }
+
+        public void close() throws IOException {
+            DexDataWriter writer = null;
+            offsetsWriter.close();
+            try {
+                writer = outputAt(dataStore, startOffset);
+                writer.writeInt(restrictionsWriter.getPosition() - startOffset);
+            } finally {
+                if (writer != null) {
+                    writer.close();
+                }
+            }
+        }
     }
 
     /**
@@ -639,6 +762,9 @@ public abstract class DexWriter<
         int prevIndex = 0;
         for (FieldKey key: fields) {
             int index = fieldSection.getFieldIndex(key);
+            if (!classSection.getFieldHiddenApiRestrictions(key).isEmpty()) {
+                hasHiddenApiRestrictions = true;
+            }
             writer.writeUleb128(index - prevIndex);
             writer.writeUleb128(classSection.getFieldAccessFlags(key));
             prevIndex = index;
@@ -650,6 +776,9 @@ public abstract class DexWriter<
         int prevIndex = 0;
         for (MethodKey key: methods) {
             int index = methodSection.getMethodIndex(key);
+            if (!classSection.getMethodHiddenApiRestrictions(key).isEmpty()) {
+                hasHiddenApiRestrictions = true;
+            }
             writer.writeUleb128(index-prevIndex);
             writer.writeUleb128(classSection.getMethodAccessFlags(key));
             writer.writeUleb128(classSection.getCodeItemOffset(key));
@@ -1324,6 +1453,9 @@ public abstract class DexWriter<
         if (numClassDataItems > 0) {
             numItems++;
         }
+        if (shouldWriteHiddenApiRestrictions()) {
+            numItems++;
+        }
         // map item itself
         numItems++;
 
@@ -1364,6 +1496,11 @@ public abstract class DexWriter<
         writeMapItem(writer, ItemType.DEBUG_INFO_ITEM, numDebugInfoItems, debugSectionOffset);
         writeMapItem(writer, ItemType.CODE_ITEM, numCodeItemItems, codeSectionOffset);
         writeMapItem(writer, ItemType.CLASS_DATA_ITEM, numClassDataItems, classDataSectionOffset);
+
+        if (shouldWriteHiddenApiRestrictions()) {
+            writeMapItem(writer, ItemType.HIDDENAPI_CLASS_DATA_ITEM, 1, hiddenApiRestrictionsOffset);
+        }
+
         writeMapItem(writer, ItemType.MAP_LIST, 1, mapSectionOffset);
     }
 

dexlib2/src/main/java/org/jf/dexlib2/writer/builder/BuilderClassPool.java

@@ -35,6 +35,7 @@ import com.google.common.base.Function;
 import com.google.common.base.Predicate;
 import com.google.common.collect.*;
 import org.jf.dexlib2.DebugItemType;
+import org.jf.dexlib2.HiddenApiRestriction;
 import org.jf.dexlib2.builder.MutableMethodImplementation;
 import org.jf.dexlib2.iface.ExceptionHandler;
 import org.jf.dexlib2.iface.Field;
@@ -56,10 +57,7 @@ import org.jf.util.ExceptionWithContext;
 import javax.annotation.Nonnull;
 import javax.annotation.Nullable;
 import java.io.IOException;
-import java.util.Collection;
+import java.util.*;
-import java.util.Iterator;
-import java.util.List;
-import java.util.Map;
 import java.util.Map.Entry;
 import java.util.concurrent.ConcurrentMap;
 
@@ -198,6 +196,16 @@ public class BuilderClassPool extends BaseBuilderPool implements ClassSection<Bu
         return builderMethod.accessFlags;
     }
 
+    @Nonnull @Override
+    public Set<HiddenApiRestriction> getFieldHiddenApiRestrictions(@Nonnull BuilderField builderField) {
+        return builderField.getHiddenApiRestrictions();
+    }
+
+    @Nonnull @Override
+    public Set<HiddenApiRestriction> getMethodHiddenApiRestrictions(@Nonnull BuilderMethod builderMethod) {
+        return builderMethod.getHiddenApiRestrictions();
+    }
+
     @Nullable @Override public BuilderAnnotationSet getClassAnnotations(@Nonnull BuilderClassDef builderClassDef) {
         if (builderClassDef.annotations.isEmpty()) {
             return null;

dexlib2/src/main/java/org/jf/dexlib2/writer/builder/BuilderField.java

@@ -46,15 +46,18 @@ public class BuilderField extends BaseFieldReference implements Field {
     final int accessFlags;
     @Nullable final BuilderEncodedValue initialValue;
     @Nonnull final BuilderAnnotationSet annotations;
+    @Nonnull Set<HiddenApiRestriction> hiddenApiRestrictions;
 
     BuilderField(@Nonnull BuilderFieldReference fieldReference,
                  int accessFlags,
                  @Nullable BuilderEncodedValue initialValue,
-                 @Nonnull BuilderAnnotationSet annotations) {
+                 @Nonnull BuilderAnnotationSet annotations,
+                 @Nonnull Set<HiddenApiRestriction> hiddenApiRestrictions) {
         this.fieldReference = fieldReference;
         this.accessFlags = accessFlags;
         this.initialValue = initialValue;
         this.annotations = annotations;
+        this.hiddenApiRestrictions = hiddenApiRestrictions;
     }
 
     @Override public int getAccessFlags() {
@@ -82,6 +85,6 @@ public class BuilderField extends BaseFieldReference implements Field {
     }
 
     @Nonnull @Override public Set<HiddenApiRestriction> getHiddenApiRestrictions() {
-        return ImmutableSet.of();
+        return hiddenApiRestrictions;
     }
 }

dexlib2/src/main/java/org/jf/dexlib2/writer/builder/BuilderMethod.java

@@ -31,7 +31,6 @@
 
 package org.jf.dexlib2.writer.builder;
 
-import com.google.common.collect.ImmutableSet;
 import org.jf.dexlib2.HiddenApiRestriction;
 import org.jf.dexlib2.base.reference.BaseMethodReference;
 import org.jf.dexlib2.iface.Method;
@@ -48,6 +47,7 @@ public class BuilderMethod extends BaseMethodReference implements Method {
     @Nonnull final List<? extends BuilderMethodParameter> parameters;
     final int accessFlags;
     @Nonnull final BuilderAnnotationSet annotations;
+    @Nonnull final Set<HiddenApiRestriction> hiddenApiRestrictions;
     @Nullable final MethodImplementation methodImplementation;
 
     int annotationSetRefListOffset = DexWriter.NO_OFFSET;
@@ -57,11 +57,13 @@ public class BuilderMethod extends BaseMethodReference implements Method {
                   @Nonnull List<? extends BuilderMethodParameter> parameters,
                   int accessFlags,
                   @Nonnull BuilderAnnotationSet annotations,
+                  @Nonnull Set<HiddenApiRestriction> hiddenApiRestrictions,
                   @Nullable MethodImplementation methodImplementation) {
         this.methodReference = methodReference;
         this.parameters = parameters;
         this.accessFlags = accessFlags;
         this.annotations = annotations;
+        this.hiddenApiRestrictions = hiddenApiRestrictions;
         this.methodImplementation = methodImplementation;
     }
 
@@ -72,10 +74,6 @@ public class BuilderMethod extends BaseMethodReference implements Method {
     @Override @Nonnull public List<? extends BuilderMethodParameter> getParameters() { return parameters; }
     @Override public int getAccessFlags() { return accessFlags; }
     @Override @Nonnull public BuilderAnnotationSet getAnnotations() { return annotations; }
-
+    @Nonnull @Override public Set<HiddenApiRestriction> getHiddenApiRestrictions() { return hiddenApiRestrictions; }
-    @Nonnull @Override public Set<HiddenApiRestriction> getHiddenApiRestrictions() {
-        return ImmutableSet.of();
-    }
-
     @Override @Nullable public MethodImplementation getImplementation() { return methodImplementation; }
 }

dexlib2/src/main/java/org/jf/dexlib2/writer/builder/DexBuilder.java

@@ -33,6 +33,7 @@ package org.jf.dexlib2.writer.builder;
 
 import com.google.common.base.Function;
 import com.google.common.collect.*;
+import org.jf.dexlib2.HiddenApiRestriction;
 import org.jf.dexlib2.Opcodes;
 import org.jf.dexlib2.ValueType;
 import org.jf.dexlib2.iface.Annotation;
@@ -75,11 +76,13 @@ public class DexBuilder extends DexWriter<BuilderStringReference, BuilderStringR
                                              @Nonnull String type,
                                              int accessFlags,
                                              @Nullable EncodedValue initialValue,
-                                             @Nonnull Set<? extends Annotation> annotations) {
+                                             @Nonnull Set<? extends Annotation> annotations,
+                                             @Nonnull Set<HiddenApiRestriction> hiddenApiRestrictions) {
         return new BuilderField(fieldSection.internField(definingClass, name, type),
                 accessFlags,
                 internNullableEncodedValue(initialValue),
-                annotationSetSection.internAnnotationSet(annotations));
+                annotationSetSection.internAnnotationSet(annotations),
+                hiddenApiRestrictions);
     }
 
     @Nonnull public BuilderMethod internMethod(@Nonnull String definingClass,
@@ -88,6 +91,7 @@ public class DexBuilder extends DexWriter<BuilderStringReference, BuilderStringR
                                                @Nonnull String returnType,
                                                int accessFlags,
                                                @Nonnull Set<? extends Annotation> annotations,
+                                               @Nonnull Set<HiddenApiRestriction> hiddenApiRestrictions,
                                                @Nullable MethodImplementation methodImplementation) {
         if (parameters == null) {
             parameters = ImmutableList.of();
@@ -96,6 +100,7 @@ public class DexBuilder extends DexWriter<BuilderStringReference, BuilderStringR
                 internMethodParameters(parameters),
                 accessFlags,
                 annotationSetSection.internAnnotationSet(annotations),
+                hiddenApiRestrictions,
                 methodImplementation);
     }
 

dexlib2/src/main/java/org/jf/dexlib2/writer/pool/ClassPool.java

@@ -38,6 +38,7 @@ import com.google.common.collect.ImmutableList;
 import com.google.common.collect.Iterables;
 import com.google.common.collect.Ordering;
 import org.jf.dexlib2.DebugItemType;
+import org.jf.dexlib2.HiddenApiRestriction;
 import org.jf.dexlib2.ReferenceType;
 import org.jf.dexlib2.builder.MutableMethodImplementation;
 import org.jf.dexlib2.iface.*;
@@ -287,6 +288,14 @@ public class ClassPool extends BasePool<String, PoolClassDef> implements ClassSe
         return method.getAccessFlags();
     }
 
+    @Nonnull @Override public Set<HiddenApiRestriction> getFieldHiddenApiRestrictions(@Nonnull Field field) {
+        return field.getHiddenApiRestrictions();
+    }
+
+    @Nonnull @Override public Set<HiddenApiRestriction> getMethodHiddenApiRestrictions(@Nonnull PoolMethod poolMethod) {
+        return poolMethod.getHiddenApiRestrictions();
+    }
+
     @Nullable @Override public Set<? extends Annotation> getClassAnnotations(@Nonnull PoolClassDef classDef) {
         Set<? extends Annotation> annotations = classDef.getAnnotations();
         if (annotations.size() == 0) {

dexlib2/src/test/java/org/jf/dexlib2/writer/CallSiteTest.java

@@ -109,7 +109,7 @@ public class CallSiteTest {
                 callSite));
 
         BuilderMethod method = dexBuilder.internMethod("Lcls1", "method1", null, "V", 0, ImmutableSet.of(),
-                methodImplementationBuilder.getMethodImplementation());
+                ImmutableSet.of(), methodImplementationBuilder.getMethodImplementation());
         dexBuilder.internClassDef("Lcls1;", AccessFlags.PUBLIC.getValue(), "Ljava/lang/Object;", null, null,
                 ImmutableSet.of(), null,
                 ImmutableList.of(method));

dexlib2/src/test/java/org/jf/dexlib2/writer/JumboStringConversionTest.java

@@ -87,6 +87,7 @@ public class JumboStringConversionTest {
                                 "V",
                                 0,
                                 ImmutableSet.<Annotation>of(),
+                                ImmutableSet.of(),
                                 methodBuilder.getMethodImplementation())));
 
         MemoryDataStore dexStore = new MemoryDataStore();
@@ -184,6 +185,7 @@ public class JumboStringConversionTest {
                                 "V",
                                 0,
                                 ImmutableSet.<Annotation>of(),
+                                ImmutableSet.of(),
                                 methodImpl)));
 
         MemoryDataStore dexStore = new MemoryDataStore();

smali/src/main/antlr/smaliTreeWalker.g

@@ -264,7 +264,7 @@ field returns [BuilderField field]
     }
 
     $field = dexBuilder.internField(classType, $SIMPLE_NAME.text, $nonvoid_type_descriptor.type, $access_list.value,
-            $field_initial_value.encodedValue, $annotations.annotations);
+            $field_initial_value.encodedValue, $annotations.annotations, ImmutableSet.of());
   };
 
 
@@ -467,6 +467,7 @@ method returns[BuilderMethod ret]
             $method_name_and_prototype.returnType,
             accessFlags,
             $annotations.annotations,
+            ImmutableSet.of(),
             methodImplementation);
   };