From 7f21b733ed466b9492b2b17d811d2844a12bca2e Mon Sep 17 00:00:00 2001
From: Roger Ferdinan <rogerferdian@gmail.com>
Date: Sun, 1 Feb 2026 19:05:33 +0700
Subject: [PATCH] feat: adding ping & pong handler

---
 v1/client/client.go        | 42 ++++++++++++++++++++------------------
 v1/examples/client/main.go |  8 ++++++++
 v1/examples/server/main.go |  4 ++--
 v1/server/server.go        |  1 +
 4 files changed, 33 insertions(+), 22 deletions(-)

diff --git a/v1/client/client.go b/v1/client/client.go
index da5c818..2a86580 100644
--- a/v1/client/client.go
+++ b/v1/client/client.go
@@ -169,7 +169,6 @@ type SafeWebsocketClient struct {
 	doneMap        *safemap.SafeMap[string, chan struct{}]
 
 	writeChan chan Message
-	pongChan  chan error
 }
 
 func (wsClient *SafeWebsocketClient) connect() error {
@@ -196,7 +195,6 @@ func (wsClient *SafeWebsocketClient) connect() error {
 
 	if wsClient.headers != nil {
 		for k, v := range *wsClient.headers {
-			fmt.Println(k, v)
 			header.Set(k, v)
 		}
 	}
@@ -220,23 +218,30 @@ func (wsClient *SafeWebsocketClient) connect() error {
 	wsClient.conn = conn
 	wsClient.isConnected = true
 
+	conn.SetPingHandler(func(pingData string) error {
+		if err := conn.SetReadDeadline(time.Now().Add(readDeadline)); err != nil {
+			log.Printf("error on read deadline: %v\n", err)
+			return err
+		}
+
+		wsClient.writeChan <- Message{
+			MessageType: MessageTypePong,
+			Data:        []byte(pingData),
+		}
+		return nil
+	})
+
+	conn.SetPongHandler(func(pingData string) error {
+		if err := conn.SetReadDeadline(time.Now().Add(readDeadline)); err != nil {
+			log.Printf("error on read deadline: %v\n", err)
+			return err
+		}
+		return nil
+	})
+
 	go wsClient.writePump()
 	go wsClient.readPump()
 
-	// conn.SetPingHandler(func(pingData string) error {
-	// 	wsClient.writeChan <- Message{
-	// 		MessageType: MessageTypePong,
-	// 		Data:        []byte(pingData),
-	// 	}
-
-	// 	select {
-	// 	case err := <-wsClient.pongChan:
-	// 		return err
-	// 	default:
-	// 	}
-	// 	return nil
-	// })
-
 	return nil
 }
 
@@ -265,9 +270,6 @@ func (wsClient *SafeWebsocketClient) writePump() {
 
 			if err := c.WriteMessage(int(data.MessageType), data.Data); err != nil {
 				log.Printf("error on write message: %v\n", err)
-				if data.MessageType == MessageTypePong {
-					wsClient.pongChan <- err
-				}
 				return
 			}
 		}
@@ -450,7 +452,7 @@ func (wsClient *SafeWebsocketClient) Close() error {
 		wsClient.conn.Close()
 	}
 	wsClient.isConnected = false
-	// close(wsClient.dataChannel)
+	close(wsClient.dataChannel)
 
 	return nil
 }
diff --git a/v1/examples/client/main.go b/v1/examples/client/main.go
index 79a7808..359d2ba 100644
--- a/v1/examples/client/main.go
+++ b/v1/examples/client/main.go
@@ -8,10 +8,18 @@ import (
 	"os/signal"
 	"syscall"
 
+	"net/http"
+	_ "net/http/pprof"
+
 	"git.neurocipta.com/rogerferdinan/safe-web-socket/v1/client"
 )
 
 func main() {
+	go func() {
+		log.Println("Starting pprof server on :6060")
+		log.Println(http.ListenAndServe(":6060", nil))
+	}()
+
 	sigChan := make(chan os.Signal, 1)
 	signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM)
 
diff --git a/v1/examples/server/main.go b/v1/examples/server/main.go
index 3b57234..9e34b72 100644
--- a/v1/examples/server/main.go
+++ b/v1/examples/server/main.go
@@ -19,7 +19,7 @@ func main() {
 		BasePort(8080).
 		ApiKey("abcd").
 		HandleFuncWebsocket("/ws/test/", "data_1", func(c chan []byte) {
-			ticker := time.NewTicker(10 * time.Millisecond)
+			ticker := time.NewTicker(100 * time.Millisecond)
 			for range ticker.C {
 				jsonBytes, err := json.Marshal(ExampleData{
 					Time: time.Now(),
@@ -32,7 +32,7 @@ func main() {
 			}
 		}).
 		HandleFuncWebsocket("/ws/test/", "data_2", func(c chan []byte) {
-			ticker := time.NewTicker(10 * time.Millisecond)
+			ticker := time.NewTicker(100 * time.Millisecond)
 			for range ticker.C {
 				jsonBytes, err := json.Marshal(ExampleData{
 					Time: time.Now(),
diff --git a/v1/server/server.go b/v1/server/server.go
index 4a17ccd..60a61e4 100644
--- a/v1/server/server.go
+++ b/v1/server/server.go
@@ -115,6 +115,7 @@ func (s *SafeWebsocketServer) AuthMiddleware(next http.Handler) http.Handler {
 				StatusCode(http.StatusForbidden).
 				Message("X-MBX-APIKEY is missing").
 				Build())
+			return
 		}
 		next.ServeHTTP(w, r)
 	})