diff --git a/server/control.go b/server/control.go index 7eb5f718..e7d6f3bb 100644 --- a/server/control.go +++ b/server/control.go @@ -210,7 +210,7 @@ func (ctl *Control) Start() { go ctl.stoper() } -func (ctl *Control) RegisterWorkConn(conn net.Conn) { +func (ctl *Control) RegisterWorkConn(conn net.Conn) error { xl := ctl.xl defer func() { if err := recover(); err != nil { @@ -222,9 +222,10 @@ func (ctl *Control) RegisterWorkConn(conn net.Conn) { select { case ctl.workConnCh <- conn: xl.Debug("new work connection registered") + return nil default: xl.Debug("work connection pool is full, discarding") - conn.Close() + return fmt.Errorf("work connection pool is full, discarding") } } diff --git a/server/service.go b/server/service.go index 47e1ba68..d867289a 100644 --- a/server/service.go +++ b/server/service.go @@ -332,7 +332,9 @@ func (svr *Service) HandleListener(l net.Listener) { conn.Close() } case *msg.NewWorkConn: - svr.RegisterWorkConn(conn, m) + if err := svr.RegisterWorkConn(conn, m); err != nil { + conn.Close() + } case *msg.NewVisitorConn: if err = svr.RegisterVisitorConn(conn, m); err != nil { xl.Warn("register visitor conn error: %v", err) @@ -428,12 +430,12 @@ func (svr *Service) RegisterControl(ctlConn net.Conn, loginMsg *msg.Login) (err } // RegisterWorkConn register a new work connection to control and proxies need it. -func (svr *Service) RegisterWorkConn(workConn net.Conn, newMsg *msg.NewWorkConn) { +func (svr *Service) RegisterWorkConn(workConn net.Conn, newMsg *msg.NewWorkConn) error { xl := frpNet.NewLogFromConn(workConn) ctl, exist := svr.ctlManager.GetById(newMsg.RunId) if !exist { xl.Warn("No client control found for run id [%s]", newMsg.RunId) - return + return fmt.Errorf("no client control found for run id [%s]", newMsg.RunId) } // Check auth. if err := svr.authVerifier.VerifyNewWorkConn(newMsg); err != nil { @@ -441,10 +443,9 @@ func (svr *Service) RegisterWorkConn(workConn net.Conn, newMsg *msg.NewWorkConn) msg.WriteMsg(workConn, &msg.StartWorkConn{ Error: "invalid authentication in NewWorkConn", }) - return + return fmt.Errorf("invalid authentication in NewWorkConn message on run id [%s]", newMsg.RunId) } - ctl.RegisterWorkConn(workConn) - return + return ctl.RegisterWorkConn(workConn) } func (svr *Service) RegisterVisitorConn(visitorConn net.Conn, newMsg *msg.NewVisitorConn) error {