Fix StreamBufferingEncoder GOAWAY bug

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

@@ -149,10 +149,14 @@ public ChannelFuture writeHeaders(ChannelHandlerContext ctx, int streamId, Http2
         if (closed) {
             return promise.setFailure(new Http2ChannelClosedException());
         }
-        if (isExistingStream(streamId) || connection().goAwayReceived()) {
+        if (isExistingStream(streamId)) {
             return super.writeHeaders(ctx, streamId, headers, streamDependency, weight,
                     exclusive, padding, endOfStream, promise);
         }
+        if (connection().goAwayReceived()) {
+            promise.setFailure(new Http2Exception(Http2Error.NO_ERROR, "GOAWAY received"));
+            return promise;
+        }
         if (canCreateStream()) {
             return super.writeHeaders(ctx, streamId, headers, streamDependency, weight,
                     exclusive, padding, endOfStream, promise);

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

@@ -111,6 +111,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 +172,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))
@@ -253,10 +258,24 @@ public void receivingGoAwayFailsBufferedStreams() throws Http2Exception {
         int failCount = 0;
         for (ChannelFuture f : futures) {
             if (f.cause() != null) {
+                assertTrue(f.cause() instanceof Http2Exception);
+                Http2Exception e = (Http2Exception) f.cause();
+                assertEquals(Http2Error.STREAM_CLOSED, e.error());
                 failCount++;
             }
         }
-        assertEquals(9, failCount);
+    }
+
+    @Test
+    public void receivingGoAwayFailsNewStreams() throws Http2Exception {
+        encoder.writeSettingsAck(ctx, newPromise());
+        connection.goAwayReceived(11, 8, EMPTY_BUFFER);
+        ChannelFuture f = encoderWriteHeaders(3, newPromise());
+
+        assertTrue(f.cause() instanceof Http2Exception);
+        Http2Exception e = (Http2Exception) f.cause();
+        assertEquals(Http2Error.NO_ERROR, e.error());
+        assertEquals("GOAWAY received", e.getMessage());
         assertEquals(0, encoder.numBufferedStreams());
     }
 
@@ -533,6 +552,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);
     }