Bug fix

Author: songquanpeng

handler/client.go

@@ -8,25 +8,30 @@ import (
 )
 
 func PublicPort(localPort, remotePort int) {
+	fmt.Println("Go Public client started.")
 	conn, err := net.Dial("tcp", common.ClientConfig.Host+":"+strconv.Itoa(common.ClientConfig.Port))
 	if err != nil {
 		fmt.Println("Failed to connect to server:", err.Error())
 		return
 	}
-	fmt.Println("Connected to server with remote port", remotePort)
+	fmt.Printf("Connected to %s:%d\n", common.ClientConfig.Host, common.ClientConfig.Port)
 	err = sendHelloPacket(conn, remotePort)
 	if err != nil {
 		fmt.Println("Failed to send hello:", err.Error())
 		return
 	}
-	fmt.Println("Hello packet sent")
-	buf := make([]byte, ConnPacketSize)
+	fmt.Printf("Ready to public local port %d to %s:%d\n", localPort, common.ClientConfig.Host, remotePort)
 	for {
-		_, err := conn.Read(buf)
+		buf := make([]byte, ConnPacketSize) // do not move buf outside this for loop, it brings a bug
+		n, err := conn.Read(buf)
 		if err != nil {
-			fmt.Println(err.Error())
+			fmt.Println("Failed to read connection packet:", err.Error())
+			fmt.Println("Abort.")
 			return
 		}
+		if n != ConnPacketSize {
+			fmt.Println("Warning: mismatched packet size.")
+		}
 		// Should be a connection packet
 		localConn, err := net.Dial("tcp", "localhost:"+strconv.Itoa(localPort))
 		if err != nil {
@@ -38,13 +43,15 @@ func PublicPort(localPort, remotePort int) {
 			fmt.Println("Failed to connect to remote server:", err.Error())
 			continue
 		}
+		fmt.Printf("Connection established: %s <-> %s\n", localConn.LocalAddr().String(), remoteConn.LocalAddr().String())
 		go func() {
 			_, err := remoteConn.Write(buf)
 			if err != nil {
 				fmt.Println("Failed to send connection packet:", err.Error())
+				return
 			}
+			go forward(localConn, remoteConn)
+			go forward(remoteConn, localConn)
 		}()
-		go forward(localConn, remoteConn)
-		go forward(remoteConn, localConn)
 	}
 }

handler/server.go

@@ -8,7 +8,7 @@ import (
 )
 
 func ServeForever() {
-	fmt.Println("Go Public server started.")
+	fmt.Printf("Go Public server started at port %d.\n", common.ServerConfig.Port)
 	addr := net.TCPAddr{
 		IP:   nil,
 		Port: common.ServerConfig.Port,
@@ -19,7 +19,7 @@ func ServeForever() {
 		os.Exit(1)
 	}
 	defer listener.Close()
-	fmt.Printf("Server listening on port %d, waiting for connections...\n", common.ServerConfig.Port)
+	fmt.Printf("Waiting for connections...\n")
 	for {
 		conn, err := listener.Accept()
 		if err != nil {
@@ -68,7 +68,7 @@ func handleHelloPacket(conn net.Conn, buf []byte) {
 			println(err.Error())
 			continue
 		}
-		fmt.Printf("[%d] Connection %p established.\n", port, userConn)
+		//fmt.Printf("[%d] Connection %p established.\n", port, userConn)
 		token := common.GenerateToken()
 		store.add(token, &userConn)
 		err = SendConnPacket(conn, token)
@@ -82,10 +82,12 @@ func handleHelloPacket(conn net.Conn, buf []byte) {
 func handleConnPacket(conn net.Conn, buf []byte) {
 	uuid := common.Bytes2Token(buf[1 : 1+tokenSize])
 	userConn := store.get(uuid)
+	store.remove(uuid)
 	if userConn == nil {
 		fmt.Println("Invalid UUID:", uuid)
 		return
 	}
+	fmt.Printf("Connection established: %s <-> %s\n", conn.RemoteAddr().String(), (*userConn).RemoteAddr().String())
 	go forward(*userConn, conn)
 	forward(conn, *userConn)
 }