src/main/java/com/android/tools/build/apkzlib/zip/ExtraField.java

/*
 * Copyright (C) 2016 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.tools.build.apkzlib.zip;

import com.android.tools.build.apkzlib.zip.utils.LittleEndianUtils;
import com.google.common.base.Preconditions;
import com.google.common.collect.ImmutableList;
import java.io.IOException;
import java.nio.ByteBuffer;
import java.util.ArrayList;
import java.util.List;
import java.util.stream.Collectors;
import javax.annotation.Nonnull;
import javax.annotation.Nullable;

/**
 * Contains an extra field.
 *
 * <p>According to the zip specification, the extra field is composed of a sequence of fields.
 * This class provides a way to access, parse and modify that information.
 *
 * <p>The zip specification calls fields to the fields inside the extra field. Because this
 * terminology is confusing, we use <i>segment</i> to refer to a part of the extra field. Each
 * segment is represented by an instance of {@link Segment} and contains a header ID and data.
 *
 * <p>Each instance of {@link ExtraField} is immutable. The extra field of a particular entry can
 * be changed by creating a new instanceof {@link ExtraField} and pass it to
 * {@link StoredEntry#setLocalExtra(ExtraField)}.
 *
 * <p>Instances of {@link ExtraField} can be created directly from the list of segments in it
 * or from the raw byte data. If created from the raw byte data, the data will only be parsed
 * on demand. So, if neither {@link #getSegments()} nor {@link #getSingleSegment(int)} is
 * invoked, the extra field will not be parsed. This guarantees low performance impact of the
 * using the extra field unless its contents are needed.
 */
public class ExtraField {

    /**
     * Header ID for field with zip alignment.
     */
    static final int ALIGNMENT_ZIP_EXTRA_DATA_FIELD_HEADER_ID = 0xd935;

    /**
     * The field's raw data, if it is known. Either this variable or {@link #segments} must be
     * non-{@code null}.
     */
    @Nullable
    private final byte[] rawData;

    /**
     * The list of field's segments. Will be populated if the extra field is created based on a
     * list of segments; will also be populated after parsing if the extra field is created based
     * on the raw bytes.
     */
    @Nullable
    private ImmutableList<Segment> segments;

    /**
     * Creates an extra field based on existing raw data.
     *
     * @param rawData the raw data; will not be parsed unless needed
     */
    public ExtraField(@Nonnull byte[] rawData) {
        this.rawData = rawData;
        segments = null;
    }

    /**
     * Creates a new extra field with no segments.
     */
    public ExtraField() {
        rawData = null;
        segments = ImmutableList.of();
    }

    /**
     * Creates a new extra field with the given segments.
     *
     * @param segments the segments
     */
    public ExtraField(@Nonnull ImmutableList<Segment> segments) {
        rawData = null;
        this.segments = segments;
    }

    /**
     * Obtains all segments in the extra field.
     *
     * @return all segments
     * @throws IOException failed to parse the extra field
     */
    public ImmutableList<Segment> getSegments() throws IOException {
        if (segments == null) {
            parseSegments();
        }

        Preconditions.checkNotNull(segments);
        return segments;
    }

    /**
     * Obtains the only segment with the provided header ID.
     *
     * @param headerId the header ID
     * @return the segment found or {@code null} if no segment contains the provided header ID
     * @throws IOException there is more than one header with the provided header ID
     */
    @Nullable
    public Segment getSingleSegment(int headerId) throws IOException {
        List<Segment> found =
                getSegments().stream()
                        .filter(s -> s.getHeaderId() == headerId)
                        .collect(Collectors.toList());
        if (found.isEmpty()) {
            return null;
        } else if (found.size() == 1) {
            return found.get(0);
        } else {
            throw new IOException(found.size() + " segments with header ID " + headerId + "found");
        }
    }

    /**
     * Parses the raw data and generates all segments in {@link #segments}.
     *
     * @throws IOException failed to parse the data
     */
    private void parseSegments() throws IOException {
        Preconditions.checkNotNull(rawData);
        Preconditions.checkState(segments == null);

        List<Segment> segments = new ArrayList<>();
        ByteBuffer buffer = ByteBuffer.wrap(rawData);

        while (buffer.remaining() > 0) {
            int headerId = LittleEndianUtils.readUnsigned2Le(buffer);
            int dataSize = LittleEndianUtils.readUnsigned2Le(buffer);
            if (dataSize < 0) {
                throw new IOException(
                        "Invalid data size for extra field segment with header ID "
                                + headerId
                                + ": "
                                + dataSize);
            }

            byte[] data = new byte[dataSize];
            if (buffer.remaining() < dataSize) {
                throw new IOException(
                        "Invalid data size for extra field segment with header ID "
                                + headerId
                                + ": "
                                + dataSize
                                + " (only "
                                + buffer.remaining()
                                + " bytes are available)");
            }
            buffer.get(data);

            SegmentFactory factory = identifySegmentFactory(headerId);
            Segment seg = factory.make(headerId, data);
            segments.add(seg);
        }

        this.segments = ImmutableList.copyOf(segments);
    }

    /**
     * Obtains the size of the extra field.
     *
     * @return the size
     */
    public int size() {
        if (rawData != null) {
            return rawData.length;
        } else {
            Preconditions.checkNotNull(segments);
            int sz = 0;
            for (Segment s : segments) {
                sz += s.size();
            }

            return sz;
        }
    }

