now using in wstotcp as well.

cmd/wstotcp/wstotcp.go

@@ -1,8 +1,10 @@
 package main
 
 import (
+	"cidebug/pkg/iowrappers"
 	"fmt"
 	"github.com/gorilla/websocket"
+	"io"
 	"log"
 	"net"
 	"net/http"
@@ -33,11 +35,12 @@ func main() {
 
 func handleWebSocket(w http.ResponseWriter, r *http.Request, tcpConn net.Conn) {
 	conn, err := upgrader.Upgrade(w, r, nil)
+	wsConn := iowrappers.NewWebSocketConn(conn)
 	if err != nil {
 		log.Println("Error upgrading to WebSocket:", err)
 		return
 	}
-	defer conn.Close()
+	defer wsConn.Close()
 
 	waitChannel := make(chan bool)
 
@@ -46,17 +49,9 @@ func handleWebSocket(w http.ResponseWriter, r *http.Request, tcpConn net.Conn) {
 		defer func() {
 			waitChannel <- true
 		}()
-		for {
+		_, err := io.Copy(tcpConn, wsConn)
-			_, message, err := conn.ReadMessage()
+		if err != nil {
-			if err != nil {
+			log.Printf("Error occurred %v", err)
-				log.Println("Error reading WebSocket message:", err)
-				return
-			}
-			_, err = tcpConn.Write(message)
-			if err != nil {
-				log.Println("Error sending message to TCP server:", err)
-				return
-			}
 		}
 	}()
 
@@ -64,18 +59,9 @@ func handleWebSocket(w http.ResponseWriter, r *http.Request, tcpConn net.Conn) {
 		defer func() {
 			waitChannel <- true
 		}()
-		for {
+		_, err := io.Copy(wsConn, tcpConn)
-			buffer := make([]byte, 1024)
+		if err != nil {
-			n, err := tcpConn.Read(buffer)
+			log.Printf("Error occurred %v", err)
-			if err != nil {
-				log.Println("Error reading from TCP server:", err)
-				return
-			}
-			err = conn.WriteMessage(websocket.TextMessage, buffer[:n])
-			if err != nil {
-				log.Println("Error writing WebSocket message:", err)
-				return
-			}
 		}
 	}()