    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.utils;
    020
    021import java.io.FileOutputStream;
    022import java.io.IOException;
    023import java.io.OutputStream;
    024import java.nio.ByteBuffer;
    025import java.nio.ByteOrder;
    026import java.nio.channels.ClosedChannelException;
    027import java.nio.channels.WritableByteChannel;
    028import java.util.concurrent.atomic.AtomicBoolean;
    029
    030/**
    031 * This class supports writing to an OutputStream or WritableByteChannel in fixed length blocks.
    032 * <p>It can be be used to support output to devices such as tape drives that require output in this
    033 * format. If the final block does not have enough content to fill an entire block, the output will
    034 * be padded to a full block size.</p>
    035 *
    036 * <p>This class can be used to support TAR,PAX, and CPIO blocked output to character special devices.
    037 * It is not recommended that this class be used unless writing to such devices, as the padding
    038 * serves no useful purpose in such cases.</p>
    039 *
    040 * <p>This class should normally wrap a FileOutputStream or associated WritableByteChannel directly.
    041 * If there is an intervening filter that modified the output, such as a CompressorOutputStream, or
    042 * performs its own buffering, such as BufferedOutputStream,  output to the device may
    043 * no longer be of the specified size.</p>
    044 *
    045 * <p>Any content written to this stream should be self-delimiting and should tolerate any padding
    046 * added to fill the last block.</p>
    047 *
    048 * @since 1.15
    049 */
    050public class FixedLengthBlockOutputStream extends OutputStream implements WritableByteChannel {
    051
    052    /**
    053     * Helper class to provide channel wrapper for arbitrary output stream that doesn't alter the
    054     * size of writes.  We can't use Channels.newChannel, because for non FileOutputStreams, it
    055     * breaks up writes into 8KB max chunks. Since the purpose of this class is to always write
    056     * complete blocks, we need to write a simple class to take care of it.
    057     */
    058    private static class BufferAtATimeOutputChannel implements WritableByteChannel {
    059
    060        private final OutputStream out;
    061        private final AtomicBoolean closed = new AtomicBoolean(false);
    062
    063        private BufferAtATimeOutputChannel(final OutputStream out) {
    064            this.out = out;
    065        }
    066
    067        @Override
    068        public void close() throws IOException {
    069            if (closed.compareAndSet(false, true)) {
    070                out.close();
    071            }
    072        }
    073
    074        @Override
    075        public boolean isOpen() {
    076            return !closed.get();
    077        }
    078
    079        @Override
    080        public int write(final ByteBuffer buffer) throws IOException {
    081            if (!isOpen()) {
    082                throw new ClosedChannelException();
    083            }
    084            if (!buffer.hasArray()) {
    085                throw new IOException("Direct buffer somehow written to BufferAtATimeOutputChannel");
    086            }
    087
    088            try {
    089                final int pos = buffer.position();
    090                final int len = buffer.limit() - pos;
    091                out.write(buffer.array(), buffer.arrayOffset() + pos, len);
    092                buffer.position(buffer.limit());
    093                return len;
    094            } catch (final IOException e) {
    095                try {
    096                    close();
    097                } catch (final IOException ignored) { //NOSONAR
    098                }
    099                throw e;
    100            }
    101        }
    102
    103    }
    104    private final WritableByteChannel out;
    105    private final int blockSize;
    106    private final ByteBuffer buffer;
    107
    108    private final AtomicBoolean closed = new AtomicBoolean(false);
    109     /**
    110     * Create a fixed length block output stream with given destination stream and block size
    111     * @param os   The stream to wrap.
    112     * @param blockSize The block size to use.
    113     */
    114    public FixedLengthBlockOutputStream(final OutputStream os, final int blockSize) {
    115        if (os instanceof FileOutputStream) {
    116            final FileOutputStream fileOutputStream = (FileOutputStream) os;
    117            out = fileOutputStream.getChannel();
    118            buffer = ByteBuffer.allocateDirect(blockSize);
    119        } else {
    120            out = new BufferAtATimeOutputChannel(os);
    121            buffer = ByteBuffer.allocate(blockSize);
    122        }
    123        this.blockSize = blockSize;
    124    }
    125
    126    /**
    127      * Create a fixed length block output stream with given destination writable byte channel and block size
    128     * @param out   The writable byte channel to wrap.
    129     * @param blockSize The block size to use.
    130     */
    131    public FixedLengthBlockOutputStream(final WritableByteChannel out, final int blockSize) {
    132        this.out = out;
    133        this.blockSize = blockSize;
    134        this.buffer = ByteBuffer.allocateDirect(blockSize);
    135    }
    136
    137    @Override
    138    public void close() throws IOException {
    139        if (closed.compareAndSet(false, true)) {
    140            try {
    141                flushBlock();
    142            } finally {
    143                out.close();
    144            }
    145        }
    146    }
    147
    148    /**
    149     * Potentially pads and then writes the current block to the underlying stream.
    150     * @throws IOException if writing fails
    151     */
    152    public void flushBlock() throws IOException {
    153        if (buffer.position() != 0) {
    154            padBlock();
    155            writeBlock();
    156        }
    157    }
    158
    159    @Override
    160    public boolean isOpen() {
    161        if (!out.isOpen()) {
    162            closed.set(true);
    163        }
    164        return !closed.get();
    165    }
    166
    167    private void maybeFlush() throws IOException {
    168        if (!buffer.hasRemaining()) {
    169            writeBlock();
    170        }
    171    }
    172
    173    private void padBlock() {
    174        buffer.order(ByteOrder.nativeOrder());
    175        int bytesToWrite = buffer.remaining();
    176        if (bytesToWrite > 8) {
    177            final int align = buffer.position() & 7;
    178            if (align != 0) {
    179                final int limit = 8 - align;
    180                for (int i = 0; i < limit; i++) {
    181                    buffer.put((byte) 0);
    182                }
    183                bytesToWrite -= limit;
    184            }
    185
    186            while (bytesToWrite >= 8) {
    187                buffer.putLong(0L);
    188                bytesToWrite -= 8;
    189            }
    190        }
    191        while (buffer.hasRemaining()) {
    192            buffer.put((byte) 0);
    193        }
    194    }
    195
    196    @Override
    197    public void write(final byte[] b, final int offset, final int length) throws IOException {
    198        if (!isOpen()) {
    199            throw new ClosedChannelException();
    200        }
    201        int off = offset;
    202        int len = length;
    203        while (len > 0) {
    204            final int n = Math.min(len, buffer.remaining());
    205            buffer.put(b, off, n);
    206            maybeFlush();
    207            len -= n;
    208            off += n;
    209        }
    210    }
    211
    212    @Override
    213    public int write(final ByteBuffer src) throws IOException {
    214        if (!isOpen()) {
    215            throw new ClosedChannelException();
    216        }
    217        final int srcRemaining = src.remaining();
    218
    219        if (srcRemaining < buffer.remaining()) {
    220            // if don't have enough bytes in src to fill up a block we must buffer
    221            buffer.put(src);
    222        } else {
    223            int srcLeft = srcRemaining;
    224            final int savedLimit = src.limit();
    225            // If we're not at the start of buffer, we have some bytes already  buffered
    226            // fill up the reset of buffer and write the block.
    227            if (buffer.position() != 0) {
    228                final int n = buffer.remaining();
    229                src.limit(src.position() + n);
    230                buffer.put(src);
    231                writeBlock();
    232                srcLeft -= n;
    233            }
    234            // whilst we have enough bytes in src for complete blocks,
    235            // write them directly from src without copying them to buffer
    236            while (srcLeft >= blockSize) {
    237                src.limit(src.position() + blockSize);
    238                out.write(src);
    239                srcLeft -= blockSize;
    240            }
    241            // copy any remaining bytes into buffer
    242            src.limit(savedLimit);
    243            buffer.put(src);
    244        }
    245        return srcRemaining;
    246    }
    247
    248    @Override
    249    public void write(final int b) throws IOException {
    250        if (!isOpen()) {
    251            throw new ClosedChannelException();
    252        }
    253        buffer.put((byte) b);
    254        maybeFlush();
    255    }
    256
    257    private void writeBlock() throws IOException {
    258        buffer.flip();
    259        final int i = out.write(buffer);
    260        final boolean hasRemaining = buffer.hasRemaining();
    261        if (i != blockSize || hasRemaining) {
    262            final String msg = String
    263                .format("Failed to write %,d bytes atomically. Only wrote  %,d",
    264                    blockSize, i);
    265            throw new IOException(msg);
    266        }
    267        buffer.clear();
    268    }
    269
    270
    271}