Skip to content

Commit 49e01b4

Browse files
dfa1claude
dfa1
and
claude
authored
fix: free native context if a stream constructor fails after creating it (#13)
* fix: free native context if a stream constructor fails after creating it The four streaming wrappers create a native ZSTD_C/DCtx and then make a further native call (set parameter / load dictionary). If that second call threw, the freshly created context was leaked: the java.io streams closed only the arena (not the cctx/dctx), and the segment drivers rethrew from create() before super() could take ownership. Free the context on every constructor error path, and move the java.io streams' post-setup buffer allocations inside the guarded block so a failure there frees the context and the arena too. Found during a memory-leak / security review of the zstd module. The rest of the module checked out: NativeObject's atomic-swap close is exactly-once and idempotent; the ctx/dict creators allocate the native handle as their last step; copy helpers guard with toIntExact and requireNative; decode entry points validate bounds and native-ness. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> * refactor: unify the stream cleanup paths Segment drivers now take ownership first (super(createCtx())) and run setup in the constructor body, so a setup failure is released by the existing close() — drops the duplicated free helper. The java.io streams share one freeCctx/freeDctx best-effort helper between the constructor error path and close(). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> --------- Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
1 parent 21a9ab4 commit 49e01b4

4 files changed

Lines changed: 91 additions & 47 deletions

File tree

‎zstd/src/main/java/io/github/dfa1/zstd/ZstdCompressStream.java‎

Lines changed: 25 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -41,39 +41,50 @@ public ZstdCompressStream() {
4141
///
4242
/// @param level the compression level
4343
public ZstdCompressStream(int level) {
44-
super(create(level, null));
44+
this(level, null);
4545
}
4646

4747
/// Creates a streaming compressor at `level` using `dictionary`.
4848
///
4949
/// @param level the compression level
50-
/// @param dictionary the dictionary to compress against
50+
/// @param dictionary the dictionary to compress against, or `null` for none
5151
public ZstdCompressStream(int level, ZstdDictionary dictionary) {
52-
super(create(level, dictionary));
52+
// Own the context first, so any failure setting it up is cleaned up by
53+
// close() — one release path, no leak on a half-built stream.
54+
super(createCctx());
55+
try {
56+
Zstd.call(() -> (long) Bindings.CCTX_SET_PARAMETER.invokeExact(
57+
ptr(), ZstdCompressParameter.COMPRESSION_LEVEL.value(), level));
58+
if (dictionary != null) {
59+
loadDictionary(dictionary);
60+
}
61+
} catch (Throwable t) {
62+
close();
63+
throw rethrow(t);
64+
}
5365
}
5466

55-
private static MemorySegment create(int level, ZstdDictionary dictionary) {
67+
private static MemorySegment createCctx() {
5668
try {
5769
MemorySegment cctx = (MemorySegment) Bindings.CREATE_CCTX.invokeExact();
5870
if (MemorySegment.NULL.equals(cctx)) {
5971
throw new ZstdException("ZSTD_createCCtx returned NULL");
6072
}
61-
Zstd.call(() -> (long) Bindings.CCTX_SET_PARAMETER.invokeExact(
62-
cctx, ZstdCompressParameter.COMPRESSION_LEVEL.value(), level));
63-
if (dictionary != null) {
64-
try (Arena staging = Arena.ofConfined()) {
65-
byte[] raw = dictionary.raw();
66-
MemorySegment d = Zstd.copyIn(staging, raw);
67-
Zstd.call(() -> (long) Bindings.CCTX_LOAD_DICTIONARY.invokeExact(
68-
cctx, d, (long) raw.length));
69-
}
70-
}
7173
return cctx;
7274
} catch (Throwable t) {
7375
throw rethrow(t);
7476
}
7577
}
7678

79+
private void loadDictionary(ZstdDictionary dictionary) {
80+
try (Arena staging = Arena.ofConfined()) {
81+
byte[] raw = dictionary.raw();
82+
MemorySegment d = Zstd.copyIn(staging, raw);
83+
Zstd.call(() -> (long) Bindings.CCTX_LOAD_DICTIONARY.invokeExact(
84+
ptr(), d, (long) raw.length));
85+
}
86+
}
87+
7788
/// Compresses as much of `src` as fits into `dst` in one step.
7889
///
7990
/// Advance the source by [ZstdStreamResult#bytesConsumed()] and write out

‎zstd/src/main/java/io/github/dfa1/zstd/ZstdDecompressStream.java‎

Lines changed: 23 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -18,36 +18,47 @@ public final class ZstdDecompressStream extends NativeObject {
1818

1919
/// Creates a streaming decompressor.
2020
public ZstdDecompressStream() {
21-
super(create(null));
21+
this(null);
2222
}
2323

2424
/// Creates a streaming decompressor for frames built with `dictionary`.
2525
///
26-
/// @param dictionary the dictionary the frames were compressed against
26+
/// @param dictionary the dictionary the frames were compressed against, or `null` for none
2727
public ZstdDecompressStream(ZstdDictionary dictionary) {
28-
super(create(dictionary));
28+
// Own the context first, so any failure setting it up is cleaned up by
29+
// close() — one release path, no leak on a half-built stream.
30+
super(createDctx());
31+
try {
32+
if (dictionary != null) {
33+
loadDictionary(dictionary);
34+
}
35+
} catch (Throwable t) {
36+
close();
37+
throw rethrow(t);
38+
}
2939
}
3040

31-
private static MemorySegment create(ZstdDictionary dictionary) {
41+
private static MemorySegment createDctx() {
3242
try {
3343
MemorySegment dctx = (MemorySegment) Bindings.CREATE_DCTX.invokeExact();
3444
if (MemorySegment.NULL.equals(dctx)) {
3545
throw new ZstdException("ZSTD_createDCtx returned NULL");
3646
}
37-
if (dictionary != null) {
38-
try (Arena staging = Arena.ofConfined()) {
39-
byte[] raw = dictionary.raw();
40-
MemorySegment d = Zstd.copyIn(staging, raw);
41-
Zstd.call(() -> (long) Bindings.DCTX_LOAD_DICTIONARY.invokeExact(
42-
dctx, d, (long) raw.length));
43-
}
44-
}
4547
return dctx;
4648
} catch (Throwable t) {
4749
throw rethrow(t);
4850
}
4951
}
5052

53+
private void loadDictionary(ZstdDictionary dictionary) {
54+
try (Arena staging = Arena.ofConfined()) {
55+
byte[] raw = dictionary.raw();
56+
MemorySegment d = Zstd.copyIn(staging, raw);
57+
Zstd.call(() -> (long) Bindings.DCTX_LOAD_DICTIONARY.invokeExact(
58+
ptr(), d, (long) raw.length));
59+
}
60+
}
61+
5162
/// Decompresses as much of `src` as fits into `dst` in one step.
5263
///
5364
/// Advance the source by [ZstdStreamResult#bytesConsumed()] and read out

‎zstd/src/main/java/io/github/dfa1/zstd/ZstdInputStream.java‎

Lines changed: 22 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -57,24 +57,39 @@ public ZstdInputStream(InputStream in) {
5757
/// @param dictionary the dictionary the frame was compressed against, or `null` for none
5858
public ZstdInputStream(InputStream in, ZstdDictionary dictionary) {
5959
this.in = in;
60+
MemorySegment d = null;
6061
try {
61-
this.dctx = (MemorySegment) Bindings.CREATE_DCTX.invokeExact();
62-
if (MemorySegment.NULL.equals(dctx)) {
62+
d = (MemorySegment) Bindings.CREATE_DCTX.invokeExact();
63+
if (MemorySegment.NULL.equals(d)) {
6364
throw new ZstdException("ZSTD_createDCtx returned NULL");
6465
}
66+
this.dctx = d;
6567
if (dictionary != null) {
6668
loadDictionary(dictionary);
6769
}
6870
this.inCap = (long) Bindings.DSTREAM_IN_SIZE.invokeExact();
6971
this.outCap = (long) Bindings.DSTREAM_OUT_SIZE.invokeExact();
72+
this.inSeg = arena.allocate(inCap);
73+
this.outSeg = arena.allocate(outCap);
74+
this.feed = new byte[Math.toIntExact(inCap)];
75+
this.hold = new byte[Math.toIntExact(outCap)];
7076
} catch (Throwable t) {
77+
// Free the context if it was created, then the arena, so a failed
78+
// constructor leaks neither the native dctx nor the arena buffers.
79+
if (d != null && !MemorySegment.NULL.equals(d)) {
80+
freeDctx(d);
81+
}
7182
arena.close();
7283
throw rethrow(t);
7384
}
74-
this.inSeg = arena.allocate(inCap);
75-
this.outSeg = arena.allocate(outCap);
76-
this.feed = new byte[Math.toIntExact(inCap)];
77-
this.hold = new byte[Math.toIntExact(outCap)];
85+
}
86+
87+
private static void freeDctx(MemorySegment dctx) {
88+
try {
89+
var _ = (long) Bindings.FREE_DCTX.invokeExact(dctx);
90+
} catch (Throwable _) {
91+
// best-effort free
92+
}
7893
}
7994

8095
private void loadDictionary(ZstdDictionary dictionary) {
@@ -158,11 +173,7 @@ public void close() throws IOException {
158173
return;
159174
}
160175
closed = true;
161-
try {
162-
var _ = (long) Bindings.FREE_DCTX.invokeExact(dctx);
163-
} catch (Throwable _) {
164-
// best-effort free
165-
}
176+
freeDctx(dctx);
166177
arena.close();
167178
in.close();
168179
}

‎zstd/src/main/java/io/github/dfa1/zstd/ZstdOutputStream.java‎

Lines changed: 21 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -93,25 +93,40 @@ private void setPledgedSrcSize(long pledgedSrcSize) {
9393
/// @param dictionary the dictionary to compress against, or `null` for none
9494
public ZstdOutputStream(OutputStream out, int level, ZstdDictionary dictionary) {
9595
this.out = out;
96+
MemorySegment c = null;
9697
try {
97-
this.cctx = (MemorySegment) Bindings.CREATE_CCTX.invokeExact();
98-
if (MemorySegment.NULL.equals(cctx)) {
98+
c = (MemorySegment) Bindings.CREATE_CCTX.invokeExact();
99+
if (MemorySegment.NULL.equals(c)) {
99100
throw new ZstdException("ZSTD_createCCtx returned NULL");
100101
}
102+
this.cctx = c;
101103
Zstd.call(() -> (long) Bindings.CCTX_SET_PARAMETER.invokeExact(
102104
cctx, ZSTD_C_COMPRESSION_LEVEL, level));
103105
if (dictionary != null) {
104106
loadDictionary(dictionary);
105107
}
106108
this.inCap = (long) Bindings.CSTREAM_IN_SIZE.invokeExact();
107109
this.outCap = (long) Bindings.CSTREAM_OUT_SIZE.invokeExact();
110+
this.inSeg = arena.allocate(inCap);
111+
this.outSeg = arena.allocate(outCap);
112+
this.drain = new byte[Math.toIntExact(outCap)];
108113
} catch (Throwable t) {
114+
// Free the context if it was created, then the arena, so a failed
115+
// constructor leaks neither the native cctx nor the arena buffers.
116+
if (c != null && !MemorySegment.NULL.equals(c)) {
117+
freeCctx(c);
118+
}
109119
arena.close();
110120
throw rethrow(t);
111121
}
112-
this.inSeg = arena.allocate(inCap);
113-
this.outSeg = arena.allocate(outCap);
114-
this.drain = new byte[Math.toIntExact(outCap)];
122+
}
123+
124+
private static void freeCctx(MemorySegment cctx) {
125+
try {
126+
var _ = (long) Bindings.FREE_CCTX.invokeExact(cctx);
127+
} catch (Throwable _) {
128+
// best-effort free
129+
}
115130
}
116131

117132
private void loadDictionary(ZstdDictionary dictionary) {
@@ -172,11 +187,7 @@ public void close() throws IOException {
172187
out.flush();
173188
} finally {
174189
closed = true;
175-
try {
176-
var _ = (long) Bindings.FREE_CCTX.invokeExact(cctx);
177-
} catch (Throwable _) {
178-
// best-effort free
179-
}
190+
freeCctx(cctx);
180191
arena.close();
181192
out.close();
182193
}

0 commit comments

Comments
 (0)