cyonjan/frp
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix: close connection on worker conn registration error

Browse Source
This commit is contained in:
Guy Lewin 2020-02-27 09:53:56 -05:00
parent 9944c0513b
commit cfcc504aba
2 changed files with 10 additions and 8 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
server/control.go
View File

@ -210,7 +210,7 @@ func (ctl *Control) Start() {
go ctl.stoper() go ctl.stoper()
} }
func (ctl *Control) RegisterWorkConn(conn net.Conn) { func (ctl *Control) RegisterWorkConn(conn net.Conn) error {
xl := ctl.xl xl := ctl.xl
defer func() { defer func() {
if err := recover(); err != nil { if err := recover(); err != nil {
@ -222,9 +222,10 @@ func (ctl *Control) RegisterWorkConn(conn net.Conn) {
select { select {
case ctl.workConnCh <- conn: case ctl.workConnCh <- conn:
xl.Debug("new work connection registered") xl.Debug("new work connection registered")
return nil
default: default:
xl.Debug("work connection pool is full, discarding") xl.Debug("work connection pool is full, discarding")
conn.Close() return fmt.Errorf("work connection pool is full, discarding")
} }
} }

13
server/service.go
View File

@ -332,7 +332,9 @@ func (svr *Service) HandleListener(l net.Listener) {
conn.Close() conn.Close()
} }
case *msg.NewWorkConn: case *msg.NewWorkConn:
svr.RegisterWorkConn(conn, m) if err := svr.RegisterWorkConn(conn, m); err != nil {
conn.Close()
}
case *msg.NewVisitorConn: case *msg.NewVisitorConn:
if err = svr.RegisterVisitorConn(conn, m); err != nil { if err = svr.RegisterVisitorConn(conn, m); err != nil {
xl.Warn("register visitor conn error: %v", err) 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. // 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) xl := frpNet.NewLogFromConn(workConn)
ctl, exist := svr.ctlManager.GetById(newMsg.RunId) ctl, exist := svr.ctlManager.GetById(newMsg.RunId)
if !exist { if !exist {
xl.Warn("No client control found for run id [%s]", newMsg.RunId) 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. // Check auth.
if err := svr.authVerifier.VerifyNewWorkConn(newMsg); err != nil { 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{ msg.WriteMsg(workConn, &msg.StartWorkConn{
Error: "invalid authentication in NewWorkConn", Error: "invalid authentication in NewWorkConn",
}) })
return return fmt.Errorf("invalid authentication in NewWorkConn message on run id [%s]", newMsg.RunId)
} }
ctl.RegisterWorkConn(workConn) return ctl.RegisterWorkConn(workConn)
return
} }
func (svr *Service) RegisterVisitorConn(visitorConn net.Conn, newMsg *msg.NewVisitorConn) error { func (svr *Service) RegisterVisitorConn(visitorConn net.Conn, newMsg *msg.NewVisitorConn) error {