forked from Mxmilu666/frp
all: some optimization
This commit is contained in:
fatedier
@@ -28,64 +28,98 @@ import (
|
||||
"frp/utils/log"
|
||||
)
|
||||
|
||||
var connection *conn.Conn = nil
|
||||
var heartBeatTimer *time.Timer = nil
|
||||
|
||||
func ControlProcess(cli *client.ProxyClient, wait *sync.WaitGroup) {
|
||||
defer wait.Done()
|
||||
|
||||
msgSendChan := make(chan interface{}, 1024)
|
||||
|
||||
c, err := loginToServer(cli)
|
||||
if err != nil {
|
||||
log.Error("ProxyName [%s], connect to server failed!", cli.Name)
|
||||
return
|
||||
}
|
||||
connection = c
|
||||
defer connection.Close()
|
||||
defer c.Close()
|
||||
|
||||
go heartbeatSender(c, msgSendChan)
|
||||
|
||||
go msgSender(cli, c, msgSendChan)
|
||||
msgReader(cli, c, msgSendChan)
|
||||
|
||||
close(msgSendChan)
|
||||
}
|
||||
|
||||
// loop for reading messages from frpc after control connection is established
|
||||
func msgReader(cli *client.ProxyClient, c *conn.Conn, msgSendChan chan interface{}) error {
|
||||
// for heartbeat
|
||||
var heartbeatTimeout bool = false
|
||||
timer := time.AfterFunc(time.Duration(client.HeartBeatTimeout)*time.Second, func() {
|
||||
heartbeatTimeout = true
|
||||
c.Close()
|
||||
log.Error("ProxyName [%s], heartbeatRes from frps timeout", cli.Name)
|
||||
})
|
||||
defer timer.Stop()
|
||||
|
||||
for {
|
||||
// ignore response content now
|
||||
content, err := connection.ReadLine()
|
||||
if err == io.EOF || nil == connection || connection.IsClosed() {
|
||||
log.Debug("ProxyName [%s], server close this control conn", cli.Name)
|
||||
var sleepTime time.Duration = 1
|
||||
buf, err := c.ReadLine()
|
||||
if err == io.EOF || c == nil || c.IsClosed() {
|
||||
c.Close()
|
||||
log.Warn("ProxyName [%s], frps close this control conn!", cli.Name)
|
||||
var delayTime time.Duration = 1
|
||||
|
||||
// loop until connect to server
|
||||
// loop until reconnect to frps
|
||||
for {
|
||||
log.Debug("ProxyName [%s], try to reconnect to server[%s:%d]...", cli.Name, client.ServerAddr, client.ServerPort)
|
||||
tmpConn, err := loginToServer(cli)
|
||||
log.Info("ProxyName [%s], try to reconnect to frps [%s:%d]...", cli.Name, client.ServerAddr, client.ServerPort)
|
||||
c, err = loginToServer(cli)
|
||||
if err == nil {
|
||||
connection.Close()
|
||||
connection = tmpConn
|
||||
go heartbeatSender(c, msgSendChan)
|
||||
break
|
||||
}
|
||||
|
||||
if sleepTime < 60 {
|
||||
sleepTime = sleepTime * 2
|
||||
if delayTime < 60 {
|
||||
delayTime = delayTime * 2
|
||||
}
|
||||
time.Sleep(sleepTime * time.Second)
|
||||
time.Sleep(delayTime * time.Second)
|
||||
}
|
||||
continue
|
||||
} else if err != nil {
|
||||
log.Warn("ProxyName [%s], read from server error, %v", cli.Name, err)
|
||||
log.Warn("ProxyName [%s], read from frps error: %v", cli.Name, err)
|
||||
continue
|
||||
}
|
||||
|
||||
clientCtlRes := &msg.ClientCtlRes{}
|
||||
if err := json.Unmarshal([]byte(content), clientCtlRes); err != nil {
|
||||
log.Warn("Parse err: %v : %s", err, content)
|
||||
continue
|
||||
}
|
||||
if consts.SCHeartBeatRes == clientCtlRes.GeneralRes.Code {
|
||||
if heartBeatTimer != nil {
|
||||
log.Debug("Client rcv heartbeat response")
|
||||
heartBeatTimer.Reset(time.Duration(client.HeartBeatTimeout) * time.Second)
|
||||
} else {
|
||||
log.Error("heartBeatTimer is nil")
|
||||
}
|
||||
ctlRes := &msg.ControlRes{}
|
||||
if err := json.Unmarshal([]byte(buf), &ctlRes); err != nil {
|
||||
log.Warn("ProxyName [%s], parse msg from frps error: %v : %s", cli.Name, err, buf)
|
||||
continue
|
||||
}
|
||||
|
||||
cli.StartTunnel(client.ServerAddr, client.ServerPort)
|
||||
switch ctlRes.Type {
|
||||
case consts.HeartbeatRes:
|
||||
log.Debug("ProxyName [%s], receive heartbeat response", cli.Name)
|
||||
timer.Reset(time.Duration(client.HeartBeatTimeout) * time.Second)
|
||||
case consts.NoticeUserConn:
|
||||
log.Debug("ProxyName [%s], new user connection", cli.Name)
|
||||
cli.StartTunnel(client.ServerAddr, client.ServerPort)
|
||||
default:
|
||||
log.Warn("ProxyName [%s}, unsupport msgType [%d]", cli.Name, ctlRes.Type)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// loop for sending messages from channel to frps
|
||||
func msgSender(cli *client.ProxyClient, c *conn.Conn, msgSendChan chan interface{}) {
|
||||
for {
|
||||
msg, ok := <-msgSendChan
|
||||
if !ok {
|
||||
break
|
||||
}
|
||||
|
||||
buf, _ := json.Marshal(msg)
|
||||
err := c.Write(string(buf) + "\n")
|
||||
if err != nil {
|
||||
log.Warn("ProxyName [%s], write to client error, proxy exit", cli.Name)
|
||||
c.Close()
|
||||
break
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -96,8 +130,8 @@ func loginToServer(cli *client.ProxyClient) (c *conn.Conn, err error) {
|
||||
return
|
||||
}
|
||||
|
||||
req := &msg.ClientCtlReq{
|
||||
Type: consts.CtlConn,
|
||||
req := &msg.ControlReq{
|
||||
Type: consts.NewCtlConn,
|
||||
ProxyName: cli.Name,
|
||||
Passwd: cli.Passwd,
|
||||
}
|
||||
@@ -115,53 +149,31 @@ func loginToServer(cli *client.ProxyClient) (c *conn.Conn, err error) {
|
||||
}
|
||||
log.Debug("ProxyName [%s], read [%s]", cli.Name, res)
|
||||
|
||||
clientCtlRes := &msg.ClientCtlRes{}
|
||||
if err = json.Unmarshal([]byte(res), &clientCtlRes); err != nil {
|
||||
ctlRes := &msg.ControlRes{}
|
||||
if err = json.Unmarshal([]byte(res), &ctlRes); err != nil {
|
||||
log.Error("ProxyName [%s], format server response error, %v", cli.Name, err)
|
||||
return
|
||||
}
|
||||
|
||||
if clientCtlRes.Code != 0 {
|
||||
log.Error("ProxyName [%s], start proxy error, %s", cli.Name, clientCtlRes.Msg)
|
||||
return c, fmt.Errorf("%s", clientCtlRes.Msg)
|
||||
if ctlRes.Code != 0 {
|
||||
log.Error("ProxyName [%s], start proxy error, %s", cli.Name, ctlRes.Msg)
|
||||
return c, fmt.Errorf("%s", ctlRes.Msg)
|
||||
}
|
||||
|
||||
go startHeartBeat(c)
|
||||
log.Debug("ProxyName [%s], connect to server[%s:%d] success!", cli.Name, client.ServerAddr, client.ServerPort)
|
||||
|
||||
log.Debug("ProxyName [%s], connect to server [%s:%d] success!", cli.Name, client.ServerAddr, client.ServerPort)
|
||||
return
|
||||
}
|
||||
|
||||
func startHeartBeat(c *conn.Conn) {
|
||||
f := func() {
|
||||
log.Error("HeartBeat timeout!")
|
||||
if c != nil {
|
||||
c.Close()
|
||||
}
|
||||
func heartbeatSender(c *conn.Conn, msgSendChan chan interface{}) {
|
||||
heartbeatReq := &msg.ControlReq{
|
||||
Type: consts.HeartbeatReq,
|
||||
}
|
||||
heartBeatTimer = time.AfterFunc(time.Duration(client.HeartBeatTimeout)*time.Second, f)
|
||||
defer heartBeatTimer.Stop()
|
||||
|
||||
clientCtlReq := &msg.ClientCtlReq{
|
||||
Type: consts.CSHeartBeatReq,
|
||||
ProxyName: "",
|
||||
Passwd: "",
|
||||
}
|
||||
request, err := json.Marshal(clientCtlReq)
|
||||
if err != nil {
|
||||
log.Warn("Serialize clientCtlReq err! Err: %v", err)
|
||||
}
|
||||
|
||||
log.Debug("Start to send heartbeat")
|
||||
log.Info("Start to send heartbeat to frps")
|
||||
for {
|
||||
time.Sleep(time.Duration(client.HeartBeatInterval) * time.Second)
|
||||
if c != nil && !c.IsClosed() {
|
||||
log.Debug("Send heartbeat to server")
|
||||
err = c.Write(string(request) + "\n")
|
||||
if err != nil {
|
||||
log.Error("Send hearbeat to server failed! Err:%v", err)
|
||||
continue
|
||||
}
|
||||
msgSendChan <- heartbeatReq
|
||||
} else {
|
||||
break
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user