Fix StreamBufferingEncoder GOAWAY bug

codec-http2/src/main/java/io/netty/handler/codec/http2/StreamBufferingEncoder.java

@@ -69,33 +69,45 @@ public Http2ChannelClosedException() {
         }
     }
 
+    private static final class GoAwayDetail {
+        private final int lastStreamId;
+        private final long errorCode;
+        private final byte[] debugData;
+
+        GoAwayDetail(int lastStreamId, long errorCode, byte[] debugData) {
+            this.lastStreamId = lastStreamId;
+            this.errorCode = errorCode;
+            this.debugData = debugData.clone();
+        }
+    }
+
     /**
      * Thrown by {@link StreamBufferingEncoder} if buffered streams are terminated due to
      * receipt of a {@code GOAWAY}.
      */
     public static final class Http2GoAwayException extends Http2Exception {
         private static final long serialVersionUID = 1326785622777291198L;
-        private final int lastStreamId;
-        private final long errorCode;
-        private final byte[] debugData;
+        private GoAwayDetail goAwayDetail;
 
         public Http2GoAwayException(int lastStreamId, long errorCode, byte[] debugData) {
+            this(new GoAwayDetail(lastStreamId, errorCode, debugData));
+        }
+
+        private Http2GoAwayException(GoAwayDetail goAwayDetail) {
             super(Http2Error.STREAM_CLOSED);
-            this.lastStreamId = lastStreamId;
-            this.errorCode = errorCode;
-            this.debugData = debugData;
+            this.goAwayDetail = goAwayDetail;
         }
 
         public int lastStreamId() {
-            return lastStreamId;
+            return goAwayDetail.lastStreamId;
         }
 
         public long errorCode() {
-            return errorCode;
+            return goAwayDetail.errorCode;
         }
 
         public byte[] debugData() {
-            return debugData;
+            return goAwayDetail.debugData;
         }
     }
 
