From 5858b1398fb33c15dd59b8a24a479c9c942aee1f Mon Sep 17 00:00:00 2001
From: Paul Lecuq <paul@paulbsd.com>
Date: Tue, 26 Dec 2023 13:15:11 +0100
Subject: [PATCH] added conn close on ping call

---
 src/ws/pubsub.go |  1 +
 src/ws/reqrep.go | 16 +++++++++++++++-
 2 files changed, 16 insertions(+), 1 deletion(-)

diff --git a/src/ws/pubsub.go b/src/ws/pubsub.go
index 901b1bc..8128cd3 100644
--- a/src/ws/pubsub.go
+++ b/src/ws/pubsub.go
@@ -33,6 +33,7 @@ func HandleWSPS(c *echo.Context, cfg *config.Config) (err error) {
 		err = websocket.Message.Receive(ws, "OK")
 		if err != nil {
 			log.Printf("disconnect: %s (from pubsub channel)\n", welcome.Hostname)
+			ws.Close()
 		}
 
 		for {
diff --git a/src/ws/reqrep.go b/src/ws/reqrep.go
index 6948e7a..b789811 100644
--- a/src/ws/reqrep.go
+++ b/src/ws/reqrep.go
@@ -57,6 +57,7 @@ func HandleWSRR(c *echo.Context, cfg *config.Config) error {
 						err = websocket.Message.Send(value.(*connectionInfo).ConnectionPS, msg)
 						if err != nil {
 							log.Println(err)
+							ws.Close()
 							gcConnOnError(ws)
 						}
 					}
@@ -77,6 +78,7 @@ func HandleWSRR(c *echo.Context, cfg *config.Config) error {
 						err = websocket.Message.Send(value.(*connectionInfo).ConnectionPS, msg)
 						if err != nil {
 							log.Println(err)
+							ws.Close()
 							gcConnOnError(ws)
 						}
 					}
@@ -89,19 +91,31 @@ func HandleWSRR(c *echo.Context, cfg *config.Config) error {
 						err = websocket.Message.Send(value.(*connectionInfo).ConnectionPS, msg)
 						if err != nil {
 							log.Println(err)
+							ws.Close()
 							gcConnOnError(ws)
 						}
 					}
 					return true
 				})
 			case "ping":
-				log.Printf("ping: %s\n", apievent.Hostname)
+				listeners.Range(func(index, value interface{}) bool {
+					if index == apievent.Hostname && value.(*connectionInfo).ConnectionPS != nil {
+						err = websocket.Message.Send(value.(*connectionInfo).ConnectionPS, msg)
+						if err != nil {
+							log.Println(err)
+							ws.Close()
+							gcConnOnError(ws)
+						}
+					}
+					return true
+				})
 			default:
 			}
 
 			err = websocket.Message.Send(ws, "OK")
 			if err != nil {
 				log.Println(err)
+				gcConnOnError(ws)
 			}
 		}
 	}).ServeHTTP((*c).Response(), (*c).Request())