Skip to content

Commit

Permalink
ZstdInputStreamNoFinalizer / ZstdOutputStreamNoFinalizer may leak buf…
Browse files Browse the repository at this point in the history
…fers when these are not arrray based.

Motivation:

We need to put the buffers back in the pool in case of validation failure as otherwise the buffer will leak.

Modifications:

Handle validation and buffer lifecycle via a static helper method

Result:

No more leak possible
  • Loading branch information
normanmaurer authored and luben committed Mar 22, 2024
1 parent 20014d1 commit 26ad7ed
Show file tree
Hide file tree
Showing 3 changed files with 13 additions and 14 deletions.
9 changes: 7 additions & 2 deletions src/main/java/com/github/luben/zstd/Zstd.java
Original file line number Diff line number Diff line change
Expand Up @@ -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;
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -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();
Expand Down Expand Up @@ -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) {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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();
}

/**
Expand Down

0 comments on commit 26ad7ed

Please sign in to comment.