From 5376df543908b7b8cbe4dfbfa14d26787fd2b027 Mon Sep 17 00:00:00 2001
From: Adam Chalmers <achalmers@cloudflare.com>
Date: Fri, 6 Mar 2020 17:25:34 -0600
Subject: [PATCH] TUN-2788: cloudflared should store one ConnDigest per HA
 connection

---
 origin/supervisor.go | 17 +++++++++--------
 origin/tunnel.go     |  8 ++++----
 2 files changed, 13 insertions(+), 12 deletions(-)

diff --git a/origin/supervisor.go b/origin/supervisor.go
index 1c7e7dfe..1cb7c336 100644
--- a/origin/supervisor.go
+++ b/origin/supervisor.go
@@ -39,7 +39,6 @@ const (
 var (
 	errJWTUnset         = errors.New("JWT unset")
 	errEventDigestUnset = errors.New("event digest unset")
-	errConnDigestUnset  = errors.New("conn digest unset")
 )
 
 // Supervisor manages non-declarative tunnels. Establishes TCP connections with the edge, and
@@ -66,7 +65,7 @@ type Supervisor struct {
 	eventDigest     []byte
 
 	connDigestLock sync.RWMutex
-	connDigest     []byte
+	connDigest     map[uint8][]byte
 
 	bufferPool *buffer.Pool
 }
@@ -101,6 +100,7 @@ func NewSupervisor(config *TunnelConfig, u uuid.UUID) (*Supervisor, error) {
 		tunnelErrors:      make(chan tunnelError),
 		tunnelsConnecting: map[int]chan struct{}{},
 		logger:            config.Logger.WithField("subsystem", "supervisor"),
+		connDigest:        make(map[uint8][]byte),
 		bufferPool:        buffer.NewPool(512 * 1024),
 	}, nil
 }
@@ -334,19 +334,20 @@ func (s *Supervisor) SetEventDigest(eventDigest []byte) {
 	s.eventDigest = eventDigest
 }
 
-func (s *Supervisor) ConnDigest() ([]byte, error) {
+func (s *Supervisor) ConnDigest(connID uint8) ([]byte, error) {
 	s.connDigestLock.RLock()
 	defer s.connDigestLock.RUnlock()
-	if s.connDigest == nil {
-		return nil, errConnDigestUnset
+	digest, ok := s.connDigest[connID]
+	if !ok {
+		return nil, fmt.Errorf("no connection digest for connection %v", connID)
 	}
-	return s.connDigest, nil
+	return digest, nil
 }
 
-func (s *Supervisor) SetConnDigest(connDigest []byte) {
+func (s *Supervisor) SetConnDigest(connID uint8, connDigest []byte) {
 	s.connDigestLock.Lock()
 	defer s.connDigestLock.Unlock()
-	s.connDigest = connDigest
+	s.connDigest[connID] = connDigest
 }
 
 func (s *Supervisor) refreshAuth(
diff --git a/origin/tunnel.go b/origin/tunnel.go
index 2aaba194..55a9e4c1 100644
--- a/origin/tunnel.go
+++ b/origin/tunnel.go
@@ -95,8 +95,8 @@ type ReconnectTunnelCredentialManager interface {
 	ReconnectToken() ([]byte, error)
 	EventDigest() ([]byte, error)
 	SetEventDigest(eventDigest []byte)
-	ConnDigest() ([]byte, error)
-	SetConnDigest(connDigest []byte)
+	ConnDigest(connID uint8) ([]byte, error)
+	SetConnDigest(connID uint8, connDigest []byte)
 }
 
 type dupConnRegisterTunnelError struct{}
@@ -286,7 +286,7 @@ func ServeTunnel(
 
 				// check if we can use Quick Reconnects
 				if config.UseQuickReconnects {
-					if digest, connDigestErr := credentialManager.ConnDigest(); connDigestErr == nil {
+					if digest, connDigestErr := credentialManager.ConnDigest(connectionID); connDigestErr == nil {
 						connDigest = digest
 					}
 				}
@@ -392,7 +392,7 @@ func RegisterTunnel(
 		return processRegisterTunnelError(registrationErr, config.Metrics, register)
 	}
 	credentialManager.SetEventDigest(registration.EventDigest)
-	credentialManager.SetConnDigest(registration.ConnDigest)
+	credentialManager.SetConnDigest(connectionID, registration.ConnDigest)
 	return processRegistrationSuccess(config, logger, connectionID, registration, register)
 }