001/**
002 * Java Web Archive Toolkit - Software to read and validate ARC, WARC
003 * and GZip files. (http://jwat.org/)
004 * Copyright 2011-2012 Netarkivet.dk (http://netarkivet.dk/)
005 *
006 * Licensed under the Apache License, Version 2.0 (the "License");
007 * you may not use this file except in compliance with the License.
008 * 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, software
013 * distributed under the License is distributed on an "AS IS" BASIS,
014 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
015 * See the License for the specific language governing permissions and
016 * limitations under the License.
017 */
018package org.jwat.common;
019
020import java.io.IOException;
021import java.io.InputStream;
022import java.io.OutputStream;
023import java.nio.ByteBuffer;
024import java.util.concurrent.Semaphore;
025
026/**
027 * Simple class to allow the use of an in-memory buffer as an alternate input or output stream.
028 *
029 * @author nicl
030 */
031public class ByteArrayIOStream {
032
033    /** Default buffer size if none specified. */
034    public static final int DEFAULT_BUFFER_SIZE = 10*1024*1024;
035
036    /** Buffer lock. */
037    protected Semaphore lock = new Semaphore(1);
038
039    /** Internal byte array. */
040    protected byte[] bytes;
041
042    /** Internal <code>ByteBuffer</code> wrapper of internal byte array. */
043    protected ByteBuffer byteBuffer;
044
045    /** Current <code>ByteBuffer</code> limit. */
046    protected int limit = 0;
047
048    /**
049     * Construct object with default buffer size.
050     */
051    public ByteArrayIOStream() {
052        this(DEFAULT_BUFFER_SIZE);
053    }
054
055    /**
056     * Construct object with supplied buffer size.
057     * @param bufferSize buffer size to use
058     */
059    public ByteArrayIOStream(int bufferSize) {
060        bytes = new byte[bufferSize];
061        byteBuffer = ByteBuffer.wrap(bytes);
062    }
063
064    /**
065     * Return <code>OutputStream</code> and lock this buffer
066     * @return <code>OutputStream</code> wrapper of this buffer
067     */
068    public OutputStream getOutputStream() {
069        if (!lock.tryAcquire()) {
070            throw new IllegalStateException();
071        }
072        byteBuffer.clear();
073        limit = 0;
074        return new OutputStreamImpl(this);
075    }
076
077    /**
078     * Return internal byte array.
079     * @return internal byte array
080     */
081    public byte[] getBytes() {
082        return bytes;
083    }
084
085    /**
086     * Return allocated length of internal buffer.
087     * @return allocated length of internal buffer
088     */
089    public int getLength() {
090        return bytes.length;
091    }
092
093    /**
094     * Return internal <code>ByteBuffer</code> limit.
095     * @return internal <code>ByteBuffer</code> limit
096     */
097    public int getLimit() {
098        return limit;
099    }
100
101    /**
102     * Return internal byte array wrapped in a new <code>ByteBuffer</code> object.
103     * @return internal byte array wrapped in a new <code>ByteBuffer</code> object
104     */
105    public ByteBuffer getByteBuffer() {
106        ByteBuffer buffer = ByteBuffer.wrap(bytes);
107        buffer.position(0);
108        buffer.limit(limit);
109        return buffer;
110    }
111
112    /**
113     * Return <code>InputStream</code> and lock this buffer
114     * @return <code>InputStream</code> wrapper of this buffer
115     */
116    public InputStream getInputStream() {
117        if (!lock.tryAcquire()) {
118            throw new IllegalStateException();
119        }
120        byteBuffer.clear();
121        byteBuffer.limit(limit);
122        return new InputStreamImpl(this);
123    }
124
125    /**
126     * Release lock, called when input/output stream is closed.
127     */
128    protected void release() {
129        lock = null;
130        byteBuffer = null;
131        bytes = null;
132        limit = 0;
133    }
134
135    /**
136     * Byte array OutputStream.
137     */
138    public static class OutputStreamImpl extends OutputStream {
139
140        /** Owning <code>ByteArrayIOStream</code>. */
141        protected ByteArrayIOStream baios;
142        /** <code>ByteArrayIOStream</code>'s <code>ByteBuffer</code>. */
143        protected ByteBuffer byteBuffer;
144
145        /**
146         * Construct <code>OutputStream</code> object.
147         * @param baios owning <code>ByteArrayIOStream</code>
148         */
149        protected OutputStreamImpl(ByteArrayIOStream baios) {
150            this.baios = baios;
151            this.byteBuffer = baios.byteBuffer;
152        }
153
154        @Override
155        public void close() {
156            if (baios != null) {
157                baios.limit = baios.byteBuffer.position();
158                baios.lock.release();
159                baios = null;
160                byteBuffer = null;
161            }
162        }
163
164        @Override
165        public void flush() throws IOException {
166        }
167
168        @Override
169        public void write(byte[] b, int off, int len) throws IOException {
170            byteBuffer.put(b, off, len);
171        }
172
173        @Override
174        public void write(byte[] b) throws IOException {
175            byteBuffer.put(b);
176        }
177
178        @Override
179        public void write(int b) throws IOException {
180            byteBuffer.put((byte)b);
181        }
182
183    }
184
185    /**
186     * Byte array InputStream.
187     */
188    public static class InputStreamImpl extends InputStream {
189
190        /** Owning <code>ByteArrayIOStream</code>. */
191        protected ByteArrayIOStream baios;
192        /** <code>ByteArrayIOStream</code>'s <code>ByteBuffer</code>. */
193        protected ByteBuffer byteBuffer;
194
195        /**
196         * Construct <code>InputStream</code> object.
197         * @param baios owning <code>ByteArrayIOStream</code>
198         */
199        protected InputStreamImpl(ByteArrayIOStream baios) {
200            this.baios = baios;
201            this.byteBuffer = baios.byteBuffer;
202        }
203
204        @Override
205        public void close() {
206            if (baios != null) {
207                baios.lock.release();
208                baios = null;
209                byteBuffer = null;
210            }
211        }
212
213        @Override
214        public int available() throws IOException {
215            return byteBuffer.limit() - byteBuffer.position();
216        }
217
218        @Override
219        public boolean markSupported() {
220            return false;
221        }
222
223        @Override
224        public synchronized void mark(int readlimit) {
225            throw new UnsupportedOperationException();
226        }
227
228        @Override
229        public synchronized void reset() throws IOException {
230            throw new UnsupportedOperationException();
231        }
232
233        @Override
234        public int read() throws IOException {
235            if (byteBuffer.remaining() > 0) {
236                return byteBuffer.get();
237            } else {
238                return -1;
239            }
240        }
241
242        @Override
243        public int read(byte[] b, int off, int len) throws IOException {
244            if (len == 0) {
245                return 0;
246            }
247            int remaining = byteBuffer.remaining();
248            if (len > remaining) {
249                len = remaining;
250            }
251            if (len > 0) {
252                byteBuffer.get(b, off, len);
253                return len;
254            } else {
255                return -1;
256            }
257        }
258
259        @Override
260        public int read(byte[] b) throws IOException {
261            int len = b.length;
262            if (len == 0) {
263                return 0;
264            }
265            int remaining = byteBuffer.remaining();
266            if (len > remaining) {
267                len = remaining;
268            }
269            if (len > 0) {
270                byteBuffer.get(b, 0, len);
271                return len;
272            } else {
273                return -1;
274            }
275        }
276
277        @Override
278        public long skip(long n) throws IOException {
279            int remaining = byteBuffer.remaining();
280            if (n > remaining) {
281                n = remaining;
282            }
283            if (n > 0) {
284                byteBuffer.position(byteBuffer.position() + (int)n);
285            }
286            return n;
287        }
288
289    }
290
291}