    001/*
    002 * Licensed to the Apache Software Foundation (ASF) under one
    003 * or more contributor license agreements.  See the NOTICE file
    004 * distributed with this work for additional information
    005 * regarding copyright ownership.  The ASF licenses this file
    006 * to you under the Apache License, Version 2.0 (the
    007 * "License"); you may not use this file except in compliance
    008 * with the License.  You may obtain a copy of the License at
    009 *
    010 * http://www.apache.org/licenses/LICENSE-2.0
    011 *
    012 * Unless required by applicable law or agreed to in writing,
    013 * software distributed under the License is distributed on an
    014 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
    015 * KIND, either express or implied.  See the License for the
    016 * specific language governing permissions and limitations
    017 * under the License.
    018 */
    019package org.apache.commons.compress.compressors.lz4;
    020
    021import java.io.ByteArrayOutputStream;
    022import java.io.IOException;
    023import java.io.OutputStream;
    024
    025import org.apache.commons.compress.compressors.CompressorOutputStream;
    026import org.apache.commons.compress.utils.ByteUtils;
    027
    028/**
    029 * CompressorOutputStream for the LZ4 frame format.
    030 *
    031 * <p>
    032 * Based on the "spec" in the version "1.5.1 (31/03/2015)"
    033 * </p>
    034 *
    035 * @see <a href="https://lz4.github.io/lz4/lz4_Frame_format.html">LZ4 Frame Format Description</a>
    036 * @since 1.14
    037 * @NotThreadSafe
    038 */
    039public class FramedLZ4CompressorOutputStream extends CompressorOutputStream<OutputStream> {
    040
    041    /**
    042     * The block sizes supported by the format.
    043     */
    044    public enum BlockSize {
    045        /** Block size of 64K */
    046        K64(64 * 1024, 4),
    047        /** Block size of 256K */
    048        K256(256 * 1024, 5),
    049        /** Block size of 1M */
    050        M1(1024 * 1024, 6),
    051        /** Block size of 4M */
    052        M4(4096 * 1024, 7);
    053
    054        private final int size, index;
    055
    056        BlockSize(final int size, final int index) {
    057            this.size = size;
    058            this.index = index;
    059        }
    060
    061        int getIndex() {
    062            return index;
    063        }
    064
    065        int getSize() {
    066            return size;
    067        }
    068    }
    069
    070    /**
    071     * Parameters of the LZ4 frame format.
    072     */
    073    public static class Parameters {
    074        /**
    075         * The default parameters of 4M block size, enabled content checksum, disabled block checksums and independent blocks.
    076         *
    077         * <p>
    078         * This matches the defaults of the lz4 command line utility.
    079         * </p>
    080         */
    081        public static final Parameters DEFAULT = new Parameters(BlockSize.M4, true, false, false);
    082        private final BlockSize blockSize;
    083        private final boolean withContentChecksum, withBlockChecksum, withBlockDependency;
    084
    085        private final org.apache.commons.compress.compressors.lz77support.Parameters lz77params;
    086
    087        /**
    088         * Sets up custom a custom block size for the LZ4 stream but otherwise uses the defaults of enabled content checksum, disabled block checksums and
    089         * independent blocks.
    090         *
    091         * @param blockSize the size of a single block.
    092         */
    093        public Parameters(final BlockSize blockSize) {
    094            this(blockSize, true, false, false);
    095        }
    096
    097        /**
    098         * Sets up custom parameters for the LZ4 stream.
    099         *
    100         * @param blockSize           the size of a single block.
    101         * @param withContentChecksum whether to write a content checksum
    102         * @param withBlockChecksum   whether to write a block checksum. Note that block checksums are not supported by the lz4 command line utility
    103         * @param withBlockDependency whether a block may depend on the content of a previous block. Enabling this may improve compression ratio but makes it
    104         *                            impossible to decompress the output in parallel.
    105         */
    106        public Parameters(final BlockSize blockSize, final boolean withContentChecksum, final boolean withBlockChecksum, final boolean withBlockDependency) {
    107            this(blockSize, withContentChecksum, withBlockChecksum, withBlockDependency, BlockLZ4CompressorOutputStream.createParameterBuilder().build());
    108        }
    109
    110        /**
    111         * Sets up custom parameters for the LZ4 stream.
    112         *
    113         * @param blockSize           the size of a single block.
    114         * @param withContentChecksum whether to write a content checksum
    115         * @param withBlockChecksum   whether to write a block checksum. Note that block checksums are not supported by the lz4 command line utility
    116         * @param withBlockDependency whether a block may depend on the content of a previous block. Enabling this may improve compression ratio but makes it
    117         *                            impossible to decompress the output in parallel.
    118         * @param lz77params          parameters used to fine-tune compression, in particular to balance compression ratio vs compression speed.
    119         */
    120        public Parameters(final BlockSize blockSize, final boolean withContentChecksum, final boolean withBlockChecksum, final boolean withBlockDependency,
    121                final org.apache.commons.compress.compressors.lz77support.Parameters lz77params) {
    122            this.blockSize = blockSize;
    123            this.withContentChecksum = withContentChecksum;
    124            this.withBlockChecksum = withBlockChecksum;
    125            this.withBlockDependency = withBlockDependency;
    126            this.lz77params = lz77params;
    127        }
    128
    129        /**
    130         * Sets up custom a custom block size for the LZ4 stream but otherwise uses the defaults of enabled content checksum, disabled block checksums and
    131         * independent blocks.
    132         *
    133         * @param blockSize  the size of a single block.
    134         * @param lz77params parameters used to fine-tune compression, in particular to balance compression ratio vs compression speed.
    135         */
    136        public Parameters(final BlockSize blockSize, final org.apache.commons.compress.compressors.lz77support.Parameters lz77params) {
    137            this(blockSize, true, false, false, lz77params);
    138        }
    139
    140        @Override
    141        public String toString() {
    142            return "LZ4 Parameters with BlockSize " + blockSize + ", withContentChecksum " + withContentChecksum + ", withBlockChecksum " + withBlockChecksum
    143                    + ", withBlockDependency " + withBlockDependency;
    144        }
    145    }
    146
    147    private static final byte[] END_MARK = new byte[4];
    148    // used in one-arg write method
    149    private final byte[] oneByte = new byte[1];
    150    private final byte[] blockData;
    151    private final Parameters params;
    152
    153    private boolean finished;
    154
    155    // used for frame header checksum and content checksum, if requested
    156    private final org.apache.commons.codec.digest.XXHash32 contentHash = new org.apache.commons.codec.digest.XXHash32();
    157    // used for block checksum, if requested
    158    private final org.apache.commons.codec.digest.XXHash32 blockHash;
    159
    160    // only created if the config requires block dependency
    161    private final byte[] blockDependencyBuffer;
    162
    163    private int collectedBlockDependencyBytes;
    164    private int currentIndex;
    165
    166    /**
    167     * Constructs a new output stream that compresses data using the LZ4 frame format using the default block size of 4MB.
    168     *
    169     * @param out the OutputStream to which to write the compressed data
    170     * @throws IOException if writing the signature fails
    171     */
    172    public FramedLZ4CompressorOutputStream(final OutputStream out) throws IOException {
    173        this(out, Parameters.DEFAULT);
    174    }
    175
    176    /**
    177     * Constructs a new output stream that compresses data using the LZ4 frame format using the given block size.
    178     *
    179     * @param out    the OutputStream to which to write the compressed data
    180     * @param params the parameters to use
    181     * @throws IOException if writing the signature fails
    182     */
    183    public FramedLZ4CompressorOutputStream(final OutputStream out, final Parameters params) throws IOException {
    184        super(out);
    185        this.params = params;
    186        blockData = new byte[params.blockSize.getSize()];
    187        blockHash = params.withBlockChecksum ? new org.apache.commons.codec.digest.XXHash32() : null;
    188        out.write(FramedLZ4CompressorInputStream.LZ4_SIGNATURE);
    189        writeFrameDescriptor();
    190        blockDependencyBuffer = params.withBlockDependency ? new byte[BlockLZ4CompressorInputStream.WINDOW_SIZE] : null;
    191    }
    192
    193    private void appendToBlockDependencyBuffer(final byte[] b, final int off, int len) {
    194        len = Math.min(len, blockDependencyBuffer.length);
    195        if (len > 0) {
    196            final int keep = blockDependencyBuffer.length - len;
    197            if (keep > 0) {
    198                // move last keep bytes towards the start of the buffer
    199                System.arraycopy(blockDependencyBuffer, len, blockDependencyBuffer, 0, keep);
    200            }
    201            // append new data
    202            System.arraycopy(b, off, blockDependencyBuffer, keep, len);
    203            collectedBlockDependencyBytes = Math.min(collectedBlockDependencyBytes + len, blockDependencyBuffer.length);
    204        }
    205    }
    206
    207    @Override
    208    public void close() throws IOException {
    209        try {
    210            finish();
    211        } finally {
    212            out.close();
    213        }
    214    }
    215
    216    /**
    217     * Compresses all blockDataRemaining data and writes it to the stream, doesn't close the underlying stream.
    218     *
    219     * @throws IOException if an error occurs
    220     */
    221    public void finish() throws IOException {
    222        if (!finished) {
    223            flushBlock();
    224            writeTrailer();
    225            finished = true;
    226        }
    227    }
    228
    229    private void flushBlock() throws IOException {
    230        if (currentIndex == 0) {
    231            return;
    232        }
    233        final boolean withBlockDependency = params.withBlockDependency;
    234        final ByteArrayOutputStream baos = new ByteArrayOutputStream();
    235        try (BlockLZ4CompressorOutputStream o = new BlockLZ4CompressorOutputStream(baos, params.lz77params)) {
    236            if (withBlockDependency) {
    237                o.prefill(blockDependencyBuffer, blockDependencyBuffer.length - collectedBlockDependencyBytes, collectedBlockDependencyBytes);
    238            }
    239            o.write(blockData, 0, currentIndex);
    240        }
    241        if (withBlockDependency) {
    242            appendToBlockDependencyBuffer(blockData, 0, currentIndex);
    243        }
    244        final byte[] b = baos.toByteArray();
    245        if (b.length > currentIndex) { // compression increased size, maybe beyond blocksize
    246            ByteUtils.toLittleEndian(out, currentIndex | FramedLZ4CompressorInputStream.UNCOMPRESSED_FLAG_MASK, 4);
    247            out.write(blockData, 0, currentIndex);
    248            if (params.withBlockChecksum) {
    249                blockHash.update(blockData, 0, currentIndex);
    250            }
    251        } else {
    252            ByteUtils.toLittleEndian(out, b.length, 4);
    253            out.write(b);
    254            if (params.withBlockChecksum) {
    255                blockHash.update(b, 0, b.length);
    256            }
    257        }
    258        if (params.withBlockChecksum) {
    259            ByteUtils.toLittleEndian(out, blockHash.getValue(), 4);
    260            blockHash.reset();
    261        }
    262        currentIndex = 0;
    263    }
    264
    265    @Override
    266    public void write(final byte[] data, int off, int len) throws IOException {
    267        if (params.withContentChecksum) {
    268            contentHash.update(data, off, len);
    269        }
    270        int blockDataRemaining = blockData.length - currentIndex;
    271        while (len > 0) {
    272            final int copyLen = Math.min(len, blockDataRemaining);
    273            System.arraycopy(data, off, blockData, currentIndex, copyLen);
    274            off += copyLen;
    275            blockDataRemaining -= copyLen;
    276            len -= copyLen;
    277            currentIndex += copyLen;
    278            if (blockDataRemaining == 0) {
    279                flushBlock();
    280                blockDataRemaining = blockData.length;
    281            }
    282        }
    283    }
    284
    285    @Override
    286    public void write(final int b) throws IOException {
    287        oneByte[0] = (byte) (b & 0xff);
    288        write(oneByte);
    289    }
    290
    291    private void writeFrameDescriptor() throws IOException {
    292        int flags = FramedLZ4CompressorInputStream.SUPPORTED_VERSION;
    293        if (!params.withBlockDependency) {
    294            flags |= FramedLZ4CompressorInputStream.BLOCK_INDEPENDENCE_MASK;
    295        }
    296        if (params.withContentChecksum) {
    297            flags |= FramedLZ4CompressorInputStream.CONTENT_CHECKSUM_MASK;
    298        }
    299        if (params.withBlockChecksum) {
    300            flags |= FramedLZ4CompressorInputStream.BLOCK_CHECKSUM_MASK;
    301        }
    302        out.write(flags);
    303        contentHash.update(flags);
    304        final int bd = params.blockSize.getIndex() << 4 & FramedLZ4CompressorInputStream.BLOCK_MAX_SIZE_MASK;
    305        out.write(bd);
    306        contentHash.update(bd);
    307        out.write((int) (contentHash.getValue() >> 8 & 0xff));
    308        contentHash.reset();
    309    }
    310
    311    private void writeTrailer() throws IOException {
    312        out.write(END_MARK);
    313        if (params.withContentChecksum) {
    314            ByteUtils.toLittleEndian(out, contentHash.getValue(), 4);
    315        }
    316    }
    317
    318}