diff --git a/services/core/java/com/android/server/integrity/serializer/ByteTrackedOutputStream.java b/services/core/java/com/android/server/integrity/serializer/ByteTrackedOutputStream.java new file mode 100644 index 0000000000000..c8d318f5096fb --- /dev/null +++ b/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; + } +} diff --git a/services/tests/servicestests/src/com/android/server/integrity/serializer/ByteTrackedOutputStreamTest.java b/services/tests/servicestests/src/com/android/server/integrity/serializer/ByteTrackedOutputStreamTest.java new file mode 100644 index 0000000000000..5ecb8b5c81695 --- /dev/null +++ b/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); + } +}