From 6430afcfa53fec6bb5a6cf85e0e66a4ad4d07902 Mon Sep 17 00:00:00 2001
From: fatedier <fatedier@gmail.com>
Date: Tue, 25 Jul 2023 15:12:40 +0800
Subject: [PATCH] fix a goroutine leak issue caused by Login plugin timeout
 (#3547)

---
 Release.md        |  4 ++++
 server/control.go | 36 +++++++++++++++++++++++++-----------
 server/service.go |  2 +-
 3 files changed, 30 insertions(+), 12 deletions(-)

diff --git a/Release.md b/Release.md
index ceaa7b0..8b59baf 100644
--- a/Release.md
+++ b/Release.md
@@ -1,3 +1,7 @@
 ### Features
 
 * Adds a `completion` command for shell completions.
+
+### Fixes
+
+* fix a goroutine leak issue caused by Login plugin timeout.
diff --git a/server/control.go b/server/control.go
index 9764b7c..abf17ab 100644
--- a/server/control.go
+++ b/server/control.go
@@ -29,7 +29,6 @@ import (
 
 	"github.com/fatedier/frp/pkg/auth"
 	"github.com/fatedier/frp/pkg/config"
-	"github.com/fatedier/frp/pkg/consts"
 	pkgerr "github.com/fatedier/frp/pkg/errors"
 	"github.com/fatedier/frp/pkg/msg"
 	plugin "github.com/fatedier/frp/pkg/plugin/server"
@@ -55,13 +54,14 @@ func NewControlManager() *ControlManager {
 	}
 }
 
-func (cm *ControlManager) Add(runID string, ctl *Control) (oldCtl *Control) {
+func (cm *ControlManager) Add(runID string, ctl *Control) (old *Control) {
 	cm.mu.Lock()
 	defer cm.mu.Unlock()
 
-	oldCtl, ok := cm.ctlsByRunID[runID]
+	var ok bool
+	old, ok = cm.ctlsByRunID[runID]
 	if ok {
-		oldCtl.Replaced(ctl)
+		old.Replaced(ctl)
 	}
 	cm.ctlsByRunID[runID] = ctl
 	return
@@ -141,14 +141,13 @@ type Control struct {
 	// replace old controller instantly.
 	runID string
 
-	// control status
-	status string
-
 	readerShutdown  *shutdown.Shutdown
 	writerShutdown  *shutdown.Shutdown
 	managerShutdown *shutdown.Shutdown
 	allShutdown     *shutdown.Shutdown
 
+	started bool
+
 	mu sync.RWMutex
 
 	// Server configuration information
@@ -187,7 +186,6 @@ func NewControl(
 		portsUsedNum:    0,
 		lastPing:        time.Now(),
 		runID:           loginMsg.RunID,
-		status:          consts.Working,
 		readerShutdown:  shutdown.New(),
 		writerShutdown:  shutdown.New(),
 		managerShutdown: shutdown.New(),
@@ -208,11 +206,19 @@ func (ctl *Control) Start() {
 		Error:   "",
 	}
 	_ = msg.WriteMsg(ctl.conn, loginRespMsg)
+	ctl.mu.Lock()
+	ctl.started = true
+	ctl.mu.Unlock()
 
 	go ctl.writer()
-	for i := 0; i < ctl.poolCount; i++ {
-		ctl.sendCh <- &msg.ReqWorkConn{}
-	}
+	go func() {
+		for i := 0; i < ctl.poolCount; i++ {
+			// ignore error here, that means that this control is closed
+			_ = errors.PanicToError(func() {
+				ctl.sendCh <- &msg.ReqWorkConn{}
+			})
+		}
+	}()
 
 	go ctl.manager()
 	go ctl.reader()
@@ -418,6 +424,14 @@ func (ctl *Control) stoper() {
 
 // block until Control closed
 func (ctl *Control) WaitClosed() {
+	ctl.mu.RLock()
+	started := ctl.started
+	ctl.mu.RUnlock()
+
+	if !started {
+		ctl.allShutdown.Done()
+		return
+	}
 	ctl.allShutdown.WaitDone()
 }
 
diff --git a/server/service.go b/server/service.go
index 5c93b66..0024688 100644
--- a/server/service.go
+++ b/server/service.go
@@ -552,7 +552,7 @@ func (svr *Service) RegisterControl(ctlConn net.Conn, loginMsg *msg.Login) (err
 
 	ctl := NewControl(ctx, svr.rc, svr.pxyManager, svr.pluginManager, svr.authVerifier, ctlConn, loginMsg, svr.cfg)
 	if oldCtl := svr.ctlManager.Add(loginMsg.RunID, ctl); oldCtl != nil {
-		oldCtl.allShutdown.WaitDone()
+		oldCtl.WaitClosed()
 	}
 
 	ctl.Start()