Merge "Introduce a new output stream that is capable of tracking the number of bytes that are written to the output stream."

services/core/java/com/android/server/integrity/serializer/ByteTrackedOutputStream.java

@@ -0,0 +1,53 @@
+/*
+ * Copyright (C) 2020 The Android Open Source Project
+ *
+ * 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 com.android.server.integrity.serializer;
+
+import java.io.IOException;
+import java.io.OutputStream;
+
+/**
+ * An output stream that tracks the total number written bytes since construction and allows
+ * querying this value any time during the execution.
+ *
+ * This class is used for constructing the rule indexing.
+ */
+public class ByteTrackedOutputStream {
+
+    private static long sWrittenBytesCount;
+    private static OutputStream sOutputStream;
+
+    public ByteTrackedOutputStream(OutputStream outputStream) {
+        sWrittenBytesCount = 0;
+        sOutputStream = outputStream;
+    }
+
+    /**
+     * Writes the given bytes into the output stream provided in constructor and updates the
+     * total number of written bytes.
+     */
+    public void write(byte[] bytes) throws IOException {
+        sWrittenBytesCount += bytes.length;
+        sOutputStream.write(bytes);
+    }
+
+    /**
+     * Returns the total number of bytes written into the output stream at the requested time.
+     */
+    public long getWrittenBytesCount() {
+        return sWrittenBytesCount;
+    }
+}

services/tests/servicestests/src/com/android/server/integrity/serializer/ByteTrackedOutputStreamTest.java

@@ -0,0 +1,72 @@
+/*
+ * Copyright (C) 2020 The Android Open Source Project
+ *
+ * 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 com.android.server.integrity.serializer;
+
+import static com.google.common.truth.Truth.assertThat;
+
+import com.android.server.integrity.model.BitOutputStream;
+
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+import java.io.ByteArrayOutputStream;
+
+@RunWith(JUnit4.class)
+public class ByteTrackedOutputStreamTest {
+
+    @Test
+    public void testConstructorStartsWithZeroBytesWritten() {
+        ByteTrackedOutputStream byteTrackedOutputStream =
+                new ByteTrackedOutputStream(new ByteArrayOutputStream());
+
+        assertThat(byteTrackedOutputStream.getWrittenBytesCount()).isEqualTo(0);
+    }
+
+    @Test
+    public void testSuccessfulWriteAndValidateWrittenBytesCount_directFromByteArray()
+            throws Exception {
+        ByteArrayOutputStream outputStream = new ByteArrayOutputStream();
+        ByteTrackedOutputStream byteTrackedOutputStream = new ByteTrackedOutputStream(outputStream);
+
+        byte[] outputContent = "This is going to be outputed for tests.".getBytes();
+        byteTrackedOutputStream.write(outputContent);
+
+        assertThat(byteTrackedOutputStream.getWrittenBytesCount()).isEqualTo(outputContent.length);
+        assertThat(outputStream.toByteArray().length).isEqualTo(outputContent.length);
+    }
+
+    @Test
+    public void testSuccessfulWriteAndValidateWrittenBytesCount_fromBitStream() throws Exception {
+        ByteArrayOutputStream outputStream = new ByteArrayOutputStream();
+        ByteTrackedOutputStream byteTrackedOutputStream = new ByteTrackedOutputStream(outputStream);
+
+        BitOutputStream bitOutputStream = new BitOutputStream();
+        bitOutputStream.setNext(/* numOfBits= */5, /* value= */1);
+        byteTrackedOutputStream.write(bitOutputStream.toByteArray());
+
+        // Even though we wrote 5 bits, this will complete to 1 byte.
+        assertThat(byteTrackedOutputStream.getWrittenBytesCount()).isEqualTo(1);
+
+        // Add a bit less than 2 bytes (10 bits).
+        bitOutputStream.clear();
+        bitOutputStream.setNext(/* numOfBits= */10, /* value= */1);
+        byteTrackedOutputStream.write(bitOutputStream.toByteArray());
+
+        assertThat(outputStream.toByteArray().length).isEqualTo(3);
+    }
+}