diff --git a/src/main/java/com/github/luben/zstd/Zstd.java b/src/main/java/com/github/luben/zstd/Zstd.java index 1a7c81e..0954cdb 100644 --- a/src/main/java/com/github/luben/zstd/Zstd.java +++ b/src/main/java/com/github/luben/zstd/Zstd.java @@ -1531,10 +1531,15 @@ public static ByteBuffer decompress(ByteBuffer srcBuff, ZstdDictDecompress dict, } } - static final byte[] extractArray(ByteBuffer buffer) { + static ByteBuffer getArrayBackedBuffer(BufferPool bufferPool, int size) throws ZstdIOException { + ByteBuffer buffer = bufferPool.get(size); + if (buffer == null) { + throw new ZstdIOException(Zstd.errMemoryAllocation(), "Cannot get ByteBuffer of size " + size + " from the BufferPool"); + } if (!buffer.hasArray() || buffer.arrayOffset() != 0) { + bufferPool.release(buffer); throw new IllegalArgumentException("provided ByteBuffer lacks array or has non-zero arrayOffset"); } - return buffer.array(); + return buffer; } } diff --git a/src/main/java/com/github/luben/zstd/ZstdInputStreamNoFinalizer.java b/src/main/java/com/github/luben/zstd/ZstdInputStreamNoFinalizer.java index d40076b..c6061a5 100644 --- a/src/main/java/com/github/luben/zstd/ZstdInputStreamNoFinalizer.java +++ b/src/main/java/com/github/luben/zstd/ZstdInputStreamNoFinalizer.java @@ -62,11 +62,8 @@ public ZstdInputStreamNoFinalizer(InputStream inStream) throws IOException { public ZstdInputStreamNoFinalizer(InputStream inStream, BufferPool bufferPool) throws IOException { super(inStream); this.bufferPool = bufferPool; - this.srcByteBuffer = bufferPool.get(srcBuffSize); - if (this.srcByteBuffer == null) { - throw new ZstdIOException(Zstd.errMemoryAllocation(), "Cannot get ByteBuffer of size " + srcBuffSize + " from the BufferPool"); - } - this.src = Zstd.extractArray(srcByteBuffer); + this.srcByteBuffer = Zstd.getArrayBackedBuffer(bufferPool, srcBuffSize); + this.src = srcByteBuffer.array(); // memory barrier synchronized(this) { this.stream = createDStream(); @@ -244,10 +241,10 @@ public synchronized long skip(long numBytes) throws IOException { if (bufferLen > numBytes) { bufferLen = (int) numBytes; } - ByteBuffer buf = bufferPool.get(bufferLen); + ByteBuffer buf = Zstd.getArrayBackedBuffer(bufferPool, bufferLen); long toSkip = numBytes; try { - byte data[] = Zstd.extractArray(buf); + byte data[] = buf.array(); while (toSkip > 0) { int read = read(data, 0, (int) Math.min((long) bufferLen, toSkip)); if (read < 0) { diff --git a/src/main/java/com/github/luben/zstd/ZstdOutputStreamNoFinalizer.java b/src/main/java/com/github/luben/zstd/ZstdOutputStreamNoFinalizer.java index 7602ea4..a176a97 100644 --- a/src/main/java/com/github/luben/zstd/ZstdOutputStreamNoFinalizer.java +++ b/src/main/java/com/github/luben/zstd/ZstdOutputStreamNoFinalizer.java @@ -80,11 +80,8 @@ public ZstdOutputStreamNoFinalizer(OutputStream outStream, BufferPool bufferPool // create compression context this.stream = createCStream(); this.bufferPool = bufferPool; - this.dstByteBuffer = bufferPool.get(dstSize); - if (this.dstByteBuffer == null) { - throw new ZstdIOException(Zstd.errMemoryAllocation(), "Cannot get ByteBuffer of size " + dstSize + " from the BufferPool"); - } - this.dst = Zstd.extractArray(dstByteBuffer); + this.dstByteBuffer = Zstd.getArrayBackedBuffer(bufferPool, dstSize); + this.dst = dstByteBuffer.array(); } /**