(1)分离出连接服务器的逻辑 (2)新增client断线重连功能

This commit is contained in:
Hurricanezwf 2016-02-04 10:55:10 +08:00
parent 793624c379
commit f065562ec3
1 changed files with 79 additions and 36 deletions

View File

@ -1,62 +1,54 @@
package main package main
import ( import (
"encoding/json"
"io" "io"
"sync" "sync"
"encoding/json" "time"
"frp/pkg/models" "frp/pkg/models"
"frp/pkg/utils/conn" "frp/pkg/utils/conn"
"frp/pkg/utils/log" "frp/pkg/utils/log"
) )
// 重连时的间隔时间区间
const (
sleepMinDuration = 1
sleepMaxDuration = 60
)
func ControlProcess(cli *models.ProxyClient, wait *sync.WaitGroup) { func ControlProcess(cli *models.ProxyClient, wait *sync.WaitGroup) {
defer wait.Done() defer wait.Done()
c := &conn.Conn{} c := loginToServer(cli)
err := c.ConnectServer(ServerAddr, ServerPort) if c == nil {
if err != nil { log.Error("ProxyName [%s], connect to server failed!", cli.Name)
log.Error("ProxyName [%s], connect to server [%s:%d] error, %v", cli.Name, ServerAddr, ServerPort, err)
return return
} }
defer c.Close() defer c.Close()
req := &models.ClientCtlReq{
Type: models.ControlConn,
ProxyName: cli.Name,
Passwd: cli.Passwd,
}
buf, _ := json.Marshal(req)
err = c.Write(string(buf) + "\n")
if err != nil {
log.Error("ProxyName [%s], write to server error, %v", cli.Name, err)
return
}
res, err := c.ReadLine()
if err != nil {
log.Error("ProxyName [%s], read from server error, %v", cli.Name, err)
return
}
log.Debug("ProxyName [%s], read [%s]", cli.Name, res)
clientCtlRes := &models.ClientCtlRes{}
if err = json.Unmarshal([]byte(res), &clientCtlRes); 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
}
for { for {
// ignore response content now // ignore response content now
_, err := c.ReadLine() _, err := c.ReadLine()
if err == io.EOF { if err == io.EOF {
// reconnect when disconnect
log.Debug("ProxyName [%s], server close this control conn", cli.Name) log.Debug("ProxyName [%s], server close this control conn", cli.Name)
break var sleepTime time.Duration = 1
for {
log.Debug("ProxyName [%s], try to reconnect to server[%s:%d]...", cli.Name, ServerAddr, ServerPort)
tmpConn := loginToServer(cli)
if tmpConn != nil {
c.Close()
c = tmpConn
break
}
if sleepTime < 60 {
sleepTime++
}
time.Sleep(sleepTime * time.Second)
}
continue
} else if err != nil { } else if err != nil {
log.Warn("ProxyName [%s], read from server error, %v", cli.Name, err) log.Warn("ProxyName [%s], read from server error, %v", cli.Name, err)
continue continue
@ -65,3 +57,54 @@ func ControlProcess(cli *models.ProxyClient, wait *sync.WaitGroup) {
cli.StartTunnel(ServerAddr, ServerPort) cli.StartTunnel(ServerAddr, ServerPort)
} }
} }
func loginToServer(cli *models.ProxyClient) (connection *conn.Conn) {
c := &conn.Conn{}
connection = nil
for i := 0; i < 1; i++ { // ZWF: 此处的for作为控制流使用
err := c.ConnectServer(ServerAddr, ServerPort)
if err != nil {
log.Error("ProxyName [%s], connect to server [%s:%d] error, %v", cli.Name, ServerAddr, ServerPort, err)
break
}
req := &models.ClientCtlReq{
Type: models.ControlConn,
ProxyName: cli.Name,
Passwd: cli.Passwd,
}
buf, _ := json.Marshal(req)
err = c.Write(string(buf) + "\n")
if err != nil {
log.Error("ProxyName [%s], write to server error, %v", cli.Name, err)
break
}
res, err := c.ReadLine()
if err != nil {
log.Error("ProxyName [%s], read from server error, %v", cli.Name, err)
break
}
log.Debug("ProxyName [%s], read [%s]", cli.Name, res)
clientCtlRes := &models.ClientCtlRes{}
if err = json.Unmarshal([]byte(res), &clientCtlRes); err != nil {
log.Error("ProxyName [%s], format server response error, %v", cli.Name, err)
break
}
if clientCtlRes.Code != 0 {
log.Error("ProxyName [%s], start proxy error, %s", cli.Name, clientCtlRes.Msg)
break
}
connection = c
}
if connection == nil {
c.Close()
}
return
}