Skip to content

Commit 91f3134

Browse files
tabish121tabish121
committed
ARTEMIS-5689 Fix tunneled core message writer encoding
Ensure that the read size from the delivery annotations and the core header and properties encoding buffer is calculated correctly so that when spanning multiple frames an error is not thrown from the copy.
1 parent 7296823 commit 91f3134

2 files changed

Lines changed: 177 additions & 11 deletions

File tree

‎artemis-protocols/artemis-amqp-protocol/src/main/java/org/apache/activemq/artemis/protocol/amqp/proton/AMQPTunneledCoreLargeMessageWriter.java‎

Lines changed: 7 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -219,15 +219,14 @@ private boolean trySendDeliveryAnnotations(ByteBuf frameBuffer, NettyReadable fr
219219
for (; protonSender.getLocalState() != EndpointState.CLOSED && state == State.STREAMING_DELIVERY_ANNOTATIONS; ) {
220220
if (annotations != null && annotations.getValue() != null && !annotations.getValue().isEmpty()) {
221221
if (!connection.flowControl(this::resume)) {
222-
break; // Resume will restart writing the headers section from where we left off.
222+
break; // Resume will restart writing the delivery annotations section from where we left off.
223223
}
224224

225225
final ByteBuf annotationsBuffer = getOrCreateDeliveryAnnotationsBuffer();
226-
final int readSize = (int) Math.min(frameBuffer.writableBytes(), annotationsBuffer.readableBytes() - position);
226+
final int readSize = Math.min(frameBuffer.writableBytes(), annotationsBuffer.readableBytes());
227227

228-
position += readSize;
229-
230-
annotationsBuffer.readBytes(frameBuffer, readSize);
228+
annotationsBuffer.readBytes(frameBuffer.internalNioBuffer(frameBuffer.writerIndex(), readSize));
229+
frameBuffer.writerIndex(frameBuffer.writerIndex() + readSize);
231230

232231
// In case the Delivery Annotations encoding exceed the AMQP frame size we
233232
// flush and keep sending until done or until flow controlled.
@@ -239,7 +238,6 @@ private boolean trySendDeliveryAnnotations(ByteBuf frameBuffer, NettyReadable fr
239238

240239
if (!annotationsBuffer.isReadable()) {
241240
encodingBuffer = null;
242-
position = 0;
243241
state = State.STREAMING_CORE_HEADERS;
244242
}
245243
} else {
@@ -259,11 +257,10 @@ private boolean trySendHeadersAndProperties(ByteBuf frameBuffer, NettyReadable f
259257
}
260258

261259
final ByteBuf headerBuffer = getOrCreateMessageHeaderBuffer();
262-
final int readSize = (int) Math.min(frameBuffer.writableBytes(), headerBuffer.readableBytes() - position);
263-
264-
position += readSize;
260+
final int readSize = Math.min(frameBuffer.writableBytes(), headerBuffer.readableBytes());
265261

266-
headerBuffer.readBytes(frameBuffer, readSize);
262+
headerBuffer.readBytes(frameBuffer.internalNioBuffer(frameBuffer.writerIndex(), readSize));
263+
frameBuffer.writerIndex(frameBuffer.writerIndex() + readSize);
267264

268265
// In case the Core message header and properties exceed the AMQP frame size we
269266
// flush and keep sending until done or until flow controlled.
@@ -275,7 +272,6 @@ private boolean trySendHeadersAndProperties(ByteBuf frameBuffer, NettyReadable f
275272

276273
if (!headerBuffer.isReadable()) {
277274
encodingBuffer = null;
278-
position = 0;
279275
state = State.STREAMING_BODY;
280276
}
281277
}

‎artemis-protocols/artemis-amqp-protocol/src/test/java/org/apache/activemq/artemis/protocol/amqp/proton/AMQPTunneledCoreLargeMessageWriterTest.java‎

Lines changed: 170 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -30,7 +30,10 @@
3030
import static org.mockito.Mockito.when;
3131

3232
import java.nio.ByteBuffer;
33+
import java.nio.charset.StandardCharsets;
3334
import java.util.HashMap;
35+
import java.util.concurrent.atomic.AtomicInteger;
36+
3437
import org.apache.activemq.artemis.core.message.LargeBodyReader;
3538
import org.apache.activemq.artemis.core.persistence.impl.journal.LargeServerMessageImpl;
3639
import org.apache.activemq.artemis.core.persistence.impl.nullpm.NullStorageManager;
@@ -273,6 +276,173 @@ private void doTestMessageEncodingWrittenToDeliveryWithAnnotations(boolean deliv
273276
verifyNoMoreInteractions(protonDelivery);
274277
}
275278

279+
@Test
280+
public void testMessageEncodingWrittenToDeliveryWithDeliveryAnnotationsThatExceedFrameSize() throws Exception {
281+
final byte[] headersBytes = new byte[4];
282+
283+
headersBytes[0] = 4;
284+
headersBytes[1] = 5;
285+
headersBytes[2] = 6;
286+
headersBytes[3] = 7;
287+
288+
final byte[] payloadBytes = new byte[4];
289+
290+
payloadBytes[0] = 1;
291+
payloadBytes[1] = 2;
292+
payloadBytes[2] = 3;
293+
payloadBytes[3] = 4;
294+
295+
final DeliveryAnnotations annotations = new DeliveryAnnotations(new HashMap<>());
296+
297+
annotations.getValue().put(Symbol.valueOf("a"), "a".repeat(1024));
298+
annotations.getValue().put(Symbol.valueOf("b"), "b".repeat(1024));
299+
annotations.getValue().put(Symbol.valueOf("c"), "c".repeat(1024));
300+
annotations.getValue().put(Symbol.valueOf("d"), "d".repeat(1024));
301+
annotations.getValue().put(Symbol.valueOf("e"), "e".repeat(1024));
302+
303+
doTestMessageEncodingForTunneledCoreLargeMessage(annotations, headersBytes, payloadBytes);
304+
}
305+
306+
@Test
307+
public void testMessageEncodingWrittenToDeliveryWithCoreHeaderEncodingThatExceedsFrameSize() throws Exception {
308+
final byte[] headersBytes = "A".repeat(4097).getBytes(StandardCharsets.US_ASCII);
309+
final byte[] payloadBytes = new byte[4];
310+
311+
payloadBytes[0] = 1;
312+
payloadBytes[1] = 2;
313+
payloadBytes[2] = 3;
314+
payloadBytes[3] = 4;
315+
316+
final DeliveryAnnotations annotations = new DeliveryAnnotations(new HashMap<>());
317+
318+
annotations.getValue().put(Symbol.valueOf("a"), "a");
319+
annotations.getValue().put(Symbol.valueOf("b"), "b");
320+
annotations.getValue().put(Symbol.valueOf("c"), "c");
321+
annotations.getValue().put(Symbol.valueOf("d"), "d");
322+
annotations.getValue().put(Symbol.valueOf("e"), "e");
323+
324+
doTestMessageEncodingForTunneledCoreLargeMessage(annotations, headersBytes, payloadBytes);
325+
}
326+
327+
@Test
328+
public void testMessageEncodingWrittenToDeliveryWithBothDAandCoreHeadersExceedingFrameSize() throws Exception {
329+
final byte[] headersBytes = "A".repeat(4097).getBytes(StandardCharsets.US_ASCII);
330+
final byte[] payloadBytes = new byte[4];
331+
332+
payloadBytes[0] = 1;
333+
payloadBytes[1] = 2;
334+
payloadBytes[2] = 3;
335+
payloadBytes[3] = 4;
336+
337+
final DeliveryAnnotations annotations = new DeliveryAnnotations(new HashMap<>());
338+
339+
annotations.getValue().put(Symbol.valueOf("a"), "a".repeat(1024));
340+
annotations.getValue().put(Symbol.valueOf("b"), "b".repeat(1024));
341+
annotations.getValue().put(Symbol.valueOf("c"), "c".repeat(1024));
342+
annotations.getValue().put(Symbol.valueOf("d"), "d".repeat(1024));
343+
annotations.getValue().put(Symbol.valueOf("e"), "e".repeat(1024));
344+
345+
doTestMessageEncodingForTunneledCoreLargeMessage(annotations, headersBytes, payloadBytes);
346+
}
347+
348+
@Test
349+
public void testMessageEncodingWrittenToDeliveryWithAllSectionsExceedFrameSize() throws Exception {
350+
final byte[] headersBytes = "A".repeat(4097).getBytes(StandardCharsets.US_ASCII);
351+
final byte[] payloadBytes = "B".repeat(4097).getBytes(StandardCharsets.US_ASCII);
352+
final DeliveryAnnotations annotations = new DeliveryAnnotations(new HashMap<>());
353+
354+
annotations.getValue().put(Symbol.valueOf("a"), "a".repeat(1024));
355+
annotations.getValue().put(Symbol.valueOf("b"), "b".repeat(1024));
356+
annotations.getValue().put(Symbol.valueOf("c"), "c".repeat(1024));
357+
annotations.getValue().put(Symbol.valueOf("d"), "d".repeat(1024));
358+
annotations.getValue().put(Symbol.valueOf("e"), "e".repeat(1024));
359+
360+
doTestMessageEncodingForTunneledCoreLargeMessage(annotations, headersBytes, payloadBytes);
361+
}
362+
363+
private void doTestMessageEncodingForTunneledCoreLargeMessage(DeliveryAnnotations annotations, byte[] headersBytes, byte[] payloadBytes) throws Exception {
364+
when(protonTransport.getOutboundFrameSizeLimit()).thenReturn(4096);
365+
366+
AMQPTunneledCoreLargeMessageWriter writer = new AMQPTunneledCoreLargeMessageWriter(serverSender);
367+
368+
writer.open(Mockito.mock(MessageReference.class));
369+
370+
final ByteBuf expectedEncoding = Unpooled.buffer();
371+
final AtomicInteger payloadReaderPosition = new AtomicInteger();
372+
373+
writeDeliveryAnnotations(expectedEncoding, annotations);
374+
375+
when(reference.getProtocolData(any())).thenReturn(annotations);
376+
377+
writeDataSection(expectedEncoding, headersBytes);
378+
writeDataSection(expectedEncoding, payloadBytes);
379+
380+
when(protonSender.getLocalState()).thenReturn(EndpointState.ACTIVE);
381+
when(protonDelivery.isPartial()).thenReturn(true);
382+
when(message.getHeadersAndPropertiesEncodeSize()).thenReturn(headersBytes.length);
383+
384+
// Provides the simulated encoded core headers and properties
385+
doAnswer(invocation -> {
386+
final ByteBuf buffer = invocation.getArgument(0);
387+
388+
buffer.writeBytes(headersBytes);
389+
390+
return null;
391+
}).when(message).encodeHeadersAndProperties(any(ByteBuf.class));
392+
393+
when(bodyReader.getSize()).thenReturn((long) payloadBytes.length);
394+
395+
final ByteBuf encodedByteBuf = Unpooled.buffer();
396+
final NettyWritable encodedBytes = new NettyWritable(encodedByteBuf);
397+
398+
// Answer back with the amount of writable bytes
399+
doAnswer(invocation -> {
400+
final ByteBuffer buffer = invocation.getArgument(0);
401+
402+
final int readSize = Math.min(buffer.remaining(), payloadBytes.length - payloadReaderPosition.get());
403+
404+
buffer.put(payloadBytes, payloadReaderPosition.get(), readSize);
405+
406+
payloadReaderPosition.addAndGet(readSize);
407+
408+
return readSize;
409+
}).when(bodyReader).readInto(any());
410+
411+
// Capture the write for comparison, this avoid issues with released netty buffers
412+
doAnswer(invocation -> {
413+
final ReadableBuffer buffer = invocation.getArgument(0);
414+
415+
encodedBytes.put(buffer);
416+
417+
return null;
418+
}).when(protonSender).send(any());
419+
420+
try {
421+
writer.writeBytes(reference);
422+
} catch (IllegalStateException e) {
423+
fail("Should not throw as the delivery is completed so no data should be written.");
424+
}
425+
426+
verify(message).usageUp();
427+
verify(message).getLargeBodyReader();
428+
verify(message).getHeadersAndPropertiesEncodeSize();
429+
verify(message).encodeHeadersAndProperties(any(ByteBuf.class));
430+
verify(reference).getMessage();
431+
verify(reference).getProtocolData(any());
432+
verify(protonSender).getSession();
433+
verify(protonDelivery).getTag();
434+
verify(protonSender, atLeastOnce()).getLocalState();
435+
verify(protonSender, atLeastOnce()).send(any(ReadableBuffer.class));
436+
437+
assertTrue(encodedByteBuf.isReadable());
438+
assertEquals(expectedEncoding.readableBytes(), encodedByteBuf.readableBytes());
439+
assertEquals(expectedEncoding, encodedByteBuf);
440+
441+
verifyNoMoreInteractions(message);
442+
verifyNoMoreInteractions(reference);
443+
verifyNoMoreInteractions(protonDelivery);
444+
}
445+
276446
@Test
277447
public void testLargeMessageUsageLoweredOnCloseWhenWriteNotCompleted() throws Exception {
278448
AMQPTunneledCoreLargeMessageWriter writer = new AMQPTunneledCoreLargeMessageWriter(serverSender);

0 commit comments

Comments
 (0)