Home | History | Annotate | Download | only in util 
      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.internal.util;
     18 
     19 import com.android.apksig.util.DataSink;
     20 import com.android.apksig.util.DataSource;
     21 import java.io.IOException;
     22 import java.nio.ByteBuffer;
     23 import java.util.ArrayList;
     24 import java.util.Arrays;
     25 
     26 /** Pseudo {@link DataSource} that chains the given {@link DataSource} as a continuous one. */
     27 public class ChainedDataSource implements DataSource {
     28 
     29     private final DataSource[] mSources;
     30     private final long mTotalSize;
     31 
     32     public ChainedDataSource(DataSource... sources) {
     33         mSources = sources;
     34         mTotalSize = Arrays.stream(sources).mapToLong(src -> src.size()).sum();
     35     }
     36 
     37     @Override
     38     public long size() {
     39         return mTotalSize;
     40     }
     41 
     42     @Override
     43     public void feed(long offset, long size, DataSink sink) throws IOException {
     44         if (offset + size > mTotalSize) {
     45             throw new IndexOutOfBoundsException("Requested more than available");
     46         }
     47 
     48         for (DataSource src : mSources) {
     49             // Offset is beyond the current source. Skip.
     50             if (offset >= src.size()) {
     51                 offset -= src.size();
     52                 continue;
     53             }
     54 
     55             // If the remaining is enough, finish it.
     56             long remaining = src.size() - offset;
     57             if (remaining >= size) {
     58                 src.feed(offset, size, sink);
     59                 break;
     60             }
     61 
     62             // If the remaining is not enough, consume all.
     63             src.feed(offset, remaining, sink);
     64             size -= remaining;
     65             offset = 0;
     66         }
     67     }
     68 
     69     @Override
     70     public ByteBuffer getByteBuffer(long offset, int size) throws IOException {
     71         if (offset + size > mTotalSize) {
     72             throw new IndexOutOfBoundsException("Requested more than available");
     73         }
     74 
     75         // Skip to the first DataSource we need.
     76         Pair<Integer, Long> firstSource = locateDataSource(offset);
     77         int i = firstSource.getFirst();
     78         offset = firstSource.getSecond();
     79 
     80         // Return the current source's ByteBuffer if it fits.
     81         if (offset + size <= mSources[i].size()) {
     82             return mSources[i].getByteBuffer(offset, size);
     83         }
     84 
     85         // Otherwise, read into a new buffer.
     86         ByteBuffer buffer = ByteBuffer.allocate(size);
     87         for (; i < mSources.length && buffer.hasRemaining(); i++) {
     88             long sizeToCopy = Math.min(mSources[i].size() - offset, buffer.remaining());
     89             mSources[i].copyTo(offset, Math.toIntExact(sizeToCopy), buffer);
     90             offset = 0;  // may not be zero for the first source, but reset after that.
     91         }
     92         buffer.rewind();
     93         return buffer;
     94     }
     95 
     96     @Override
     97     public void copyTo(long offset, int size, ByteBuffer dest) throws IOException {
     98         feed(offset, size, new ByteBufferSink(dest));
     99     }
    100 
    101     @Override
    102     public DataSource slice(long offset, long size) {
    103         // Find the first slice.
    104         Pair<Integer, Long> firstSource = locateDataSource(offset);
    105         int beginIndex = firstSource.getFirst();
    106         long beginLocalOffset = firstSource.getSecond();
    107         DataSource beginSource = mSources[beginIndex];
    108 
    109         if (beginLocalOffset + size <= beginSource.size()) {
    110             return beginSource.slice(beginLocalOffset, size);
    111         }
    112 
    113         // Add the first slice to chaining, followed by the middle full slices, then the last.
    114         ArrayList<DataSource> sources = new ArrayList<>();
    115         sources.add(beginSource.slice(
    116                 beginLocalOffset, beginSource.size() - beginLocalOffset));
    117 
    118         Pair<Integer, Long> lastSource = locateDataSource(offset + size);
    119         int endIndex = lastSource.getFirst();
    120         long endLocalOffset = lastSource.getSecond();
    121 
    122         for (int i = beginIndex + 1; i < endIndex - 1; i++) {
    123             sources.add(mSources[i]);
    124         }
    125 
    126         sources.add(mSources[endIndex].slice(0, endLocalOffset));
    127         return new ChainedDataSource(sources.toArray(new DataSource[0]));
    128     }
    129 
    130     /**
    131      * Find the index of DataSource that offset is at.
    132      * @return Pair of DataSource index and the local offset in the DataSource.
    133      */
    134     private Pair<Integer, Long> locateDataSource(long offset) {
    135         long localOffset = offset;
    136         for (int i = 0; i < mSources.length; i++) {
    137             if (localOffset < mSources[i].size()) {
    138                 return Pair.of(i, localOffset);
    139             }
    140             localOffset -= mSources[i].size();
    141         }
    142         throw new IndexOutOfBoundsException("Access is out of bound, offset: " + offset +
    143                 ", totalSize: " + mTotalSize);
    144     }
    145 }
    146