V7 backport #9958 and #10368

lib/cache/cache.go

@@ -303,7 +303,7 @@ type Cache struct {
 
 	// fnCache is used to perform short ttl-based caching of the results of
 	// regularly called methods.
-	fnCache *fnCache
+	fnCache *utils.FnCache
 
 	trustCache         services.Trust
 	clusterConfigCache services.ClusterConfiguration
@@ -568,6 +568,14 @@ func New(config Config) (*Cache, error) {
 		return nil, trace.Wrap(err)
 	}
 
+	fnCache, err := utils.NewFnCache(utils.FnCacheConfig{
+		TTL:   time.Second,
+		Clock: config.Clock,
+	})
+	if err != nil {
+		return nil, trace.Wrap(err)
+	}
+
 	ctx, cancel := context.WithCancel(config.Context)
 	cs := &Cache{
 		wrapper:            wrapper,
@@ -576,7 +584,7 @@ func New(config Config) (*Cache, error) {
 		Config:             config,
 		generation:         atomic.NewUint64(0),
 		initC:              make(chan struct{}),
-		fnCache:            newFnCache(time.Second),
+		fnCache:            fnCache,
 		trustCache:         local.NewCAService(wrapper),
 		clusterConfigCache: clusterConfigCache,
 		provisionerCache:   local.NewProvisioningService(wrapper),

lib/reversetunnel/agent.go

@@ -345,12 +345,19 @@ func (a *Agent) run() {
 		a.log.Warningf("Failed to create remote tunnel: %v, conn: %v.", err, conn)
 		return
 	}
-	defer conn.Close()
+
+	local := conn.LocalAddr().String()
+	remote := conn.RemoteAddr().String()
+	defer func() {
+		if err := conn.Close(); err != nil {
+			a.log.Warnf("Failed to close remote tunnel: %v, local addr: %s remote addr: %s", err, local, remote)
+		}
+	}()
 
 	// Successfully connected to remote cluster.
 	a.log.WithFields(log.Fields{
-		"addr":        conn.LocalAddr().String(),
-		"remote-addr": conn.RemoteAddr().String(),
+		"addr":        local,
+		"remote-addr": remote,
 	}).Info("Connected.")
 
 	// wrap up remaining business logic in closure for easy
@@ -375,14 +382,14 @@ func (a *Agent) run() {
 		// or permanent loss of a proxy.
 		err = a.processRequests(conn)
 		if err != nil {
-			a.log.Warnf("Unable to continue processesing requests: %v.", err)
+			a.log.Warnf("Unable to continue processioning requests: %v.", err)
 			return
 		}
 	}
 	// if Tracker was provided, then the agent shouldn't continue unless
 	// no other agents hold a claim.
 	if a.Tracker != nil {
-		if !a.Tracker.WithProxy(doWork, a.Lease, a.getPrincipalsList()...) {
+		if !a.Tracker.WithProxy(doWork, a.getPrincipalsList()...) {
 			a.log.Debugf("Proxy already held by other agent: %v, releasing.", a.getPrincipalsList())
 		}
 	} else {
@@ -478,15 +485,19 @@ func (a *Agent) processRequests(conn *ssh.Client) error {
 	}
 }
 
-// handleDisovery receives discovery requests from the reverse tunnel
+// handleDiscovery receives discovery requests from the reverse tunnel
 // server, that informs agent about proxies registered in the remote
 // cluster and the reverse tunnels already established
 //
 // ch   : SSH channel which received "teleport-transport" out-of-band request
 // reqC : request payload
 func (a *Agent) handleDiscovery(ch ssh.Channel, reqC <-chan *ssh.Request) {
 	a.log.Debugf("handleDiscovery requests channel.")
-	defer ch.Close()
+	defer func() {
+		if err := ch.Close(); err != nil {
+			a.log.Warnf("Failed to close discovery channel:: %v", err)
+		}
+	}()
 
 	for {
 		var req *ssh.Request
@@ -507,7 +518,7 @@ func (a *Agent) handleDiscovery(ch ssh.Channel, reqC <-chan *ssh.Request) {
 			if a.Tracker != nil {
 				// Notify tracker of all known proxies.
 				for _, p := range r.Proxies {
-					a.Tracker.TrackExpected(a.Lease, p.GetName())
+					a.Tracker.TrackExpected(p.GetName())
 				}
 			}
 		}

lib/reversetunnel/agentpool.go

@@ -59,7 +59,7 @@ type AgentPool struct {
 	spawnLimiter utils.Retry
 
 	mu     sync.Mutex
-	agents map[utils.NetAddr][]*Agent
+	agents []*Agent
 }
 
 // AgentPoolConfig holds configuration parameters for the agent pool
@@ -89,8 +89,8 @@ type AgentPoolConfig struct {
 	Component string
 	// ReverseTunnelServer holds all reverse tunnel connections.
 	ReverseTunnelServer Server
-	// ProxyAddr points to the address of the ssh proxy
-	ProxyAddr string
+	// Resolver retrieves the reverse tunnel address
+	Resolver Resolver
 	// Cluster is a cluster name of the proxy.
 	Cluster string
 	// FIPS indicates if Teleport was started in FIPS mode.
@@ -135,11 +135,6 @@ func NewAgentPool(ctx context.Context, cfg AgentPoolConfig) (*AgentPool, error)
 		return nil, trace.Wrap(err)
 	}
 
-	proxyAddr, err := utils.ParseAddr(cfg.ProxyAddr)
-	if err != nil {
-		return nil, trace.Wrap(err)
-	}
-
 	ctx, cancel := context.WithCancel(ctx)
 	tr, err := track.New(ctx, track.Config{ClusterName: cfg.Cluster})
 	if err != nil {
@@ -148,7 +143,6 @@ func NewAgentPool(ctx context.Context, cfg AgentPoolConfig) (*AgentPool, error)
 	}
 
 	pool := &AgentPool{
-		agents:       make(map[utils.NetAddr][]*Agent),
 		proxyTracker: tr,
 		cfg:          cfg,
 		ctx:          ctx,
@@ -161,7 +155,7 @@ func NewAgentPool(ctx context.Context, cfg AgentPoolConfig) (*AgentPool, error)
 			},
 		}),
 	}
-	pool.proxyTracker.Start(*proxyAddr)
+	pool.proxyTracker.Start()
 	return pool, nil
 }
 
@@ -204,11 +198,11 @@ func (m *AgentPool) processSeekEvents() {
 		// The proxy tracker has given us permission to act on a given
 		// tunnel address
 		case lease := <-m.proxyTracker.Acquire():
-			m.log.Debugf("Seeking: %+v.", lease.Key())
 			m.withLock(func() {
 				// Note that ownership of the lease is transferred to agent
 				// pool for the lifetime of the connection
 				if err := m.addAgent(lease); err != nil {
+					lease.Release()
 					m.log.WithError(err).Errorf("Failed to add agent.")
 				}
 			})
@@ -232,12 +226,7 @@ func (m *AgentPool) withLock(f func()) {
 type matchAgentFn func(a *Agent) bool
 
 func (m *AgentPool) closeAgents() {
-	for key, agents := range m.agents {
-		m.agents[key] = filterAndClose(agents, func(*Agent) bool { return true })
-		if len(m.agents[key]) == 0 {
-			delete(m.agents, key)
-		}
-	}
+	m.agents = filterAndClose(m.agents, func(*Agent) bool { return true })
 }
 
 func filterAndClose(agents []*Agent, matchAgent matchAgentFn) []*Agent {
@@ -246,11 +235,18 @@ func filterAndClose(agents []*Agent, matchAgent matchAgentFn) []*Agent {
 		agent := agents[i]
 		if matchAgent(agent) {
 			agent.log.Debugf("Pool is closing agent.")
-			agent.Close()
+			if err := agent.Close(); err != nil {
+				agent.log.WithError(err).Warnf("Failed to close agent")
+			}
 		} else {
 			filtered = append(filtered, agent)
 		}
 	}
+
+	if len(filtered) <= 0 {
+		return nil
+	}
+
 	return filtered
 }
 
@@ -273,9 +269,13 @@ func (m *AgentPool) pollAndSyncAgents() {
 // transfers into the AgentPool, and will be released when the AgentPool
 // is done with it.
 func (m *AgentPool) addAgent(lease track.Lease) error {
-	addr := lease.Key().(utils.NetAddr)
+	addr, err := m.cfg.Resolver()
+	if err != nil {
+		return trace.Wrap(err)
+	}
+
 	agent, err := NewAgent(AgentConfig{
-		Addr:                addr,
+		Addr:                *addr,
 		ClusterName:         m.cfg.Cluster,
 		Username:            m.cfg.HostUUID,
 		Signer:              m.cfg.HostSigner,
@@ -293,28 +293,25 @@ func (m *AgentPool) addAgent(lease track.Lease) error {
 	})
 	if err != nil {
 		// ensure that lease has been released; OK to call multiple times.
-		lease.Release()
 		return trace.Wrap(err)
 	}
 	m.log.Debugf("Adding %v.", agent)
 	// start the agent in a goroutine. no need to handle Start() errors: Start() will be
 	// retrying itself until the agent is closed
 	go agent.Start()
-	m.agents[addr] = append(m.agents[addr], agent)
+	m.agents = append(m.agents, agent)
 	return nil
 }
 
-// Counts returns a count of the number of proxies a outbound tunnel is
+// Count returns a count of the number of proxies an outbound tunnel is
 // connected to. Used in tests to determine if a proxy has been found and/or
 // removed.
 func (m *AgentPool) Count() int {
 	var out int
 	m.withLock(func() {
-		for _, agents := range m.agents {
-			for _, agent := range agents {
-				if agent.getState() == agentStateConnected {
-					out++
-				}
+		for _, agent := range m.agents {
+			if agent.getState() == agentStateConnected {
+				out++
 			}
 		}
 	})
@@ -325,19 +322,10 @@ func (m *AgentPool) Count() int {
 // removeDisconnected removes disconnected agents from the list of agents.
 // This function should be called under a lock.
 func (m *AgentPool) removeDisconnected() {
-	for agentKey, agentSlice := range m.agents {
-		// Filter and close all disconnected agents.
-		validAgents := filterAndClose(agentSlice, func(agent *Agent) bool {
-			return agent.getState() == agentStateDisconnected
-		})
-
-		// Update (or delete) agent key with filter applied.
-		if len(validAgents) > 0 {
-			m.agents[agentKey] = validAgents
-		} else {
-			delete(m.agents, agentKey)
-		}
-	}
+	// Filter and close all disconnected agents.
+	m.agents = filterAndClose(m.agents, func(agent *Agent) bool {
+		return agent.getState() == agentStateDisconnected
+	})
 }
 
 // Make sure ServerHandlerToListener implements both interfaces.