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()
}
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")
}
}

13
server/service.go
View File

@ -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 {