From a52f47c9a99c6e4d58485e2aff6c782092ae5efa Mon Sep 17 00:00:00 2001
From: Chung-Ting Huang <chungting@cloudflare.com>
Date: Wed, 23 Oct 2019 16:02:33 -0500
Subject: [PATCH] TUN-2460: Configure according to the ClientConfig recevied
 from a successful Connect

---
 connection/manager.go      |  8 ++++++--
 connection/manager_test.go | 14 ++++++--------
 2 files changed, 12 insertions(+), 10 deletions(-)

diff --git a/connection/manager.go b/connection/manager.go
index b59d6ee8..61f26c36 100644
--- a/connection/manager.go
+++ b/connection/manager.go
@@ -10,7 +10,9 @@ import (
 
 	"github.com/cloudflare/cloudflared/cmd/cloudflared/buildinfo"
 	"github.com/cloudflare/cloudflared/h2mux"
+	"github.com/cloudflare/cloudflared/streamhandler"
 	"github.com/cloudflare/cloudflared/tunnelrpc/pogs"
+
 	"github.com/google/uuid"
 	"github.com/pkg/errors"
 	"github.com/sirupsen/logrus"
@@ -25,7 +27,7 @@ const (
 // EdgeManager manages connections with the edge
 type EdgeManager struct {
 	// streamHandler handles stream opened by the edge
-	streamHandler h2mux.MuxedStreamHandler
+	streamHandler *streamhandler.StreamHandler
 	// TLSConfig is the TLS configuration to connect with edge
 	tlsConfig *tls.Config
 	// cloudflaredConfig is the cloudflared configuration that is determined when the process first starts
@@ -52,7 +54,7 @@ type CloudflaredConfig struct {
 }
 
 func NewEdgeManager(
-	streamHandler h2mux.MuxedStreamHandler,
+	streamHandler *streamhandler.StreamHandler,
 	edgeConnMgrConfigurable *EdgeManagerConfigurable,
 	userCredential []byte,
 	tlsConfig *tls.Config,
@@ -155,6 +157,8 @@ func (em *EdgeManager) newConnection(ctx context.Context) *pogs.ConnectError {
 
 	em.state.newConnection(h2muxConn)
 	em.logger.Infof("connected to %s", connResult.ConnectedTo())
+
+	em.streamHandler.UseConfiguration(ctx, connResult.ClientConfig())
 	return nil
 }
 
diff --git a/connection/manager_test.go b/connection/manager_test.go
index 7565567f..6732ce24 100644
--- a/connection/manager_test.go
+++ b/connection/manager_test.go
@@ -8,7 +8,9 @@ import (
 	"github.com/stretchr/testify/assert"
 
 	"github.com/cloudflare/cloudflared/h2mux"
+	"github.com/cloudflare/cloudflared/streamhandler"
 	"github.com/cloudflare/cloudflared/tunnelrpc/pogs"
+
 	"github.com/google/uuid"
 	"github.com/sirupsen/logrus"
 )
@@ -42,16 +44,12 @@ var (
 	}
 )
 
-type mockStreamHandler struct {
-}
-
-func (msh *mockStreamHandler) ServeStream(*h2mux.MuxedStream) error {
-	return nil
-}
-
 func mockEdgeManager() *EdgeManager {
+	newConfigChan := make(chan<- *pogs.ClientConfig)
+	useConfigResultChan := make(<-chan *pogs.UseConfigurationResult)
+	logger := logrus.New()
 	return NewEdgeManager(
-		&mockStreamHandler{},
+		streamhandler.NewStreamHandler(newConfigChan, useConfigResultChan, logger),
 		configurable,
 		[]byte{},
 		nil,