content: /*
   * 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 - 1);
          int endIndex = lastSource.getFirst();
          long endLocalOffset = lastSource.getSecond();
  
          for (int i = beginIndex + 1; i < endIndex; i++) {
              sources.add(mSources[i]);
          }
  
          sources.add(mSources[endIndex].slice(0, endLocalOffset + 1));
          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);
      }
  }
  