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}