fix control delete error

YorkOrz · Jan 26, 2019 · 70ac7d3 · 70ac7d3
1 parent 385c4d3
commit 70ac7d3
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 3 deletions.
diff --git a/server/control.go b/server/control.go
@@ -62,10 +62,13 @@ func (cm *ControlManager) Add(runId string, ctl *Control) (oldCtl *Control) {
 	return
 }
 
-func (cm *ControlManager) Del(runId string) {
+// we should make sure if it's the same control to prevent delete a new one
+func (cm *ControlManager) Del(runId string, ctl *Control) {
 	cm.mu.Lock()
 	defer cm.mu.Unlock()
-	delete(cm.ctlsByRunId, runId)
+	if c, ok := cm.ctlsByRunId[runId]; ok && c == ctl {
+		delete(cm.ctlsByRunId, runId)
+	}
 }
 
 func (cm *ControlManager) GetById(runId string) (ctl *Control, ok bool) {

diff --git a/server/service.go b/server/service.go
@@ -353,7 +353,7 @@ func (svr *Service) RegisterControl(ctlConn frpNet.Conn, loginMsg *msg.Login) (e
 	go func() {
 		// block until control closed
 		ctl.WaitClosed()
-		svr.ctlManager.Del(loginMsg.RunId)
+		svr.ctlManager.Del(loginMsg.RunId, ctl)
 	}()
 	return
 }