src/main/java/com/android/apksig/internal/util/ChainedDataSource.java

/*
 * Copyright (C) 2017 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.apksig.internal.util;

import com.android.apksig.util.DataSink;
import com.android.apksig.util.DataSource;
import java.io.IOException;
import java.nio.ByteBuffer;
import java.util.ArrayList;
import java.util.Arrays;

/** Pseudo {@link DataSource} that chains the given {@link DataSource} as a continuous one. */
public class ChainedDataSource implements DataSource {

    private final DataSource[] mSources;
    private final long mTotalSize;

    public ChainedDataSource(DataSource... sources) {
        mSources = sources;
        mTotalSize = Arrays.stream(sources).mapToLong(src -> src.size()).sum();
    }

    @Override
    public long size() {
        return mTotalSize;
    }

    @Override
    public void feed(long offset, long size, DataSink sink) throws IOException {
        if (offset + size > mTotalSize) {
            throw new IndexOutOfBoundsException("Requested more than available");
        }

        for (DataSource src : mSources) {
            // Offset is beyond the current source. Skip.
            if (offset >= src.size()) {
                offset -= src.size();
                continue;
            }

            // If the remaining is enough, finish it.
            long remaining = src.size() - offset;
            if (remaining >= size) {
                src.feed(offset, size, sink);
                break;
            }

            // If the remaining is not enough, consume all.
            src.feed(offset, remaining, sink);
            size -= remaining;
            offset = 0;
        }
    }

    @Override
    public ByteBuffer getByteBuffer(long offset, int size) throws IOException {
        if (offset + size > mTotalSize) {
            throw new IndexOutOfBoundsException("Requested more than available");
        }

        // Skip to the first DataSource we need.
        Pair<Integer, Long> firstSource = locateDataSource(offset);
        int i = firstSource.getFirst();
        offset = firstSource.getSecond();

        // Return the current source's ByteBuffer if it fits.
        if (offset + size <= mSources[i].size()) {
            return mSources[i].getByteBuffer(offset, size);
        }

        // Otherwise, read into a new buffer.
        ByteBuffer buffer = ByteBuffer.allocate(size);
        for (; i < mSources.length && buffer.hasRemaining(); i++) {
            long sizeToCopy = Math.min(mSources[i].size() - offset, buffer.remaining());
            mSources[i].copyTo(offset, Math.toIntExact(sizeToCopy), buffer);
            offset = 0;  // may not be zero for the first source, but reset after that.
        }
        buffer.rewind();
        return buffer;
    }

    @Override
    public void copyTo(long offset, int size, ByteBuffer dest) throws IOException {
        feed(offset, size, new ByteBufferSink(dest));
    }

    @Override
    public DataSource slice(long offset, long size) {
        // Find the first slice.
        Pair<Integer, Long> firstSource = locateDataSource(offset);
        int beginIndex = firstSource.getFirst();
        long beginLocalOffset = firstSource.getSecond();
        DataSource beginSource = mSources[beginIndex];

        if (beginLocalOffset + size <= beginSource.size()) {
            return beginSource.slice(beginLocalOffset, size);
        }

        // Add the first slice to chaining, followed by the middle full slices, then the last.
        ArrayList<DataSource> sources = new ArrayList<>();
        sources.add(beginSource.slice(
                beginLocalOffset, beginSource.size() - beginLocalOffset));

        Pair<Integer, Long> lastSource = locateDataSource(offset + size);
        int endIndex = lastSource.getFirst();
        long endLocalOffset = lastSource.getSecond();

        for (int i = beginIndex + 1; i < endIndex - 1; i++) {
            sources.add(mSources[i]);
        }

        sources.add(mSources[endIndex].slice(0, endLocalOffset));
        return new ChainedDataSource(sources.toArray(new DataSource[0]));
    }

    /**
     * Find the index of DataSource that offset is at.
     * @return Pair of DataSource index and the local offset in the DataSource.
     */
    private Pair<Integer, Long> locateDataSource(long offset) {
        long localOffset = offset;
        for (int i = 0; i < mSources.length; i++) {
            if (localOffset < mSources[i].size()) {
                return Pair.of(i, localOffset);
            }
            localOffset -= mSources[i].size();
        }
        throw new IndexOutOfBoundsException("Access is out of bound, offset: " + offset +
                ", totalSize: " + mTotalSize);
    }
}