@@ -106,6 +118,7 @@ public byte[] debugData() {
     private final TreeMap<Integer, PendingStream> pendingStreams = new TreeMap<Integer, PendingStream>();
     private int maxConcurrentStreams;
     private boolean closed;
+    private GoAwayDetail goAwayDetail;
 
     public StreamBufferingEncoder(Http2ConnectionEncoder delegate) {
         this(delegate, SMALLEST_MAX_CONCURRENT_STREAMS);
@@ -118,7 +131,9 @@ public StreamBufferingEncoder(Http2ConnectionEncoder delegate, int initialMaxCon
 
             @Override
             public void onGoAwayReceived(int lastStreamId, long errorCode, ByteBuf debugData) {
-                cancelGoAwayStreams(lastStreamId, errorCode, debugData);
+                goAwayDetail = new GoAwayDetail(
+                    lastStreamId, errorCode, ByteBufUtil.getBytes(debugData));
+                cancelGoAwayStreams(goAwayDetail);
             }
 
             @Override
@@ -149,13 +164,14 @@ public ChannelFuture writeHeaders(ChannelHandlerContext ctx, int streamId, Http2
         if (closed) {
             return promise.setFailure(new Http2ChannelClosedException());
         }
-        if (isExistingStream(streamId) || connection().goAwayReceived()) {
+        if (isExistingStream(streamId) || canCreateStream()) {
             return super.writeHeaders(ctx, streamId, headers, streamDependency, weight,
                     exclusive, padding, endOfStream, promise);
         }
-        if (canCreateStream()) {
-            return super.writeHeaders(ctx, streamId, headers, streamDependency, weight,
-                    exclusive, padding, endOfStream, promise);
+        if (goAwayDetail != null) {
+            promise.setFailure(new Http2GoAwayException(
+                goAwayDetail.lastStreamId, goAwayDetail.lastStreamId, goAwayDetail.debugData));
+            return promise;
         }
         PendingStream pendingStream = pendingStreams.get(streamId);
         if (pendingStream == null) {
@@ -248,12 +264,12 @@ private void tryCreatePendingStreams() {
         }
     }
 
-    private void cancelGoAwayStreams(int lastStreamId, long errorCode, ByteBuf debugData) {
+    private void cancelGoAwayStreams(GoAwayDetail goAwayDetail) {
         Iterator<PendingStream> iter = pendingStreams.values().iterator();
-        Exception e = new Http2GoAwayException(lastStreamId, errorCode, ByteBufUtil.getBytes(debugData));
+        Exception e = new Http2GoAwayException(goAwayDetail);
         while (iter.hasNext()) {
             PendingStream stream = iter.next();
-            if (stream.streamId > lastStreamId) {
+            if (stream.streamId > goAwayDetail.lastStreamId) {
                 iter.remove();
                 stream.close(e);
             }

codec-http2/src/test/java/io/netty/handler/codec/http2/StreamBufferingEncoderTest.java

@@ -49,6 +49,7 @@
 import io.netty.channel.ChannelPromise;
 import io.netty.channel.DefaultChannelPromise;
 import io.netty.channel.DefaultMessageSizeEstimator;
+import io.netty.handler.codec.http2.StreamBufferingEncoder.Http2GoAwayException;
 import io.netty.util.ReferenceCountUtil;
 import io.netty.util.concurrent.EventExecutor;
 import io.netty.util.concurrent.ImmediateEventExecutor;
@@ -111,6 +112,11 @@ public void setup() throws Exception {
         when(writer.writeGoAway(any(ChannelHandlerContext.class), anyInt(), anyLong(), any(ByteBuf.class),
                 any(ChannelPromise.class)))
                 .thenAnswer(successAnswer());
+        when(writer.writeHeaders(any(ChannelHandlerContext.class), anyInt(), any(Http2Headers.class),
+            anyInt(), anyBoolean(), any(ChannelPromise.class))).thenAnswer(noopAnswer());
+        when(writer.writeHeaders(any(ChannelHandlerContext.class), anyInt(), any(Http2Headers.class),
+            anyInt(), anyShort(), anyBoolean(), anyInt(), anyBoolean(), any(ChannelPromise.class)))
+            .thenAnswer(noopAnswer());
 
         connection = new DefaultHttp2Connection(false);
         connection.remote().flowController(new DefaultHttp2RemoteFlowController(connection));
@@ -167,7 +173,7 @@ public void multipleWritesToActiveStream() {
         encoder.writeData(ctx, 3, data(), 0, false, newPromise());
         encoderWriteHeaders(3, newPromise());
 
-        writeVerifyWriteHeaders(times(2), 3);
+        writeVerifyWriteHeaders(times(1), 3);
         // Contiguous data writes are coalesced
         ArgumentCaptor<ByteBuf> bufCaptor = ArgumentCaptor.forClass(ByteBuf.class);
         verify(writer, times(1))
@@ -245,18 +251,32 @@ public void receivingGoAwayFailsBufferedStreams() throws Http2Exception {
             futures.add(encoderWriteHeaders(streamId, newPromise()));
             streamId += 2;
         }
+        assertEquals(5, connection.numActiveStreams());
         assertEquals(4, encoder.numBufferedStreams());
 
         connection.goAwayReceived(11, 8, EMPTY_BUFFER);
 
         assertEquals(5, connection.numActiveStreams());
+        assertEquals(0, encoder.numBufferedStreams());
         int failCount = 0;
         for (ChannelFuture f : futures) {
             if (f.cause() != null) {
+                assertTrue(f.cause() instanceof Http2GoAwayException);
                 failCount++;
             }
         }
-        assertEquals(9, failCount);
+        assertEquals(4, failCount);
+    }
+
+    @Test
+    public void receivingGoAwayFailsNewStreamIfMaxConcurrentStreamsReached() throws Http2Exception {
+        encoder.writeSettingsAck(ctx, newPromise());
+        setMaxConcurrentStreams(1);
+        encoderWriteHeaders(3, newPromise());
+        connection.goAwayReceived(11, 8, EMPTY_BUFFER);
+        ChannelFuture f = encoderWriteHeaders(5, newPromise());
+
+        assertTrue(f.cause() instanceof Http2GoAwayException);
         assertEquals(0, encoder.numBufferedStreams());
     }
 
@@ -533,6 +553,20 @@ public ChannelFuture answer(InvocationOnMock invocation) throws Throwable {
         };
     }
 
+    private Answer<ChannelFuture> noopAnswer() {
+        return new Answer<ChannelFuture>() {
+            @Override
+            public ChannelFuture answer(InvocationOnMock invocation) throws Throwable {
+                for (Object a : invocation.getArguments()) {
+                    if (a instanceof ChannelPromise) {
+                        return (ChannelFuture) a;
+                    }
+                }
+                return newPromise();
+            }
+        };
+    }
+
     private ChannelPromise newPromise() {
         return new DefaultChannelPromise(channel, ImmediateEventExecutor.INSTANCE);
     }