Update the pointer of Continuation's headerBuffer after copying the p…

…ayload (#507)
apache · Nov 27, 2024 · 191097a · 191097a
1 parent 306de8e
commit 191097a
Show file tree

Hide file tree

Showing 2 changed files with 57 additions and 5 deletions.
diff --git a/...ore5-h2/src/main/java/org/apache/hc/core5/http2/impl/nio/AbstractH2StreamMultiplexer.java b/...ore5-h2/src/main/java/org/apache/hc/core5/http2/impl/nio/AbstractH2StreamMultiplexer.java
@@ -1342,8 +1342,11 @@ void copyPayload(final ByteBuffer payload) {
             if (payload == null) {
                 return;
             }
-            headerBuffer.ensureCapacity(payload.remaining());
-            payload.get(headerBuffer.array(), headerBuffer.length(), payload.remaining());
+            final int originalLength = headerBuffer.length();
+            final int toCopy = payload.remaining();
+            headerBuffer.ensureCapacity(toCopy);
+            payload.get(headerBuffer.array(), originalLength, toCopy);
+            headerBuffer.setLength(originalLength + toCopy);
         }
 
         ByteBuffer getContent() {

diff --git a/...-h2/src/test/java/org/apache/hc/core5/http2/impl/nio/TestAbstractH2StreamMultiplexer.java b/...-h2/src/test/java/org/apache/hc/core5/http2/impl/nio/TestAbstractH2StreamMultiplexer.java
@@ -29,9 +29,15 @@
 
 import java.io.IOException;
 import java.nio.ByteBuffer;
+import java.util.ArrayList;
+import java.util.List;
 
+import org.apache.hc.core5.http.Header;
+import org.apache.hc.core5.http.HttpException;
 import org.apache.hc.core5.http.config.CharCodingConfig;
 import org.apache.hc.core5.http.impl.BasicHttpConnectionMetrics;
+import org.apache.hc.core5.http.impl.CharCodingSupport;
+import org.apache.hc.core5.http.message.BasicHeader;
 import org.apache.hc.core5.http.nio.AsyncPushConsumer;
 import org.apache.hc.core5.http.nio.HandlerFactory;
 import org.apache.hc.core5.http.nio.command.ExecutableCommand;
@@ -45,23 +51,31 @@
 import org.apache.hc.core5.http2.frame.FrameType;
 import org.apache.hc.core5.http2.frame.RawFrame;
 import org.apache.hc.core5.http2.frame.StreamIdGenerator;
+import org.apache.hc.core5.http2.hpack.HPackEncoder;
 import org.apache.hc.core5.reactor.ProtocolIOSession;
+import org.apache.hc.core5.util.ByteArrayBuffer;
 import org.junit.jupiter.api.Assertions;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
+import org.mockito.ArgumentCaptor;
 import org.mockito.ArgumentMatchers;
+import org.mockito.Captor;
 import org.mockito.Mock;
 import org.mockito.Mockito;
 import org.mockito.MockitoAnnotations;
 
 class TestAbstractH2StreamMultiplexer {
+    private static final FrameFactory FRAME_FACTORY = DefaultFrameFactory.INSTANCE;
+    private static final H2StreamHandler STREAM_HANDLER = Mockito.mock(H2StreamHandler.class);
 
     @Mock
     ProtocolIOSession protocolIOSession;
     @Mock
     HttpProcessor httpProcessor;
     @Mock
     H2StreamListener h2StreamListener;
+    @Captor
+    ArgumentCaptor<List<Header>> headersCaptor;
 
     @BeforeEach
     void prepareMocks() {
@@ -99,7 +113,7 @@ H2StreamHandler createRemotelyInitiatedStream(
                 final HttpProcessor httpProcessor,
                 final BasicHttpConnectionMetrics connMetrics,
                 final HandlerFactory<AsyncPushConsumer> pushHandlerFactory) throws IOException {
-            return null;
+            return STREAM_HANDLER;
         }
 
         @Override
@@ -128,7 +142,7 @@ void testInputOneFrame() throws Exception {
 
         final AbstractH2StreamMultiplexer streamMultiplexer = new H2StreamMultiplexerImpl(
                 protocolIOSession,
-                DefaultFrameFactory.INSTANCE,
+                FRAME_FACTORY,
                 StreamIdGenerator.ODD,
                 httpProcessor,
                 CharCodingConfig.DEFAULT,
@@ -179,7 +193,7 @@ void testInputMultipleFrames() throws Exception {
 
         final AbstractH2StreamMultiplexer streamMultiplexer = new H2StreamMultiplexerImpl(
                 protocolIOSession,
-                DefaultFrameFactory.INSTANCE,
+                FRAME_FACTORY,
                 StreamIdGenerator.ODD,
                 httpProcessor,
                 CharCodingConfig.DEFAULT,
@@ -212,5 +226,40 @@ void testInputMultipleFrames() throws Exception {
         });
     }
 
+    @Test
+    void testInputHeaderContinuationFrame() throws IOException, HttpException {
+        final H2Config h2Config = H2Config.custom().setMaxFrameSize(FrameConsts.MIN_FRAME_SIZE)
+                .build();
+
+        final ByteArrayBuffer buf = new ByteArrayBuffer(19);
+        final HPackEncoder encoder = new HPackEncoder(H2Config.INIT.getHeaderTableSize(), CharCodingSupport.createEncoder(CharCodingConfig.DEFAULT));
+        final List<Header> headers = new ArrayList<>();
+        headers.add(new BasicHeader("test-header-key", "value"));
+        headers.add(new BasicHeader(":status", "200"));
+        encoder.encodeHeaders(buf, headers, h2Config.isCompressionEnabled());
+
+        final WritableByteChannelMock writableChannel = new WritableByteChannelMock(1024);
+        final FrameOutputBuffer outBuffer = new FrameOutputBuffer(16 * 1024);
+
+        final RawFrame headerFrame = FRAME_FACTORY.createHeaders(2, ByteBuffer.wrap(buf.array(), 0, 10), false, false);
+        outBuffer.write(headerFrame, writableChannel);
+        final RawFrame continuationFrame = FRAME_FACTORY.createContinuation(2, ByteBuffer.wrap(buf.array(), 10, 9), true);
+        outBuffer.write(continuationFrame, writableChannel);
+        final byte[] bytes = writableChannel.toByteArray();
+
+        final AbstractH2StreamMultiplexer streamMultiplexer = new H2StreamMultiplexerImpl(
+                protocolIOSession,
+                FRAME_FACTORY,
+                StreamIdGenerator.ODD,
+                httpProcessor,
+                CharCodingConfig.DEFAULT,
+                h2Config,
+                h2StreamListener
+        );
+
+        streamMultiplexer.onInput(ByteBuffer.wrap(bytes));
+        Mockito.verify(STREAM_HANDLER).consumeHeader(headersCaptor.capture(), ArgumentMatchers.eq(false));
+        Assertions.assertFalse(headersCaptor.getValue().isEmpty());
+    }
 }