From 160d6efc8984cfcfc613b8f8ff302dfabd5171a0 Mon Sep 17 00:00:00 2001
From: esrrhs <esrrhs@163.com>
Date: Wed, 30 Oct 2019 19:45:19 +0800
Subject: [PATCH] add

---
 client.go     | 26 ++++++++++++++++++++++----
 pingtunnel.go |  8 ++++++--
 server.go     | 20 ++++++++++++++++----
 3 files changed, 44 insertions(+), 10 deletions(-)

diff --git a/client.go b/client.go
index c78808b..563868e 100644
--- a/client.go
+++ b/client.go
@@ -8,6 +8,7 @@ import (
 	"math"
 	"math/rand"
 	"net"
+	"sync"
 	"time"
 )
 
@@ -65,9 +66,10 @@ func NewClient(addr string, server string, target string, timeout int, key int,
 }
 
 type Client struct {
-	exit     bool
-	rtt      time.Duration
-	interval *time.Ticker
+	exit           bool
+	rtt            time.Duration
+	interval       *time.Ticker
+	workResultLock sync.WaitGroup
 
 	id       int
 	sequence int
@@ -176,11 +178,14 @@ func (p *Client) Run() error {
 	}
 
 	recv := make(chan *Packet, 10000)
-	go recvICMP(*p.conn, recv)
+	go recvICMP(&p.workResultLock, &p.exit, *p.conn, recv)
 
 	p.interval = time.NewTicker(time.Second)
 
 	go func() {
+		p.workResultLock.Add(1)
+		defer p.workResultLock.Done()
+
 		for !p.exit {
 			select {
 			case <-p.interval.C:
@@ -198,6 +203,7 @@ func (p *Client) Run() error {
 
 func (p *Client) Stop() {
 	p.exit = true
+	p.workResultLock.Wait()
 	p.conn.Close()
 	if p.tcplistenConn != nil {
 		p.tcplistenConn.Close()
@@ -210,6 +216,9 @@ func (p *Client) Stop() {
 
 func (p *Client) AcceptTcp() error {
 
+	p.workResultLock.Add(1)
+	defer p.workResultLock.Done()
+
 	loggo.Info("client waiting local accept tcp")
 
 	for !p.exit {
@@ -237,6 +246,9 @@ func (p *Client) AcceptTcp() error {
 
 func (p *Client) AcceptTcpConn(conn *net.TCPConn, targetAddr string) {
 
+	p.workResultLock.Add(1)
+	defer p.workResultLock.Done()
+
 	uuid := UniqueId()
 	tcpsrcaddr := conn.RemoteAddr().(*net.TCPAddr)
 
@@ -426,6 +438,9 @@ func (p *Client) AcceptTcpConn(conn *net.TCPConn, targetAddr string) {
 
 func (p *Client) Accept() error {
 
+	p.workResultLock.Add(1)
+	defer p.workResultLock.Done()
+
 	loggo.Info("client waiting local accept udp")
 
 	bytes := make([]byte, 10240)
@@ -580,6 +595,9 @@ func (p *Client) showNet() {
 
 func (p *Client) AcceptSock5Conn(conn *net.TCPConn) {
 
+	p.workResultLock.Add(1)
+	defer p.workResultLock.Done()
+
 	var err error = nil
 	if err = sock5Handshake(conn); err != nil {
 		loggo.Error("socks handshake: %s", err)
diff --git a/pingtunnel.go b/pingtunnel.go
index ae078e5..4681002 100644
--- a/pingtunnel.go
+++ b/pingtunnel.go
@@ -12,6 +12,7 @@ import (
 	"golang.org/x/net/ipv4"
 	"io"
 	"net"
+	"sync"
 	"syscall"
 	"time"
 )
@@ -77,10 +78,13 @@ func sendICMP(id int, sequence int, conn icmp.PacketConn, server *net.IPAddr, ta
 	return
 }
 
-func recvICMP(conn icmp.PacketConn, recv chan<- *Packet) {
+func recvICMP(workResultLock *sync.WaitGroup, exit *bool, conn icmp.PacketConn, recv chan<- *Packet) {
+
+	(*workResultLock).Add(1)
+	defer (*workResultLock).Done()
 
 	bytes := make([]byte, 10240)
-	for {
+	for !*exit {
 		conn.SetReadDeadline(time.Now().Add(time.Millisecond * 100))
 		n, srcaddr, err := conn.ReadFrom(bytes)
 
diff --git a/server.go b/server.go
index 78b228b..4fd79a8 100644
--- a/server.go
+++ b/server.go
@@ -6,6 +6,7 @@ import (
 	"github.com/golang/protobuf/proto"
 	"golang.org/x/net/icmp"
 	"net"
+	"sync"
 	"time"
 )
 
@@ -17,9 +18,10 @@ func NewServer(key int) (*Server, error) {
 }
 
 type Server struct {
-	exit     bool
-	key      int
-	interval *time.Ticker
+	exit           bool
+	key            int
+	interval       *time.Ticker
+	workResultLock sync.WaitGroup
 
 	conn *icmp.PacketConn
 
@@ -61,11 +63,14 @@ func (p *Server) Run() error {
 	p.localConnMap = make(map[string]*ServerConn)
 
 	recv := make(chan *Packet, 10000)
-	go recvICMP(*p.conn, recv)
+	go recvICMP(&p.workResultLock, &p.exit, *p.conn, recv)
 
 	p.interval = time.NewTicker(time.Second)
 
 	go func() {
+		p.workResultLock.Add(1)
+		defer p.workResultLock.Done()
+
 		for !p.exit {
 			select {
 			case <-p.interval.C:
@@ -82,6 +87,7 @@ func (p *Server) Run() error {
 
 func (p *Server) Stop() {
 	p.exit = true
+	p.workResultLock.Wait()
 	p.conn.Close()
 	p.interval.Stop()
 }
@@ -193,6 +199,9 @@ func (p *Server) processPacket(packet *Packet) {
 
 func (p *Server) RecvTCP(conn *ServerConn, id string, src *net.IPAddr) {
 
+	p.workResultLock.Add(1)
+	defer p.workResultLock.Done()
+
 	loggo.Info("server waiting target response %s -> %s %s", conn.tcpaddrTarget.String(), conn.id, conn.tcpconn.LocalAddr().String())
 
 	loggo.Info("start wait remote connect tcp %s %s", conn.id, conn.tcpaddrTarget.String())
@@ -370,6 +379,9 @@ func (p *Server) RecvTCP(conn *ServerConn, id string, src *net.IPAddr) {
 
 func (p *Server) Recv(conn *ServerConn, id string, src *net.IPAddr) {
 
+	p.workResultLock.Add(1)
+	defer p.workResultLock.Done()
+
 	loggo.Info("server waiting target response %s -> %s %s", conn.ipaddrTarget.String(), conn.id, conn.conn.LocalAddr().String())
 
 	for !p.exit {