MUX: Prevent goroutine leak (#5110)

app/reverse/bridge.go

@@ -9,6 +9,7 @@ import (
 	"github.com/xtls/xray-core/common/mux"
 	"github.com/xtls/xray-core/common/net"
 	"github.com/xtls/xray-core/common/session"
+	"github.com/xtls/xray-core/common/signal"
 	"github.com/xtls/xray-core/common/task"
 	"github.com/xtls/xray-core/features/routing"
 	"github.com/xtls/xray-core/transport"
@@ -53,6 +54,9 @@ func (b *Bridge) cleanup() {
 		if w.IsActive() {
 			activeWorkers = append(activeWorkers, w)
 		}
+		if w.Closed() {
+			w.Timer.SetTimeout(0)
+		}
 	}
 
 	if len(activeWorkers) != len(b.workers) {
@@ -98,6 +102,7 @@ type BridgeWorker struct {
 	Worker     *mux.ServerWorker
 	Dispatcher routing.Dispatcher
 	State      Control_State
+	Timer      *signal.ActivityTimer
 }
 
 func NewBridgeWorker(domain string, tag string, d routing.Dispatcher) (*BridgeWorker, error) {
@@ -125,6 +130,10 @@ func NewBridgeWorker(domain string, tag string, d routing.Dispatcher) (*BridgeWo
 	}
 	w.Worker = worker
 
+	terminate := func() {
+		worker.Close()
+	}
+	w.Timer = signal.CancelAfterInactivity(ctx, terminate, 60*time.Second)
 	return w, nil
 }
 
@@ -144,6 +153,10 @@ func (w *BridgeWorker) IsActive() bool {
 	return w.State == Control_ACTIVE && !w.Worker.Closed()
 }
 
+func (w *BridgeWorker) Closed() bool {
+	return w.Worker.Closed()
+}
+
 func (w *BridgeWorker) Connections() uint32 {
 	return w.Worker.ActiveConnections()
 }
@@ -153,13 +166,20 @@ func (w *BridgeWorker) handleInternalConn(link *transport.Link) {
 	for {
 		mb, err := reader.ReadMultiBuffer()
 		if err != nil {
-			break
+			if w.Closed() {
+				w.Timer.SetTimeout(0)
+			} else {
+				w.Timer.SetTimeout(24 * time.Hour)
+			}
+			return
 		}
+		w.Timer.Update()
 		for _, b := range mb {
 			var ctl Control
 			if err := proto.Unmarshal(b.Bytes(), &ctl); err != nil {
 				errors.LogInfoInner(context.Background(), err, "failed to parse proto message")
-				break
+				w.Timer.SetTimeout(0)
+				return
 			}
 			if ctl.State != w.State {
 				w.State = ctl.State

app/reverse/portal.go

@@ -12,6 +12,7 @@ import (
 	"github.com/xtls/xray-core/common/net"
 	"github.com/xtls/xray-core/common/serial"
 	"github.com/xtls/xray-core/common/session"
+	"github.com/xtls/xray-core/common/signal"
 	"github.com/xtls/xray-core/common/task"
 	"github.com/xtls/xray-core/features/outbound"
 	"github.com/xtls/xray-core/transport"
@@ -159,6 +160,8 @@ func (p *StaticMuxPicker) cleanup() error {
 	for _, w := range p.workers {
 		if !w.Closed() {
 			activeWorkers = append(activeWorkers, w)
+		} else {
+			w.timer.SetTimeout(0)
 		}
 	}
 
@@ -225,6 +228,7 @@ type PortalWorker struct {
 	reader   buf.Reader
 	draining bool
 	counter  uint32
+	timer    *signal.ActivityTimer
 }
 
 func NewPortalWorker(client *mux.ClientWorker) (*PortalWorker, error) {
@@ -244,10 +248,14 @@ func NewPortalWorker(client *mux.ClientWorker) (*PortalWorker, error) {
 	if !f {
 		return nil, errors.New("unable to dispatch control connection")
 	}
+	terminate := func() {
+		client.Close()
+	}
 	w := &PortalWorker{
 		client: client,
 		reader: downlinkReader,
 		writer: uplinkWriter,
+		timer:  signal.CancelAfterInactivity(ctx, terminate, 24*time.Hour), // // prevent leak
 	}
 	w.control = &task.Periodic{
 		Execute:  w.heartbeat,
@@ -274,7 +282,6 @@ func (w *PortalWorker) heartbeat() error {
 		msg.State = Control_DRAIN
 
 		defer func() {
-			w.client.GetTimer().Reset(time.Second * 16)
 			common.Close(w.writer)
 			common.Interrupt(w.reader)
 			w.writer = nil
@@ -286,6 +293,7 @@ func (w *PortalWorker) heartbeat() error {
 		b, err := proto.Marshal(msg)
 		common.Must(err)
 		mb := buf.MergeBytes(nil, b)
+		w.timer.Update()
 		return w.writer.WriteMultiBuffer(mb)
 	}
 	return nil