From c5f31f660ceb84db00848f902ab70af242b1f9e8 Mon Sep 17 00:00:00 2001
From: Erik Brakkee <erik@brakkee.org>
Date: Thu, 18 Jul 2024 21:18:40 +0200
Subject: [PATCH] now using in wstotcp as well.

---
 cmd/wstotcp/wstotcp.go | 34 ++++++++++------------------------
 1 file changed, 10 insertions(+), 24 deletions(-)

diff --git a/cmd/wstotcp/wstotcp.go b/cmd/wstotcp/wstotcp.go
index bdd107b..a8e401d 100644
--- a/cmd/wstotcp/wstotcp.go
+++ b/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 {
-			_, message, err := conn.ReadMessage()
-			if err != nil {
-				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
-			}
+		_, err := io.Copy(tcpConn, wsConn)
+		if err != nil {
+			log.Printf("Error occurred %v", err)
 		}
 	}()
 
@@ -64,18 +59,9 @@ func handleWebSocket(w http.ResponseWriter, r *http.Request, tcpConn net.Conn) {
 		defer func() {
 			waitChannel <- true
 		}()
-		for {
-			buffer := make([]byte, 1024)
-			n, err := tcpConn.Read(buffer)
-			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
-			}
+		_, err := io.Copy(wsConn, tcpConn)
+		if err != nil {
+			log.Printf("Error occurred %v", err)
 		}
 	}()