feat(clientproxy): automatically close boundary connect when proxy is no longer in use (#6232)

Co-authored-by: Andrew Gaffney <andrew@gaffney.cc>

api/proxy/option.go

@@ -39,6 +39,7 @@ type Options struct {
 	WithSkipSessionTeardown      bool
 	withSessionTeardownTimeout   time.Duration
 	withApiClient                *api.Client
+	withInactivityTimeout        time.Duration
 }
 
 // Option is a function that takes in an options struct and sets values or
@@ -142,3 +143,12 @@ func WithApiClient(with *api.Client) Option {
 		return nil
 	}
 }
+
+// WithInactivityTimeout provides an optional duration after which a session
+// with no active connections will be cancelled
+func WithInactivityTimeout(with time.Duration) Option {
+	return func(o *Options) error {
+		o.withInactivityTimeout = with
+		return nil
+	}
+}

api/proxy/proxy.go

@@ -34,26 +34,28 @@ import (
 const sessionCancelTimeout = 30 * time.Second
 
 type ClientProxy struct {
-	tofuToken               string
-	cachedListenerAddress   *ua.String
-	connectionsLeft         *atomic.Int32
-	connsLeftCh             chan int32
-	callerConnectionsLeftCh chan int32
-	apiClient               *api.Client
-	sessionAuthzData        *targets.SessionAuthorizationData
-	createTime              time.Time
-	expiration              time.Time
-	ctx                     context.Context
-	cancel                  context.CancelFunc
-	transport               *http.Transport
-	workerAddr              string
-	listenAddrPort          netip.AddrPort
-	listener                *atomic.Value
-	listenerCloseOnce       *sync.Once
-	clientTlsConf           *tls.Config
-	connWg                  *sync.WaitGroup
-	started                 *atomic.Bool
-	skipSessionTeardown     bool
+	tofuToken             string
+	cachedListenerAddress *ua.String
+	connectionsLeft       *atomic.Int32
+	activeConns           *atomic.Int32
+	connsLeftCh           chan int32
+	callerConnsLeftCh     chan int32
+	apiClient             *api.Client
+	sessionAuthzData      *targets.SessionAuthorizationData
+	createTime            time.Time
+	expiration            time.Time
+	ctx                   context.Context
+	cancel                context.CancelFunc
+	transport             *http.Transport
+	workerAddr            string
+	listenAddrPort        netip.AddrPort
+	listener              *atomic.Value
+	listenerCloseOnce     *sync.Once
+	clientTlsConf         *tls.Config
+	connWg                *sync.WaitGroup
+	started               *atomic.Bool
+	skipSessionTeardown   bool
+	closeReason           *atomic.Value
 }
 
 // New creates a new client proxy. The given context should be cancelable; once
@@ -90,17 +92,19 @@ func New(ctx context.Context, authzToken string, opt ...Option) (*ClientProxy, e
 	}
 
 	p := &ClientProxy{
-		cachedListenerAddress:   ua.NewString(""),
-		connsLeftCh:             make(chan int32),
-		connectionsLeft:         new(atomic.Int32),
-		listener:                new(atomic.Value),
-		listenerCloseOnce:       new(sync.Once),
-		connWg:                  new(sync.WaitGroup),
-		listenAddrPort:          opts.WithListenAddrPort,
-		callerConnectionsLeftCh: opts.WithConnectionsLeftCh,
-		started:                 new(atomic.Bool),
-		skipSessionTeardown:     opts.WithSkipSessionTeardown,
-		apiClient:               opts.withApiClient,
+		cachedListenerAddress: ua.NewString(""),
+		connsLeftCh:           make(chan int32),
+		connectionsLeft:       new(atomic.Int32),
+		activeConns:           new(atomic.Int32),
+		listener:              new(atomic.Value),
+		listenerCloseOnce:     new(sync.Once),
+		connWg:                new(sync.WaitGroup),
+		listenAddrPort:        opts.WithListenAddrPort,
+		callerConnsLeftCh:     opts.WithConnectionsLeftCh,
+		started:               new(atomic.Bool),
+		skipSessionTeardown:   opts.WithSkipSessionTeardown,
+		apiClient:             opts.withApiClient,
+		closeReason:           new(atomic.Value),
 	}
 
 	if opts.WithListener != nil {
@@ -142,7 +146,7 @@ func New(ctx context.Context, authzToken string, opt ...Option) (*ClientProxy, e
 	// We don't _rely_ on client-side timeout verification but this prevents us
 	// seeming to be ready for a connection that will immediately fail when we
 	// try to actually make it
-	p.ctx, p.cancel = context.WithDeadline(ctx, p.expiration)
+	p.ctx, p.cancel = context.WithDeadlineCause(ctx, p.expiration, fmt.Errorf("Session has expired"))
 
 	transport := cleanhttp.DefaultTransport()
 	transport.DisableKeepAlives = false
@@ -212,6 +216,17 @@ func (p *ClientProxy) Start(opt ...Option) (retErr error) {
 	// Ensure closing the listener runs on any other return condition
 	defer listenerCloseFunc()
 
+	// automatically close the proxy when inactive
+	proxyAutoClose := time.AfterFunc(10*time.Minute, func() {
+		p.cancel()
+		p.setCloseReason("Inactivity timeout reached")
+	})
+
+	activeConnCh := make(chan int32)
+	activeConnFn := func(d int32) {
+		activeConnCh <- p.activeConns.Add(d)
+	}
+
 	fin := make(chan error, 10)
 	p.connWg.Add(1)
 	go func() {
@@ -243,8 +258,10 @@ func (p *ClientProxy) Start(opt ...Option) (retErr error) {
 					return
 				}
 			}
+			activeConnFn(1)
 			p.connWg.Add(1)
 			go func() {
+				defer activeConnFn(-1)
 				defer listeningConn.Close()
 				defer p.connWg.Done()
 				wsConn, err := p.getWsConn(p.ctx)
@@ -305,27 +322,40 @@ func (p *ClientProxy) Start(opt ...Option) (retErr error) {
 		}()
 		defer p.connWg.Done()
 		defer listenerCloseFunc()
-
 		for {
 			select {
 			case <-p.ctx.Done():
+				if err := context.Cause(p.ctx); !errors.Is(err, context.Canceled) {
+					p.setCloseReason(err.Error())
+				}
 				return
 			case connsLeft := <-p.connsLeftCh:
 				p.connectionsLeft.Store(connsLeft)
-				if p.callerConnectionsLeftCh != nil {
-					p.callerConnectionsLeftCh <- connsLeft
+				if p.callerConnsLeftCh != nil {
+					p.callerConnsLeftCh <- connsLeft
 				}
 				if connsLeft == 0 {
 					// Close the listener as we can't authorize any more
 					// connections
+					p.setCloseReason("No connections left in session")
 					return
 				}
+			case activeConns := <-activeConnCh:
+				switch {
+				case activeConns > 0:
+					// always stop the timer when a new connection is made,
+					// even if timeout opt is 0
+					proxyAutoClose.Stop()
+				case opts.withInactivityTimeout <= 0:
+					// no timeout was set, timer should not be reset for inactivity
+				case activeConns == 0:
+					proxyAutoClose.Reset(opts.withInactivityTimeout)
+				}
 			}
 		}
 	}()
 
 	p.connWg.Wait()
-	defer p.cancel()
 
 	{
 		// the go funcs are done, so we can safely close the chan and range over any errors
@@ -367,6 +397,25 @@ func (p *ClientProxy) CloseSession(sessionTeardownTimeout time.Duration) error {
 	return nil
 }
 
+// CloseReason returns the reason why the proxy was closed, if the proxy closed
+// itself. If the proxy is still running or the proxy was closed externally, an
+// empty string is returned.
+func (p *ClientProxy) CloseReason() string {
+	switch r := p.closeReason.Load().(type) {
+	case string:
+		return r
+	default:
+		return ""
+	}
+}
+
+// setCloseReason updates the reason the proxy closed from an empty string to the
+// provided string. setCloseReason only accepts the first provided reason for
+// closing, all other calls are ignored.
+func (p *ClientProxy) setCloseReason(reason string) {
+	p.closeReason.CompareAndSwap(nil, reason)
+}
+
 // ListenerAddress returns the address of the client proxy listener. Because the
 // listener is started with Start(), this could be called before listening
 // occurs. To avoid returning until we have a valid value, pass a context;