Merge "Initialize BitInputStream with rule input stream"

4 files changed, 53 insertions, 30 deletions

diff --git a/services/core/java/com/android/server/integrity/model/BitInputStream.java b/services/core/java/com/android/server/integrity/model/BitInputStream.java
index 09bc7e8b9861..e768fe6626ac 100644
--- a/services/core/java/com/android/server/integrity/model/BitInputStream.java
+++ b/services/core/java/com/android/server/integrity/model/BitInputStream.java
@@ -16,15 +16,29 @@
 
 package com.android.server.integrity.model;
 
+import java.io.IOException;
+import java.io.InputStream;
+
 /** A wrapper class for reading a stream of bits. */
 public class BitInputStream {
 
-    private byte[] mRuleBytes;
     private long mBitPointer;
+    private boolean mReadFromStream;
+
+    private byte[] mRuleBytes;
+    private InputStream mRuleInputStream;
+
+    private byte mCurrentRuleByte;
 
     public BitInputStream(byte[] ruleBytes) {
         this.mRuleBytes = ruleBytes;
         this.mBitPointer = 0;
+        this.mReadFromStream = false;
+    }
+
+    public BitInputStream(InputStream ruleInputStream) {
+        this.mRuleInputStream = ruleInputStream;
+        this.mReadFromStream = true;
     }
 
     /**
@@ -33,34 +47,43 @@ public class BitInputStream {
      * @param numOfBits The number of bits to read.
      * @return The value read from the stream.
      */
-    public int getNext(int numOfBits) {
+    public int getNext(int numOfBits) throws IOException {
         int component = 0;
         int count = 0;
 
-        int idx = (int) (mBitPointer / 8);
-        int offset = 7 - (int) (mBitPointer % 8);
-
         while (count++ < numOfBits) {
-            if (idx >= mRuleBytes.length) {
-                throw new IllegalArgumentException(String.format("Invalid byte index: %d", idx));
+            if (mBitPointer % 8 == 0) {
+                mCurrentRuleByte = getNextByte();
             }
+            int offset = 7 - (int) (mBitPointer % 8);
 
             component <<= 1;
-            component |= (mRuleBytes[idx] >>> offset) & 1;
+            component |= (mCurrentRuleByte >>> offset) & 1;
 
-            offset--;
-            if (offset == -1) {
-                idx++;
-                offset = 7;
-            }
+            mBitPointer++;
         }
 
-        mBitPointer += numOfBits;
         return component;
     }
 
     /** Check if there are bits left in the stream. */
-    public boolean hasNext() {
-        return mBitPointer / 8 < mRuleBytes.length;
+    public boolean hasNext() throws IOException {
+        if (mReadFromStream) {
+            return mRuleInputStream.available() > 0;
+        } else {
+            return mBitPointer / 8 < mRuleBytes.length;
+        }
+    }
+
+    private byte getNextByte() throws IOException {
+        if (mReadFromStream) {
+            return (byte) mRuleInputStream.read();
+        } else {
+            int idx = (int) (mBitPointer / 8);
+            if (idx >= mRuleBytes.length) {
+                throw new IllegalArgumentException(String.format("Invalid byte index: %d", idx));
+            }
+            return mRuleBytes[idx];
+        }
     }
 }
diff --git a/services/core/java/com/android/server/integrity/parser/RuleBinaryParser.java b/services/core/java/com/android/server/integrity/parser/RuleBinaryParser.java
index 8aa0751af11c..3ef45a637bc1 100644
--- a/services/core/java/com/android/server/integrity/parser/RuleBinaryParser.java
+++ b/services/core/java/com/android/server/integrity/parser/RuleBinaryParser.java
@@ -36,6 +36,7 @@ import android.content.integrity.Rule;
 
 import com.android.server.integrity.model.BitInputStream;
 
+import java.io.IOException;
 import java.io.InputStream;
 import java.util.ArrayList;
 import java.util.List;
@@ -56,15 +57,14 @@ public class RuleBinaryParser implements RuleParser {
     @Override
     public List<Rule> parse(InputStream inputStream) throws RuleParseException {
         try {
-            byte[] ruleBytes = new byte[inputStream.available()];
-            inputStream.read(ruleBytes);
-            return parse(ruleBytes);
+            BitInputStream bitInputStream = new BitInputStream(inputStream);
+            return parseRules(bitInputStream);
         } catch (Exception e) {
             throw new RuleParseException(e.getMessage(), e);
         }
     }
 
-    private List<Rule> parseRules(BitInputStream bitInputStream) {
+    private List<Rule> parseRules(BitInputStream bitInputStream) throws IOException {
         List<Rule> parsedRules = new ArrayList<>();
 
         // Read the rule binary file format version.
@@ -79,7 +79,7 @@ public class RuleBinaryParser implements RuleParser {
         return parsedRules;
     }
 
-    private Rule parseRule(BitInputStream bitInputStream) {
+    private Rule parseRule(BitInputStream bitInputStream) throws IOException {
         Formula formula = parseFormula(bitInputStream);
         int effect = bitInputStream.getNext(EFFECT_BITS);
 
@@ -90,7 +90,7 @@ public class RuleBinaryParser implements RuleParser {
         return new Rule(formula, effect);
     }
 
-    private Formula parseFormula(BitInputStream bitInputStream) {
+    private Formula parseFormula(BitInputStream bitInputStream) throws IOException {
         int separator = bitInputStream.getNext(SEPARATOR_BITS);
         switch (separator) {
             case ATOMIC_FORMULA_START:
@@ -105,7 +105,7 @@ public class RuleBinaryParser implements RuleParser {
         }
     }
 
-    private CompoundFormula parseCompoundFormula(BitInputStream bitInputStream) {
+    private CompoundFormula parseCompoundFormula(BitInputStream bitInputStream) throws IOException {
         int connector = bitInputStream.getNext(CONNECTOR_BITS);
         List<Formula> formulas = new ArrayList<>();
 
@@ -118,7 +118,7 @@ public class RuleBinaryParser implements RuleParser {
         return new CompoundFormula(connector, formulas);
     }
 
-    private AtomicFormula parseAtomicFormula(BitInputStream bitInputStream) {
+    private AtomicFormula parseAtomicFormula(BitInputStream bitInputStream) throws IOException {
         int key = bitInputStream.getNext(KEY_BITS);
         int operator = bitInputStream.getNext(OPERATOR_BITS);
 
diff --git a/services/core/java/com/android/server/integrity/serializer/RuleBinarySerializer.java b/services/core/java/com/android/server/integrity/serializer/RuleBinarySerializer.java
index b988fd4c40f1..fdbb7d9df293 100644
--- a/services/core/java/com/android/server/integrity/serializer/RuleBinarySerializer.java
+++ b/services/core/java/com/android/server/integrity/serializer/RuleBinarySerializer.java
@@ -127,7 +127,7 @@ public class RuleBinarySerializer implements RuleSerializer {
 
         bitOutputStream.setNext(SEPARATOR_BITS, ATOMIC_FORMULA_START);
         bitOutputStream.setNext(KEY_BITS, atomicFormula.getKey());
-        if (atomicFormula instanceof AtomicFormula.StringAtomicFormula) {
+        if (atomicFormula.getTag() == AtomicFormula.STRING_ATOMIC_FORMULA_TAG) {
             AtomicFormula.StringAtomicFormula stringAtomicFormula =
                     (AtomicFormula.StringAtomicFormula) atomicFormula;
             bitOutputStream.setNext(OPERATOR_BITS, AtomicFormula.EQ);
@@ -135,7 +135,7 @@ public class RuleBinarySerializer implements RuleSerializer {
                     stringAtomicFormula.getValue(),
                     stringAtomicFormula.getIsHashedValue(),
                     bitOutputStream);
-        } else if (atomicFormula instanceof AtomicFormula.IntAtomicFormula) {
+        } else if (atomicFormula.getTag() == AtomicFormula.INT_ATOMIC_FORMULA_TAG) {
             AtomicFormula.IntAtomicFormula intAtomicFormula =
                     (AtomicFormula.IntAtomicFormula) atomicFormula;
             bitOutputStream.setNext(OPERATOR_BITS, intAtomicFormula.getOperator());
@@ -143,7 +143,7 @@ public class RuleBinarySerializer implements RuleSerializer {
                     String.valueOf(intAtomicFormula.getValue()),
                     /* isHashedValue= */ false,
                     bitOutputStream);
-        } else if (atomicFormula instanceof AtomicFormula.BooleanAtomicFormula) {
+        } else if (atomicFormula.getTag() == AtomicFormula.BOOLEAN_ATOMIC_FORMULA_TAG) {
             AtomicFormula.BooleanAtomicFormula booleanAtomicFormula =
                     (AtomicFormula.BooleanAtomicFormula) atomicFormula;
             bitOutputStream.setNext(OPERATOR_BITS, AtomicFormula.EQ);
diff --git a/services/core/java/com/android/server/integrity/serializer/RuleXmlSerializer.java b/services/core/java/com/android/server/integrity/serializer/RuleXmlSerializer.java
index 72068ceeb4f0..cfe50c6c8ac9 100644
--- a/services/core/java/com/android/server/integrity/serializer/RuleXmlSerializer.java
+++ b/services/core/java/com/android/server/integrity/serializer/RuleXmlSerializer.java
@@ -127,7 +127,7 @@ public class RuleXmlSerializer implements RuleSerializer {
         xmlSerializer.startTag(NAMESPACE, ATOMIC_FORMULA_TAG);
         serializeAttributeValue(
                 KEY_ATTRIBUTE, String.valueOf(atomicFormula.getKey()), xmlSerializer);
-        if (atomicFormula instanceof AtomicFormula.StringAtomicFormula) {
+        if (atomicFormula.getTag() == AtomicFormula.STRING_ATOMIC_FORMULA_TAG) {
             serializeAttributeValue(
                     VALUE_ATTRIBUTE,
                     ((AtomicFormula.StringAtomicFormula) atomicFormula).getValue(),
@@ -137,7 +137,7 @@ public class RuleXmlSerializer implements RuleSerializer {
                     String.valueOf(
                             ((AtomicFormula.StringAtomicFormula) atomicFormula).getIsHashedValue()),
                     xmlSerializer);
-        } else if (atomicFormula instanceof AtomicFormula.IntAtomicFormula) {
+        } else if (atomicFormula.getTag() == AtomicFormula.INT_ATOMIC_FORMULA_TAG) {
             serializeAttributeValue(
                     OPERATOR_ATTRIBUTE,
                     String.valueOf(((AtomicFormula.IntAtomicFormula) atomicFormula).getOperator()),
@@ -146,7 +146,7 @@ public class RuleXmlSerializer implements RuleSerializer {
                     VALUE_ATTRIBUTE,
                     String.valueOf(((AtomicFormula.IntAtomicFormula) atomicFormula).getValue()),
                     xmlSerializer);
-        } else if (atomicFormula instanceof AtomicFormula.BooleanAtomicFormula) {
+        } else if (atomicFormula.getTag() == AtomicFormula.BOOLEAN_ATOMIC_FORMULA_TAG) {
             serializeAttributeValue(
                     VALUE_ATTRIBUTE,
                     String.valueOf(((AtomicFormula.BooleanAtomicFormula) atomicFormula).getValue()),