diff --git a/bridge/bridge.go b/bridge/bridge.go index 41568dfe9..056077f6a 100755 --- a/bridge/bridge.go +++ b/bridge/bridge.go @@ -1,7 +1,6 @@ package bridge import ( - "ehang.io/nps-mux" "encoding/binary" "errors" "fmt" @@ -12,6 +11,8 @@ import ( "sync" "time" + nps_mux "ehang.io/nps-mux" + "ehang.io/nps/lib/common" "ehang.io/nps/lib/conn" "ehang.io/nps/lib/crypt" @@ -89,9 +90,25 @@ func (s *Bridge) StartTunnel() error { return nil } -//get health information form client +// get health information form client func (s *Bridge) GetHealthFromClient(id int, c *conn.Conn) { for { + // 先读取4字节标志 + buf := make([]byte, 4) + n, err := c.Read(buf) + if err != nil || n != 4 { + break + } + flag := string(buf) + + // 处理心跳包 + if flag == common.HEART_BEAT { + continue + } + + // 处理健康检查信息 + // 把标志写回缓冲区,因为GetHealthInfo会重新读取 + c.Rb = append(c.Rb, buf...) if info, status, err := c.GetHealthInfo(); err != nil { break } else if !status { //the status is true , return target to the targetArr @@ -154,7 +171,7 @@ func (s *Bridge) GetHealthFromClient(id int, c *conn.Conn) { s.DelClient(id) } -//验证失败,返回错误验证flag,并且关闭连接 +// 验证失败,返回错误验证flag,并且关闭连接 func (s *Bridge) verifyError(c *conn.Conn) { c.Write([]byte(common.VERIFY_EER)) } @@ -224,7 +241,7 @@ func (s *Bridge) DelClient(id int) { } } -//use different +// use different func (s *Bridge) typeDeal(typeVal string, c *conn.Conn, id int, vs string) { isPub := file.GetDb().IsPubClient(id) switch typeVal { @@ -303,7 +320,7 @@ func (s *Bridge) typeDeal(typeVal string, c *conn.Conn, id int, vs string) { return } -//register ip +// register ip func (s *Bridge) register(c *conn.Conn) { var hour int32 if err := binary.Read(c, binary.LittleEndian, &hour); err == nil { @@ -387,7 +404,7 @@ func (s *Bridge) ping() { } } -//get config and add task from client config +// get config and add task from client config func (s *Bridge) getConfig(c *conn.Conn, isPub bool, client *file.Client) { var fail bool loop: diff --git a/client/client.go b/client/client.go index e1e02f689..815d972c9 100755 --- a/client/client.go +++ b/client/client.go @@ -3,13 +3,14 @@ package client import ( "bufio" "bytes" - "ehang.io/nps-mux" "net" "net/http" "strconv" "sync" "time" + nps_mux "ehang.io/nps-mux" + "github.com/astaxie/beego/logs" "github.com/xtaci/kcp-go" @@ -33,7 +34,7 @@ type TRPClient struct { once sync.Once } -//new client +// new client func NewRPClient(svraddr string, vKey string, bridgeConnType string, proxyUrl string, cnf *config.Config, disconnectTime int) *TRPClient { return &TRPClient{ svrAddr: svraddr, @@ -50,7 +51,7 @@ func NewRPClient(svraddr string, vKey string, bridgeConnType string, proxyUrl st var NowStatus int var CloseClient bool -//start +// start func (s *TRPClient) Start() { CloseClient = false retry: @@ -84,7 +85,7 @@ retry: s.handleMain() } -//handle main connection +// handle main connection func (s *TRPClient) handleMain() { for { flags, err := s.signal.ReadFlag() @@ -151,7 +152,7 @@ func (s *TRPClient) newUdpConn(localAddr, rAddr string, md5Password string) { } } -//pmux tunnel +// pmux tunnel func (s *TRPClient) newChan() { tunnel, err := NewConn(s.bridgeConnType, s.vKey, s.svrAddr, common.WORK_CHAN, s.proxyUrl) if err != nil { @@ -280,13 +281,18 @@ func (s *TRPClient) handleUdp(serverConn net.Conn) { // Whether the monitor channel is closed func (s *TRPClient) ping() { s.ticker = time.NewTicker(time.Second * 5) -loop: - for { - select { - case <-s.ticker.C: - if s.tunnel != nil && s.tunnel.IsClose { + for range s.ticker.C { + if s.tunnel != nil && s.tunnel.IsClose { + s.Close() + break + } + // 添加signal连接的心跳检测 + if s.signal != nil { + // 发送心跳包 + if _, err := s.signal.Write([]byte(common.HEART_BEAT)); err != nil { + // 连接已断开,关闭客户端 s.Close() - break loop + break } } } diff --git a/lib/common/const.go b/lib/common/const.go index ffb2fa601..944f7a979 100644 --- a/lib/common/const.go +++ b/lib/common/const.go @@ -27,6 +27,7 @@ const ( CONN_TCP = "tcp" CONN_UDP = "udp" CONN_TEST = "TST" + HEART_BEAT = "heat" //heart beat UnauthorizedBytes = `HTTP/1.1 401 Unauthorized Content-Type: text/plain; charset=utf-8 WWW-Authenticate: Basic realm="easyProxy"