From 89b738f8fa0b3ca762df2607177b2ad77bf8a26f Mon Sep 17 00:00:00 2001
From: Nuno Diegues <nuno@cloudflare.com>
Date: Thu, 4 Mar 2021 18:45:39 +0000
Subject: [PATCH] TUN-4026: Fix regression where HTTP2 edge transport was no
 longer propagating control plane errors

---
 connection/http2.go      |  1 +
 connection/http2_test.go | 39 +++++++++++++++++++++++++++++++++++++++
 2 files changed, 40 insertions(+)

diff --git a/connection/http2.go b/connection/http2.go
index 086975a0..3412a05e 100644
--- a/connection/http2.go
+++ b/connection/http2.go
@@ -108,6 +108,7 @@ func (c *http2Connection) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	switch connType {
 	case TypeControlStream:
 		proxyErr = c.serveControlStream(r.Context(), respWriter)
+		c.controlStreamErr = proxyErr
 	case TypeWebsocket:
 		stripWebsocketUpgradeHeader(r)
 		proxyErr = c.config.OriginProxy.Proxy(respWriter, r, TypeWebsocket)
diff --git a/connection/http2_test.go b/connection/http2_test.go
index 8124ff86..a8d474bd 100644
--- a/connection/http2_test.go
+++ b/connection/http2_test.go
@@ -114,6 +114,7 @@ func TestServeHTTP(t *testing.T) {
 }
 
 type mockNamedTunnelRPCClient struct {
+	shouldFail   error
 	registered   chan struct{}
 	unregistered chan struct{}
 }
@@ -125,6 +126,9 @@ func (mc mockNamedTunnelRPCClient) RegisterConnection(
 	connIndex uint8,
 	observer *Observer,
 ) error {
+	if mc.shouldFail != nil {
+		return mc.shouldFail
+	}
 	close(mc.registered)
 	return nil
 }
@@ -136,12 +140,14 @@ func (mc mockNamedTunnelRPCClient) GracefulShutdown(ctx context.Context, gracePe
 func (mockNamedTunnelRPCClient) Close() {}
 
 type mockRPCClientFactory struct {
+	shouldFail   error
 	registered   chan struct{}
 	unregistered chan struct{}
 }
 
 func (mf *mockRPCClientFactory) newMockRPCClient(context.Context, io.ReadWriteCloser, *zerolog.Logger) NamedTunnelRPCClient {
 	return mockNamedTunnelRPCClient{
+		shouldFail:   mf.shouldFail,
 		registered:   mf.registered,
 		unregistered: mf.unregistered,
 	}
@@ -249,6 +255,39 @@ func TestServeControlStream(t *testing.T) {
 	wg.Wait()
 }
 
+func TestFailRegistration(t *testing.T) {
+	http2Conn, edgeConn := newTestHTTP2Connection()
+
+	rpcClientFactory := mockRPCClientFactory{
+		shouldFail:   errDuplicationConnection,
+		registered:   make(chan struct{}),
+		unregistered: make(chan struct{}),
+	}
+	http2Conn.newRPCClientFunc = rpcClientFactory.newMockRPCClient
+
+	ctx, cancel := context.WithCancel(context.Background())
+	var wg sync.WaitGroup
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+		http2Conn.Serve(ctx)
+	}()
+
+	req, err := http.NewRequestWithContext(ctx, http.MethodGet, "http://localhost:8080/", nil)
+	require.NoError(t, err)
+	req.Header.Set(internalUpgradeHeader, controlStreamUpgrade)
+
+	edgeHTTP2Conn, err := testTransport.NewClientConn(edgeConn)
+	require.NoError(t, err)
+	resp, err := edgeHTTP2Conn.RoundTrip(req)
+	require.NoError(t, err)
+	require.Equal(t, http.StatusBadGateway, resp.StatusCode)
+
+	assert.NotNil(t, http2Conn.controlStreamErr)
+	cancel()
+	wg.Wait()
+}
+
 func TestGracefulShutdownHTTP2(t *testing.T) {
 	http2Conn, edgeConn := newTestHTTP2Connection()