Skip to content

Commit 30ce27f

Browse files
bluestreak01claude
bluestreak01
and
claude
committed
Harden QWP client: volatile ioThread, decoder leak fixes
- Mark ioThread volatile in QwpQueryClient so cancel() from a thread other than the one that ran connect() observes the published reference, and a concurrent null-out from close() does not race. - In QwpResultBatchDecoder, zero decompressScratchAddr and decompressScratchCapacity between free and malloc so a throwing malloc cannot leave a dangling address paired with a non-zero capacity. The next decode would otherwise skip the first-alloc branch and trigger a use-after-free. - Wrap QwpResultBatchDecoderHardeningTest cases in assertMemoryLeak and close the decoder in finally to catch native-memory leaks in the decoder itself. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
1 parent f42254b commit 30ce27f

3 files changed

Lines changed: 144 additions & 118 deletions

File tree

‎core/src/main/java/io/questdb/client/cutlass/qwp/client/QwpQueryClient.java‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -79,7 +79,11 @@ public class QwpQueryClient implements QuietCloseable {
7979
// payload before it parks, and the client auto-replenishes by the size of
8080
// each batch as the user releases it.
8181
private long initialCreditBytes;
82-
private QwpEgressIoThread ioThread;
82+
// Volatile so a cancel() call from a thread other than the one that ran
83+
// connect() sees the published reference (and a concurrent null-out from
84+
// close() is observed without a stale-reference race). The thread-safety
85+
// contract documented on cancel() relies on this.
86+
private volatile QwpEgressIoThread ioThread;
8387
private Thread ioThreadHandle;
8488
private boolean lastCloseTimedOut;
8589
private int negotiatedQwpVersion;

‎core/src/main/java/io/questdb/client/cutlass/qwp/client/QwpResultBatchDecoder.java‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -295,9 +295,14 @@ private void decodePayload(QwpBatchBuffer buffer, long payload, int payloadLen)
295295
long decLen = Zstd.decompress(dctx, p, srcLen, decompressScratchAddr, decompressScratchCapacity);
296296
while (decLen < 0 && decompressScratchCapacity < MAX_SCRATCH) {
297297
int newCap = Math.min(decompressScratchCapacity * 2, MAX_SCRATCH);
298+
// Reset to 0 before free + malloc so a throwing malloc cannot leave
299+
// a dangling address + non-zero capacity behind; the next decode
300+
// would otherwise skip the first-alloc branch and use-after-free.
298301
Unsafe.free(decompressScratchAddr, decompressScratchCapacity, MemoryTag.NATIVE_DEFAULT);
302+
decompressScratchAddr = 0;
303+
decompressScratchCapacity = 0;
304+
decompressScratchAddr = Unsafe.malloc(newCap, MemoryTag.NATIVE_DEFAULT);
299305
decompressScratchCapacity = newCap;
300-
decompressScratchAddr = Unsafe.malloc(decompressScratchCapacity, MemoryTag.NATIVE_DEFAULT);
301306
decLen = Zstd.decompress(dctx, p, srcLen, decompressScratchAddr, decompressScratchCapacity);
302307
}
303308
if (decLen < 0) {

‎core/src/test/java/io/questdb/client/test/cutlass/qwp/client/QwpResultBatchDecoderHardeningTest.java‎

Lines changed: 133 additions & 116 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,7 @@
3232
import io.questdb.client.cutlass.qwp.protocol.QwpConstants;
3333
import io.questdb.client.std.MemoryTag;
3434
import io.questdb.client.std.Unsafe;
35+
import io.questdb.client.test.tools.TestUtils;
3536
import org.junit.Assert;
3637
import org.junit.Test;
3738

@@ -51,25 +52,28 @@ public class QwpResultBatchDecoderHardeningTest {
5152
* append nulls until OOM (or AIOOBE for negative ids cast from a high varint).
5253
*/
5354
@Test
54-
public void testHugeSchemaIdIsRejected() {
55-
QwpResultBatchDecoder decoder = new QwpResultBatchDecoder();
56-
QwpBatchBuffer buffer = new QwpBatchBuffer(256);
57-
long staging = Unsafe.malloc(256, MemoryTag.NATIVE_DEFAULT);
58-
try {
59-
// schema_id = 1_000_000_000, well above the 65_535 cap.
60-
int len = writeMinimalResultBatch(staging, /*schemaId=*/ 1_000_000_000L);
61-
buffer.copyFromPayload(staging, len);
55+
public void testHugeSchemaIdIsRejected() throws Exception {
56+
TestUtils.assertMemoryLeak(() -> {
57+
QwpResultBatchDecoder decoder = new QwpResultBatchDecoder();
58+
QwpBatchBuffer buffer = new QwpBatchBuffer(256);
59+
long staging = Unsafe.malloc(256, MemoryTag.NATIVE_DEFAULT);
6260
try {
63-
decoder.decode(buffer);
64-
Assert.fail("decoder must reject huge schema_id");
65-
} catch (QwpDecodeException expected) {
66-
Assert.assertTrue("error message should mention schema_id: " + expected.getMessage(),
67-
expected.getMessage().contains("schema_id"));
61+
// schema_id = 1_000_000_000, well above the 65_535 cap.
62+
int len = writeMinimalResultBatch(staging, /*schemaId=*/ 1_000_000_000L);
63+
buffer.copyFromPayload(staging, len);
64+
try {
65+
decoder.decode(buffer);
66+
Assert.fail("decoder must reject huge schema_id");
67+
} catch (QwpDecodeException expected) {
68+
Assert.assertTrue("error message should mention schema_id: " + expected.getMessage(),
69+
expected.getMessage().contains("schema_id"));
70+
}
71+
} finally {
72+
Unsafe.free(staging, 256, MemoryTag.NATIVE_DEFAULT);
73+
buffer.close();
74+
decoder.close();
6875
}
69-
} finally {
70-
Unsafe.free(staging, 256, MemoryTag.NATIVE_DEFAULT);
71-
buffer.close();
72-
}
76+
});
7377
}
7478

7579
/**
@@ -78,32 +82,35 @@ public void testHugeSchemaIdIsRejected() {
7882
* rejected, not silently passed to {@code getQuick(negativeIndex)}.
7983
*/
8084
@Test
81-
public void testNegativeSchemaIdIsRejected() {
82-
QwpResultBatchDecoder decoder = new QwpResultBatchDecoder();
83-
QwpBatchBuffer buffer = new QwpBatchBuffer(256);
84-
long staging = Unsafe.malloc(256, MemoryTag.NATIVE_DEFAULT);
85-
try {
86-
// 5-byte varint encoding 0x80000000 (which casts to Integer.MIN_VALUE).
87-
// varint bytes for 0x80000000:
88-
// value bits 7..0: 0x00 -> byte: 0x80 (continuation)
89-
// value bits 14..8: 0x00 -> byte: 0x80
90-
// value bits 21..15:0x00 -> byte: 0x80
91-
// value bits 28..22:0x00 -> byte: 0x80
92-
// value bits 35..29:0x08 -> byte: 0x08 (no continuation)
93-
int len = writeMinimalResultBatchWithRawSchemaIdVarint(
94-
staging, new byte[]{(byte) 0x80, (byte) 0x80, (byte) 0x80, (byte) 0x80, (byte) 0x08});
95-
buffer.copyFromPayload(staging, len);
85+
public void testNegativeSchemaIdIsRejected() throws Exception {
86+
TestUtils.assertMemoryLeak(() -> {
87+
QwpResultBatchDecoder decoder = new QwpResultBatchDecoder();
88+
QwpBatchBuffer buffer = new QwpBatchBuffer(256);
89+
long staging = Unsafe.malloc(256, MemoryTag.NATIVE_DEFAULT);
9690
try {
97-
decoder.decode(buffer);
98-
Assert.fail("decoder must reject huge/negative schema_id");
99-
} catch (QwpDecodeException expected) {
100-
Assert.assertTrue("error message should mention schema_id: " + expected.getMessage(),
101-
expected.getMessage().contains("schema_id"));
91+
// 5-byte varint encoding 0x80000000 (which casts to Integer.MIN_VALUE).
92+
// varint bytes for 0x80000000:
93+
// value bits 7..0: 0x00 -> byte: 0x80 (continuation)
94+
// value bits 14..8: 0x00 -> byte: 0x80
95+
// value bits 21..15:0x00 -> byte: 0x80
96+
// value bits 28..22:0x00 -> byte: 0x80
97+
// value bits 35..29:0x08 -> byte: 0x08 (no continuation)
98+
int len = writeMinimalResultBatchWithRawSchemaIdVarint(
99+
staging, new byte[]{(byte) 0x80, (byte) 0x80, (byte) 0x80, (byte) 0x80, (byte) 0x08});
100+
buffer.copyFromPayload(staging, len);
101+
try {
102+
decoder.decode(buffer);
103+
Assert.fail("decoder must reject huge/negative schema_id");
104+
} catch (QwpDecodeException expected) {
105+
Assert.assertTrue("error message should mention schema_id: " + expected.getMessage(),
106+
expected.getMessage().contains("schema_id"));
107+
}
108+
} finally {
109+
Unsafe.free(staging, 256, MemoryTag.NATIVE_DEFAULT);
110+
buffer.close();
111+
decoder.close();
102112
}
103-
} finally {
104-
Unsafe.free(staging, 256, MemoryTag.NATIVE_DEFAULT);
105-
buffer.close();
106-
}
113+
});
107114
}
108115

109116
/**
@@ -114,32 +121,34 @@ public void testNegativeSchemaIdIsRejected() {
114121
* detects the overrun and reports a bounded error.
115122
*/
116123
@Test
117-
public void testQueryErrorMsgLenOverrunIsRejected() {
118-
// Frame contents:
119-
// 12 bytes header (uninspected by decodeError)
120-
// 1 byte msg_kind
121-
// 8 bytes request_id
122-
// 1 byte status
123-
// 2 bytes msgLen (we set 0xFFFF)
124-
// 0 bytes of actual message body
125-
// Total payload: 24 bytes; msgLen would otherwise force reading 65535 bytes.
126-
int payloadLen = 12 + 1 + 8 + 1 + 2;
127-
long buf = Unsafe.malloc(payloadLen, MemoryTag.NATIVE_DEFAULT);
128-
try {
129-
// Zero out
130-
for (int i = 0; i < payloadLen; i++) Unsafe.getUnsafe().putByte(buf + i, (byte) 0);
131-
// Write an obviously bogus msgLen at the right offset (header + msg_kind + reqId + status).
132-
long msgLenOffset = buf + 12 + 1 + 8 + 1;
133-
Unsafe.getUnsafe().putShort(msgLenOffset, (short) 0xFFFF);
124+
public void testQueryErrorMsgLenOverrunIsRejected() throws Exception {
125+
TestUtils.assertMemoryLeak(() -> {
126+
// Frame contents:
127+
// 12 bytes header (uninspected by decodeError)
128+
// 1 byte msg_kind
129+
// 8 bytes request_id
130+
// 1 byte status
131+
// 2 bytes msgLen (we set 0xFFFF)
132+
// 0 bytes of actual message body
133+
// Total payload: 24 bytes; msgLen would otherwise force reading 65535 bytes.
134+
int payloadLen = 12 + 1 + 8 + 1 + 2;
135+
long buf = Unsafe.malloc(payloadLen, MemoryTag.NATIVE_DEFAULT);
136+
try {
137+
// Zero out
138+
for (int i = 0; i < payloadLen; i++) Unsafe.getUnsafe().putByte(buf + i, (byte) 0);
139+
// Write an obviously bogus msgLen at the right offset (header + msg_kind + reqId + status).
140+
long msgLenOffset = buf + 12 + 1 + 8 + 1;
141+
Unsafe.getUnsafe().putShort(msgLenOffset, (short) 0xFFFF);
134142

135-
QueryEvent ev = QwpEgressIoThread.decodeError(buf, payloadLen);
136-
Assert.assertEquals(QueryEvent.KIND_ERROR, ev.kind);
137-
Assert.assertNotNull(ev.errorMessage);
138-
Assert.assertTrue("error must mention msg_len overrun: " + ev.errorMessage,
139-
ev.errorMessage.contains("msg_len") && ev.errorMessage.contains("exceeds"));
140-
} finally {
141-
Unsafe.free(buf, payloadLen, MemoryTag.NATIVE_DEFAULT);
142-
}
143+
QueryEvent ev = QwpEgressIoThread.decodeError(buf, payloadLen);
144+
Assert.assertEquals(QueryEvent.KIND_ERROR, ev.kind);
145+
Assert.assertNotNull(ev.errorMessage);
146+
Assert.assertTrue("error must mention msg_len overrun: " + ev.errorMessage,
147+
ev.errorMessage.contains("msg_len") && ev.errorMessage.contains("exceeds"));
148+
} finally {
149+
Unsafe.free(buf, payloadLen, MemoryTag.NATIVE_DEFAULT);
150+
}
151+
});
143152
}
144153

145154
/**
@@ -148,28 +157,30 @@ public void testQueryErrorMsgLenOverrunIsRejected() {
148157
* confirming a real defensive guard, not a broken decoder.
149158
*/
150159
@Test
151-
public void testQueryErrorValidMessageDecodes() {
152-
byte[] msgBytes = "boom".getBytes(java.nio.charset.StandardCharsets.UTF_8);
153-
int payloadLen = 12 + 1 + 8 + 1 + 2 + msgBytes.length;
154-
long buf = Unsafe.malloc(payloadLen, MemoryTag.NATIVE_DEFAULT);
155-
try {
156-
for (int i = 0; i < payloadLen; i++) Unsafe.getUnsafe().putByte(buf + i, (byte) 0);
157-
long statusOffset = buf + 12 + 1 + 8;
158-
Unsafe.getUnsafe().putByte(statusOffset, (byte) 0x05);
159-
long msgLenOffset = statusOffset + 1;
160-
Unsafe.getUnsafe().putShort(msgLenOffset, (short) msgBytes.length);
161-
long bytesOffset = msgLenOffset + 2;
162-
for (int i = 0; i < msgBytes.length; i++) {
163-
Unsafe.getUnsafe().putByte(bytesOffset + i, msgBytes[i]);
164-
}
160+
public void testQueryErrorValidMessageDecodes() throws Exception {
161+
TestUtils.assertMemoryLeak(() -> {
162+
byte[] msgBytes = "boom".getBytes(java.nio.charset.StandardCharsets.UTF_8);
163+
int payloadLen = 12 + 1 + 8 + 1 + 2 + msgBytes.length;
164+
long buf = Unsafe.malloc(payloadLen, MemoryTag.NATIVE_DEFAULT);
165+
try {
166+
for (int i = 0; i < payloadLen; i++) Unsafe.getUnsafe().putByte(buf + i, (byte) 0);
167+
long statusOffset = buf + 12 + 1 + 8;
168+
Unsafe.getUnsafe().putByte(statusOffset, (byte) 0x05);
169+
long msgLenOffset = statusOffset + 1;
170+
Unsafe.getUnsafe().putShort(msgLenOffset, (short) msgBytes.length);
171+
long bytesOffset = msgLenOffset + 2;
172+
for (int i = 0; i < msgBytes.length; i++) {
173+
Unsafe.getUnsafe().putByte(bytesOffset + i, msgBytes[i]);
174+
}
165175

166-
QueryEvent ev = QwpEgressIoThread.decodeError(buf, payloadLen);
167-
Assert.assertEquals(QueryEvent.KIND_ERROR, ev.kind);
168-
Assert.assertEquals((byte) 0x05, ev.errorStatus);
169-
Assert.assertEquals("boom", ev.errorMessage);
170-
} finally {
171-
Unsafe.free(buf, payloadLen, MemoryTag.NATIVE_DEFAULT);
172-
}
176+
QueryEvent ev = QwpEgressIoThread.decodeError(buf, payloadLen);
177+
Assert.assertEquals(QueryEvent.KIND_ERROR, ev.kind);
178+
Assert.assertEquals((byte) 0x05, ev.errorStatus);
179+
Assert.assertEquals("boom", ev.errorMessage);
180+
} finally {
181+
Unsafe.free(buf, payloadLen, MemoryTag.NATIVE_DEFAULT);
182+
}
183+
});
173184
}
174185

175186
/**
@@ -180,24 +191,27 @@ public void testQueryErrorValidMessageDecodes() {
180191
* memory backwards.
181192
*/
182193
@Test
183-
public void testStringColumnNegativeTotalBytesIsRejected() {
184-
QwpResultBatchDecoder decoder = new QwpResultBatchDecoder();
185-
QwpBatchBuffer buffer = new QwpBatchBuffer(512);
186-
long staging = Unsafe.malloc(512, MemoryTag.NATIVE_DEFAULT);
187-
try {
188-
int len = writeStringResultBatch(staging, /*nonNull=*/ 1, /*totalBytes=*/ -1);
189-
buffer.copyFromPayload(staging, len);
194+
public void testStringColumnNegativeTotalBytesIsRejected() throws Exception {
195+
TestUtils.assertMemoryLeak(() -> {
196+
QwpResultBatchDecoder decoder = new QwpResultBatchDecoder();
197+
QwpBatchBuffer buffer = new QwpBatchBuffer(512);
198+
long staging = Unsafe.malloc(512, MemoryTag.NATIVE_DEFAULT);
190199
try {
191-
decoder.decode(buffer);
192-
Assert.fail("decoder must reject negative totalBytes");
193-
} catch (QwpDecodeException expected) {
194-
Assert.assertTrue("error message should describe invalid total bytes: " + expected.getMessage(),
195-
expected.getMessage().contains("total bytes"));
200+
int len = writeStringResultBatch(staging, /*nonNull=*/ 1, /*totalBytes=*/ -1);
201+
buffer.copyFromPayload(staging, len);
202+
try {
203+
decoder.decode(buffer);
204+
Assert.fail("decoder must reject negative totalBytes");
205+
} catch (QwpDecodeException expected) {
206+
Assert.assertTrue("error message should describe invalid total bytes: " + expected.getMessage(),
207+
expected.getMessage().contains("total bytes"));
208+
}
209+
} finally {
210+
Unsafe.free(staging, 512, MemoryTag.NATIVE_DEFAULT);
211+
buffer.close();
212+
decoder.close();
196213
}
197-
} finally {
198-
Unsafe.free(staging, 512, MemoryTag.NATIVE_DEFAULT);
199-
buffer.close();
200-
}
214+
});
201215
}
202216

203217
/**
@@ -206,19 +220,22 @@ public void testStringColumnNegativeTotalBytesIsRejected() {
206220
* the negative-value rejection above is testing the right code path.
207221
*/
208222
@Test
209-
public void testStringColumnValidTotalBytesIsAccepted() throws QwpDecodeException {
210-
QwpResultBatchDecoder decoder = new QwpResultBatchDecoder();
211-
QwpBatchBuffer buffer = new QwpBatchBuffer(512);
212-
long staging = Unsafe.malloc(512, MemoryTag.NATIVE_DEFAULT);
213-
try {
214-
int len = writeStringResultBatch(staging, /*nonNull=*/ 1, /*totalBytes=*/ 5);
215-
buffer.copyFromPayload(staging, len);
216-
decoder.decode(buffer);
217-
// no exception => the decoder accepts the valid wire bytes
218-
} finally {
219-
Unsafe.free(staging, 512, MemoryTag.NATIVE_DEFAULT);
220-
buffer.close();
221-
}
223+
public void testStringColumnValidTotalBytesIsAccepted() throws Exception {
224+
TestUtils.assertMemoryLeak(() -> {
225+
QwpResultBatchDecoder decoder = new QwpResultBatchDecoder();
226+
QwpBatchBuffer buffer = new QwpBatchBuffer(512);
227+
long staging = Unsafe.malloc(512, MemoryTag.NATIVE_DEFAULT);
228+
try {
229+
int len = writeStringResultBatch(staging, /*nonNull=*/ 1, /*totalBytes=*/ 5);
230+
buffer.copyFromPayload(staging, len);
231+
decoder.decode(buffer);
232+
// no exception => the decoder accepts the valid wire bytes
233+
} finally {
234+
Unsafe.free(staging, 512, MemoryTag.NATIVE_DEFAULT);
235+
buffer.close();
236+
decoder.close();
237+
}
238+
});
222239
}
223240

224241
// -----------------------------------------------------------------------

0 commit comments

Comments
 (0)