    /**
     * Writes the extra field to the given output buffer.
     *
     * @param out the output buffer to write the field; exactly {@link #size()} bytes will be
     * written
     * @throws IOException failed to write the extra fields
     */
    public void write(@Nonnull ByteBuffer out) throws IOException {
        if (rawData != null) {
            out.put(rawData);
        } else {
            Preconditions.checkNotNull(segments);
            for (Segment s : segments) {
                s.write(out);
            }
        }
    }

    /**
     * Identifies the factory to create the segment with the provided header ID.
     *
     * @param headerId the header ID
     * @return the segmnet factory that creates segments with the given header
     */
    @Nonnull
    private static SegmentFactory identifySegmentFactory(int headerId) {
        if (headerId == ALIGNMENT_ZIP_EXTRA_DATA_FIELD_HEADER_ID) {
            return AlignmentSegment::new;
        }

        return RawDataSegment::new;
    }

    /**
     * Field inside the extra field. A segment contains a header ID and data. Specific types of
     * segments implement this interface.
     */
    public interface Segment {

        /**
         * Obtains the segment's header ID.
         *
         * @return the segment's header ID
         */
        int getHeaderId();

        /**
         * Obtains the size of the segment including the header ID.
         *
         * @return the number of bytes needed to write the segment
         */
        int size();

        /**
         * Writes the segment to a buffer.
         *
         * @param out the buffer where to write the segment to; exactly {@link #size()} bytes will
         * be written
         * @throws IOException failed to write segment data
         */
        void write(@Nonnull ByteBuffer out) throws IOException;
    }

    /**
     * Factory that creates a segment.
     */
    @FunctionalInterface
    interface SegmentFactory {

        /**
         * Creates a new segment.
         *
         * @param headerId the header ID
         * @param data the segment's data
         * @return the created segment
         * @throws IOException failed to create the segment from the data
         */
        @Nonnull
        Segment make(int headerId, @Nonnull byte[] data) throws IOException;
    }

    /**
     * Segment of raw data: this class represents a general segment containing an array of bytes
     * as data.
     */
    public static class RawDataSegment implements Segment {

        /**
         * Header ID.
         */
        private final int headerId;

        /**
         * Data in the segment.
         */
        @Nonnull
        private final byte[] data;

        /**
         * Creates a new raw data segment.
         *
         * @param headerId the header ID
         * @param data the segment data
         */
        RawDataSegment(int headerId, @Nonnull byte[] data) {
            this.headerId = headerId;
            this.data = data;
        }

        @Override
        public int getHeaderId() {
            return headerId;
        }

        @Override
        public void write(@Nonnull ByteBuffer out) throws IOException {
            LittleEndianUtils.writeUnsigned2Le(out, headerId);
            LittleEndianUtils.writeUnsigned2Le(out, data.length);
            out.put(data);
        }

        @Override
        public int size() {
            return 4 + data.length;
        }
    }

    /**
     * Segment with information on an alignment: this segment contains information on how an entry
     * should be aligned and contains zero-filled data to force alignment.
     *
     * <p>An alignment segment contains the header ID, the size of the data, the alignment value
     * and zero bytes to pad
     */
    public static class AlignmentSegment implements Segment {

        /**
         * Minimum size for an alignment segment.
         */
        public static final int MINIMUM_SIZE = 6;

        /**
         * The alignment value.
         */
        private int alignment;

        /**
         * How many bytes of padding are in this segment?
         */
        private int padding;

        /**
         * Creates a new alignment segment.
         *
         * @param alignment the alignment value
         * @param totalSize how many bytes should this segment take?
         */
        public AlignmentSegment(int alignment, int totalSize) {
            Preconditions.checkArgument(alignment > 0, "alignment <= 0");
            Preconditions.checkArgument(totalSize >= MINIMUM_SIZE, "totalSize < MINIMUM_SIZE");

            /*
             * We have 6 bytes of fixed data: header ID (2 bytes), data size (2 bytes), alignment
             * value (2 bytes).
             */
            this.alignment = alignment;
            padding = totalSize - MINIMUM_SIZE;
        }

        /**
         * Creates a new alignment segment from extra data.
         *
         * @param headerId the header ID
         * @param data the segment data
         * @throws IOException failed to create the segment from the data
         */
        public AlignmentSegment(int headerId, @Nonnull byte[] data) throws IOException {
            Preconditions.checkArgument(headerId == ALIGNMENT_ZIP_EXTRA_DATA_FIELD_HEADER_ID);

            ByteBuffer dataBuffer = ByteBuffer.wrap(data);
            alignment = LittleEndianUtils.readUnsigned2Le(dataBuffer);
            if (alignment <= 0) {
                throw new IOException("Invalid alignment in alignment field: " + alignment);
            }

            padding = data.length - 2;
        }

        @Override
        public void write(@Nonnull ByteBuffer out) throws IOException {
            LittleEndianUtils.writeUnsigned2Le(out, ALIGNMENT_ZIP_EXTRA_DATA_FIELD_HEADER_ID);
            LittleEndianUtils.writeUnsigned2Le(out, padding + 2);
            LittleEndianUtils.writeUnsigned2Le(out, alignment);
            out.put(new byte[padding]);
        }

        @Override
        public int size() {
            return padding + 6;
        }

        @Override
        public int getHeaderId() {
            return ALIGNMENT_ZIP_EXTRA_DATA_FIELD_HEADER_ID;
        }
    }
}