1 /*
2  * Copyright (C) 2017 The Android Open Source Project
3  *
4  * Licensed under the Apache License, Version 2.0 (the "License");
5  * you may not use this file except in compliance with the License.
6  * You may obtain a copy of the License at
7  *
8  *      http://www.apache.org/licenses/LICENSE-2.0
9  *
10  * Unless required by applicable law or agreed to in writing, software
11  * distributed under the License is distributed on an "AS IS" BASIS,
12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13  * See the License for the specific language governing permissions and
14  * limitations under the License.
15  */
16 
17 package com.android.apksig.util;
18 
19 import static org.junit.Assert.assertEquals;
20 import static org.junit.Assert.fail;
21 
22 import java.io.Closeable;
23 import java.io.IOException;
24 import java.nio.ByteBuffer;
25 import java.nio.charset.StandardCharsets;
26 import org.junit.Test;
27 
28 /**
29  * Base class for testing implementations of {@link DataSink}. This class tests the contract of
30  * {@code DataSink}.
31  *
32  * <p>To subclass, provide an implementation of {@link #createDataSink()} which returns the
33  * implementation of {@code DataSink} you want to test.
34  */
35 public abstract class DataSinkTestBase<T extends DataSink> {
36     /**
37      * Returns a new {@link DataSink}.
38      */
createDataSink()39     protected abstract CloseableWithDataSink<T> createDataSink() throws IOException;
40 
41     /**
42      * Returns the contents of the data sink.
43      */
getContents(T dataSink)44     protected abstract ByteBuffer getContents(T dataSink) throws IOException;
45 
46     @Test
testConsumeFromArray()47     public void testConsumeFromArray() throws Exception {
48         try (CloseableWithDataSink<T> c = createDataSink()) {
49             T sink = c.getDataSink();
50             byte[] input = "abcdefg".getBytes(StandardCharsets.UTF_8);
51             sink.consume(input, 2, 3); // "cde"
52             sink.consume(input, 0, 1); // "a"
53             assertContentsEquals("cdea", sink);
54 
55             // Zero-length chunks
56             sink.consume(input, 0, 0);
57             sink.consume(input, 1, 0);
58             sink.consume(input, input.length - 2, 0);
59             sink.consume(input, input.length - 1, 0);
60             sink.consume(input, input.length, 0);
61 
62             // Invalid chunks
63             assertConsumeArrayThrowsIOOB(sink, input, -1, 0);
64             assertConsumeArrayThrowsIOOB(sink, input, -1, 3);
65             assertConsumeArrayThrowsIOOB(sink, input, 0, input.length + 1);
66             assertConsumeArrayThrowsIOOB(sink, input, input.length - 2, 4);
67             assertConsumeArrayThrowsIOOB(sink, input, input.length + 1, 0);
68             assertConsumeArrayThrowsIOOB(sink, input, input.length + 1, 1);
69 
70             assertContentsEquals("cdea", sink);
71         }
72     }
73 
74     @Test
testConsumeFromByteBuffer()75     public void testConsumeFromByteBuffer() throws Exception {
76         try (CloseableWithDataSink<T> c = createDataSink()) {
77             T sink = c.getDataSink();
78             ByteBuffer input = ByteBuffer.wrap("abcdefg".getBytes(StandardCharsets.UTF_8));
79             input.position(2);
80             input.limit(5);
81             sink.consume(input); // "cde"
82             assertEquals(5, input.position());
83             assertEquals(5, input.limit());
84 
85             input.position(0);
86             input.limit(1);
87             sink.consume(input); // "a"
88             assertContentsEquals("cdea", sink);
89 
90             // Empty input
91             sink.consume(input);
92             assertContentsEquals("cdea", sink);
93 
94             // ByteBuffer which isn't backed by a byte[]
95             input = ByteBuffer.allocateDirect(2);
96             input.put((byte) 'X');
97             input.put((byte) 'Z');
98             input.flip();
99             sink.consume(input);
100 
101             assertContentsEquals("cdeaXZ", sink);
102             assertEquals(2, input.position());
103             assertEquals(2, input.limit());
104 
105             // Empty input
106             sink.consume(input);
107             assertContentsEquals("cdeaXZ", sink);
108         }
109     }
110 
111     /**
112      * Returns the contents of the provided buffer as a string. The buffer's position and limit
113      * remain unchanged.
114      */
toString(ByteBuffer buf)115     private static String toString(ByteBuffer buf) {
116         return DataSourceTestBase.toString(buf);
117     }
118 
assertContentsEquals(String expectedContents, T sink)119     private void assertContentsEquals(String expectedContents, T sink) throws IOException {
120         ByteBuffer actual = getContents(sink);
121         assertEquals(expectedContents, toString(actual));
122     }
123 
assertConsumeArrayThrowsIOOB( DataSink sink, byte[] arr, int offset, int length)124     private static void assertConsumeArrayThrowsIOOB(
125             DataSink sink, byte[] arr, int offset, int length) throws IOException {
126         try {
127             sink.consume(arr, offset, length);
128             fail();
129         } catch (IndexOutOfBoundsException expected) {}
130     }
131 
132     public static class CloseableWithDataSink<T extends DataSink> implements Closeable {
133         private final T mDataSink;
134         private final Closeable mCloseable;
135 
CloseableWithDataSink(T dataSink, Closeable closeable)136         private CloseableWithDataSink(T dataSink, Closeable closeable) {
137             mDataSink = dataSink;
138             mCloseable = closeable;
139         }
140 
of(T dataSink)141         public static <T extends DataSink> CloseableWithDataSink<T> of(T dataSink) {
142             return new CloseableWithDataSink<>(dataSink, null);
143         }
144 
of( T dataSink, Closeable closeable)145         public static <T extends DataSink> CloseableWithDataSink<T> of(
146                 T dataSink, Closeable closeable) {
147             return new CloseableWithDataSink<>(dataSink, closeable);
148         }
149 
getDataSink()150         public T getDataSink() {
151             return mDataSink;
152         }
153 
getCloseable()154         public Closeable getCloseable() {
155             return mCloseable;
156         }
157 
158         @Override
close()159         public void close() throws IOException {
160             if (mCloseable != null) {
161                 mCloseable.close();
162             }
163         }
164     }
165 }
166