From b76965fe528a3d0260582454ac6fc91ab477c9cf Mon Sep 17 00:00:00 2001
From: esrrhs <esrrhs@163.com>
Date: Sat, 22 Dec 2018 16:42:45 +0800
Subject: [PATCH] add

---
 .idea/vcs.xml |  1 +
 client.go     | 41 ++++++++++++++++++++++++++++++++++++++---
 cmd/main.go   |  5 ++++-
 pingtunnel.go |  4 +++-
 server.go     |  6 ++++++
 5 files changed, 52 insertions(+), 5 deletions(-)

diff --git a/.idea/vcs.xml b/.idea/vcs.xml
index 94a25f7..16d3f00 100644
--- a/.idea/vcs.xml
+++ b/.idea/vcs.xml
@@ -2,5 +2,6 @@
 <project version="4">
   <component name="VcsDirectoryMappings">
     <mapping directory="$PROJECT_DIR$" vcs="Git" />
+    <mapping directory="$PROJECT_DIR$/src/github.com/esrrhs/pingtunnel" vcs="Git" />
   </component>
 </project>
\ No newline at end of file
diff --git a/client.go b/client.go
index 92e38e6..1deab02 100644
--- a/client.go
+++ b/client.go
@@ -1,6 +1,7 @@
 package pingtunnel
 
 import (
+	"encoding/binary"
 	"fmt"
 	"golang.org/x/net/icmp"
 	"math"
@@ -9,7 +10,7 @@ import (
 	"time"
 )
 
-func NewClient(addr string, server string, target string, timeout int, sproto int, rproto int) (*Client, error) {
+func NewClient(addr string, server string, target string, timeout int, sproto int, rproto int, hb int) (*Client, error) {
 
 	ipaddr, err := net.ResolveUDPAddr("udp", addr)
 	if err != nil {
@@ -32,6 +33,7 @@ func NewClient(addr string, server string, target string, timeout int, sproto in
 		timeout:      timeout,
 		sproto:       sproto,
 		rproto:       rproto,
+		hb:           hb,
 	}, nil
 }
 
@@ -42,6 +44,7 @@ type Client struct {
 	timeout int
 	sproto  int
 	rproto  int
+	hb      int
 
 	ipaddr *net.UDPAddr
 	addr   string
@@ -61,6 +64,8 @@ type Client struct {
 	recvPacket     uint64
 	sendPacketSize uint64
 	recvPacketSize uint64
+
+	sendHBPacket uint64
 }
 
 type ClientConn struct {
@@ -119,12 +124,24 @@ func (p *Client) Run() {
 	interval := time.NewTicker(time.Second)
 	defer interval.Stop()
 
+	inter := 1000
+	if p.hb > 0 {
+		inter := 1000 / p.hb
+		if inter <= 0 {
+			inter = 1
+		}
+	}
+	intervalHB := time.NewTicker(time.Millisecond * (time.Duration)(inter))
+	defer intervalHB.Stop()
+
 	for {
 		select {
 		case <-interval.C:
 			p.checkTimeoutConn()
 			p.ping()
 			p.showNet()
+		case <-intervalHB.C:
+			p.heartbeat()
 		case r := <-recv:
 			p.processPacket(r)
 		}
@@ -235,7 +252,7 @@ func (p *Client) checkTimeoutConn() {
 }
 
 func (p *Client) ping() {
-	if p.sendPacket == 0 {
+	if p.sendPacket == 0 && p.recvPacket == 0 {
 		now := time.Now()
 		b, _ := now.MarshalBinary()
 		sendICMP(p.id, p.sequence, *p.conn, p.ipaddrServer, p.targetAddr, "", (uint32)(PING), b, p.sproto, p.rproto)
@@ -245,9 +262,27 @@ func (p *Client) ping() {
 }
 
 func (p *Client) showNet() {
-	fmt.Printf("send %dPacket/s %dKB/s recv %dPacket/s %dKB/s\n", p.sendPacket, p.sendPacketSize/1024, p.recvPacket, p.recvPacketSize/1024)
+	fmt.Printf("send %dPacket/s %dKB/s recv %dPacket/s %dKB/s HB %d/s\n",
+		p.sendPacket, p.sendPacketSize/1024, p.recvPacket, p.recvPacketSize/1024, p.sendHBPacket)
 	p.sendPacket = 0
 	p.recvPacket = 0
 	p.sendPacketSize = 0
 	p.recvPacketSize = 0
+	p.sendHBPacket = 0
+}
+
+func (p *Client) heartbeat() {
+
+	if p.hb > 0 {
+		for _, conn := range p.localIdToConnMap {
+
+			b := make([]byte, 4)
+			binary.BigEndian.PutUint32(b[:4], rand.Uint32())
+			sendICMP(p.id, p.sequence, *p.conn, p.ipaddrServer, p.targetAddr, conn.id, (uint32)(HB), b, p.sproto, p.rproto)
+			p.sequence++
+
+			p.sendHBPacket++
+
+		}
+	}
 }
diff --git a/cmd/main.go b/cmd/main.go
index 97ffac0..25ca475 100644
--- a/cmd/main.go
+++ b/cmd/main.go
@@ -36,6 +36,8 @@ Usage:
 
     -rproto   客户端接收ping协议的协议，默认是0
               The protocol that the client receives the ping. The default is 0.
+
+    -hb       客户端保持每秒发到服务器的心跳包，用于在某些网络下更新服务器的id和seq，以接收到服务器的reply
 `
 
 func main() {
@@ -47,6 +49,7 @@ func main() {
 	timeout := flag.Int("timeout", 60, "conn timeout")
 	sproto := flag.Int("sproto", 8, "send ping proto")
 	rproto := flag.Int("rproto", 0, "recv ping proto")
+	hb := flag.Int("hb", 0, "client heartbeat")
 	flag.Usage = func() {
 		fmt.Printf(usage)
 	}
@@ -76,7 +79,7 @@ func main() {
 		fmt.Printf("server %s\n", *server)
 		fmt.Printf("target %s\n", *target)
 
-		c, err := pingtunnel.NewClient(*listen, *server, *target, *timeout, *sproto, *rproto)
+		c, err := pingtunnel.NewClient(*listen, *server, *target, *timeout, *sproto, *rproto, *hb)
 		if err != nil {
 			fmt.Printf("ERROR: %s\n", err.Error())
 			return
diff --git a/pingtunnel.go b/pingtunnel.go
index fd1c209..646d46b 100644
--- a/pingtunnel.go
+++ b/pingtunnel.go
@@ -18,6 +18,7 @@ import (
 const (
 	DATA uint32 = 0x01010101
 	PING uint32 = 0x02020202
+	HB   uint32 = 0x03030303
 	END  uint32 = 0xAAAABBBB
 )
 
@@ -197,7 +198,8 @@ func recvICMP(conn icmp.PacketConn, recv chan<- *Packet) {
 		}
 		my.Unmarshal(bytes[8:n])
 
-		if (my.TYPE != (uint32)(DATA) && my.TYPE != (uint32)(PING)) || my.ENDTYPE != (uint32)(END) {
+		if (my.TYPE != (uint32)(DATA) && my.TYPE != (uint32)(PING) && my.TYPE != (uint32)(HB)) ||
+			my.ENDTYPE != (uint32)(END) {
 			//fmt.Printf("processPacket diff type %s %d %d \n", my.ID, my.TYPE, my.ENDTYPE)
 			continue
 		}
diff --git a/server.go b/server.go
index 1ccd7e9..d4f462c 100644
--- a/server.go
+++ b/server.go
@@ -104,6 +104,12 @@ func (p *Server) processPacket(packet *Packet) {
 	udpConn.echoId = packet.echoId
 	udpConn.echoSeq = packet.echoSeq
 
+	if packet.msgType == HB {
+		udpConn.echoId = packet.echoId
+		udpConn.echoSeq = packet.echoSeq
+		return
+	}
+
 	_, err := udpConn.conn.Write(packet.data)
 	if err != nil {
 		fmt.Printf("WriteToUDP Error %s\n", err)