客户端断线重连bug修复

conn.go

@@ -88,12 +88,16 @@ func (s *Conn) ReadFlag() (string, error) {
 
 //读取host
 func (s *Conn) GetHostFromConn() (typeStr string, host string, err error) {
+retry:
 	ltype := make([]byte, 3)
 	_, err = s.Read(ltype)
 	if err != nil {
 		return
 	}
 	typeStr = string(ltype)
+	if typeStr == TEST_FLAG {
+		goto retry
+	}
 	len, err := s.GetLen()
 	if err != nil {
 		return
@@ -200,17 +204,21 @@ func (s *Conn) WriteCompress(b []byte, compress int) (n int, err error) {
 	return
 }
 
-func (s *Conn) wError() {
-	s.conn.Write([]byte(RES_MSG))
+func (s *Conn) wError() (int, error) {
+	return s.conn.Write([]byte(RES_MSG))
 }
-func (s *Conn) wSign() {
-	s.conn.Write([]byte(RES_SIGN))
+func (s *Conn) wSign() (int, error) {
+	return s.conn.Write([]byte(RES_SIGN))
 }
 
-func (s *Conn) wMain() {
-	s.conn.Write([]byte(WORK_MAIN))
+func (s *Conn) wMain() (int, error) {
+	return s.conn.Write([]byte(WORK_MAIN))
 }
 
-func (s *Conn) wChan() {
-	s.conn.Write([]byte(WORK_CHAN))
+func (s *Conn) wChan() (int, error) {
+	return s.conn.Write([]byte(WORK_CHAN))
+}
+
+func (s *Conn) wTest() (int, error) {
+	return s.conn.Write([]byte(TEST_FLAG))
 }

main.go

@@ -10,12 +10,12 @@ var (
 	configPath   = flag.String("config", "config.json", "配置文件路径")
 	tcpPort      = flag.Int("tcpport", 8284, "Socket连接或者监听的端口")
 	httpPort     = flag.Int("httpport", 8024, "当mode为server时为服务端监听端口，当为mode为client时为转发至本地客户端的端口")
-	rpMode       = flag.String("mode", "client", "启动模式，可选为client、server")
+	rpMode       = flag.String("mode", "client", "启动模式，可选为client|server")
 	tunnelTarget = flag.String("target", "10.1.50.203:80", "远程目标")
 	verifyKey    = flag.String("vkey", "", "验证密钥")
-	u            = flag.String("u", "", "sock5验证用户名")
-	p            = flag.String("p", "", "sock5验证密码")
-	compress     = flag.String("compress", "", "数据压缩（gizp|snappy）")
+	u            = flag.String("u", "", "socks5验证用户名")
+	p            = flag.String("p", "", "socks5验证密码")
+	compress     = flag.String("compress", "", "数据压缩（gzip|snappy）")
 	config       Config
 	err          error
 	DataEncode   int

server.go

@@ -15,6 +15,9 @@ const (
 	WORK_CHAN  = "chan"
 	RES_SIGN   = "sign"
 	RES_MSG    = "msg0"
+	TEST_FLAG  = "tst"
+	CONN_TCP   = "tcp"
+	CONN_UDP   = "udp"
 )
 
 type HttpModeServer struct {
@@ -172,11 +175,12 @@ func (s *TunnelModeServer) startTunnelServer() {
 //TODO：这种实现方式……
 //tcp隧道模式
 func ProcessTunnel(c *Conn, s *TunnelModeServer) error {
-retry:
 	link := s.GetTunnel()
-	if _, err := link.WriteHost("tcp", s.tunnelTarget); err != nil {
+	if _, err := link.WriteHost(CONN_TCP, s.tunnelTarget); err != nil {
 		link.Close()
-		goto retry
+		c.Close()
+		log.Println(err)
+		return err
 	}
 	go relay(link, c, DataEncode)
 	relay(c, link, DataDecode)
@@ -190,11 +194,12 @@ func ProcessHttp(c *Conn, s *TunnelModeServer) error {
 		c.Close()
 		return err
 	}
-retry:
 	link := s.GetTunnel()
 	if _, err := link.WriteHost("tcp", addr); err != nil {
+		c.Close()
 		link.Close()
-		goto retry
+		log.Println(err)
+		return err
 	}
 	if method == "CONNECT" {
 		fmt.Fprint(c, "HTTP/1.1 200 Connection established\r\n")

sock5.go

@@ -135,9 +135,9 @@ func (s *Sock5ModeServer) doConnect(c net.Conn, command uint8) (proxyConn *Conn,
 	s.sendReply(c, succeeded)
 	var ltype string
 	if command == associateMethod {
-		ltype = "udp"
+		ltype = CONN_UDP
 	} else {
-		ltype = "tcp"
+		ltype = CONN_TCP
 	}
 	_, err = client.WriteHost(ltype, addr)
 	return client, nil
@@ -154,6 +154,7 @@ func (s *Sock5ModeServer) handleConnect(c net.Conn) {
 	}
 
 }
+
 // passive mode
 func (s *Sock5ModeServer) handleBind(c net.Conn) {
 }

tunnel.go

@@ -100,5 +100,12 @@ func (s *Tunnel) GetTunnel() *Conn {
 	if len(s.tunnelList) < 10 { //新建通道
 		go s.newChan()
 	}
-	return <-s.tunnelList
+retry:
+	c := <-s.tunnelList
+	_, err := c.wTest()
+	if err != nil {
+		c.Close()
+		goto retry
+	}
+	return c
 }

udp.go

@@ -56,7 +56,10 @@ func (s *UdpModeServer) startTunnelServer() {
 
 func (s *UdpModeServer) process(addr *net.UDPAddr, data []byte) {
 	conn := s.GetTunnel()
-	conn.WriteHost("udp", s.tunnelTarget)
+	if _, err := conn.WriteHost(CONN_UDP, s.tunnelTarget);err!=nil{
+		conn.Close()
+		return
+	}
 	go func() {
 		for {
 			buf := make([]byte, 1024)