From 14d72693dd9a5747eefed8f0d1d4f6fb06a3afa5 Mon Sep 17 00:00:00 2001
From: Paul Lecuq <paul@paulbsd.com>
Date: Sun, 9 Apr 2023 15:05:31 +0200
Subject: [PATCH] more stable ipbl websocket feat

---
 src/ws/init.go   |  8 ++++----
 src/ws/pubsub.go | 13 +++++++++++--
 src/ws/reqrep.go | 44 ++++++++++++++++++++++++++------------------
 3 files changed, 41 insertions(+), 24 deletions(-)

diff --git a/src/ws/init.go b/src/ws/init.go
index ea56441..6ffee2b 100644
--- a/src/ws/init.go
+++ b/src/ws/init.go
@@ -11,10 +11,10 @@ import (
 var WebSocketChannel = make(chan models.APIEvent, 100)
 var WebSocketChannelDone = make(chan bool)
 
-var LISTENERS map[string]ConnectionInfo
+var LISTENERS map[string]*ConnectionInfo
 
 func Init(cfg *config.Config) {
-	LISTENERS = make(map[string]ConnectionInfo)
+	LISTENERS = make(map[string]*ConnectionInfo)
 }
 
 func WelcomeAgents(ws *websocket.Conn, welcome WSWelcome, t string) {
@@ -26,13 +26,13 @@ func WelcomeAgents(ws *websocket.Conn, welcome WSWelcome, t string) {
 				ConnectionPS: ws,
 				InitDate:     time.Now(),
 			}
-			LISTENERS[welcome.Hostname] = connectinfo
+			LISTENERS[welcome.Hostname] = &connectinfo
 		} else if t == "rr" {
 			connectinfo := ConnectionInfo{
 				ConnectionRR: ws,
 				InitDate:     time.Now(),
 			}
-			LISTENERS[welcome.Hostname] = connectinfo
+			LISTENERS[welcome.Hostname] = &connectinfo
 		}
 	} else {
 		if t == "ps" {
diff --git a/src/ws/pubsub.go b/src/ws/pubsub.go
index a65a5b5..eeb3b97 100644
--- a/src/ws/pubsub.go
+++ b/src/ws/pubsub.go
@@ -3,6 +3,7 @@ package ws
 import (
 	"encoding/json"
 	"log"
+	"time"
 
 	"git.paulbsd.com/paulbsd/ipbl/src/config"
 	"github.com/labstack/echo/v4"
@@ -19,7 +20,6 @@ func HandleWSPS(c echo.Context, cfg *config.Config) (err error) {
 		err := websocket.Message.Receive(ws, &msg)
 		if err != nil {
 			log.Println(err)
-			return
 		}
 
 		err = json.Unmarshal(msg, &welcome)
@@ -29,6 +29,15 @@ func HandleWSPS(c echo.Context, cfg *config.Config) (err error) {
 			WelcomeAgents(ws, welcome, "ps")
 		}
 
+		err = websocket.Message.Receive(ws, "OK")
+		if err != nil {
+			log.Println(err)
+		}
+
+		for {
+			time.Sleep(50 * time.Millisecond)
+		}
+
 		/*for {
 			var apievent, ok = <-WebSocketChannel
 			if ok {
@@ -39,7 +48,7 @@ func HandleWSPS(c echo.Context, cfg *config.Config) (err error) {
 				}
 				err = websocket.Message.Send(ws, data)
 				if err != nil {
-					gcConn(ws, "ps")
+					//gcConn(ws, "ps")
 					log.Println(LISTENERS)
 					return
 				}
diff --git a/src/ws/reqrep.go b/src/ws/reqrep.go
index 24cf8ea..9815040 100644
--- a/src/ws/reqrep.go
+++ b/src/ws/reqrep.go
@@ -20,7 +20,6 @@ func HandleWSRR(c echo.Context, cfg *config.Config) error {
 		err := websocket.Message.Receive(ws, &msg)
 		if err != nil {
 			log.Println(err)
-			return
 		}
 
 		err = json.Unmarshal(msg, &welcome)
@@ -35,12 +34,10 @@ func HandleWSRR(c echo.Context, cfg *config.Config) error {
 
 			var msg []byte
 			err := websocket.Message.Receive(ws, &msg)
-			log.Println(string(msg))
 			if err != nil {
 				log.Println(err)
 				return
 			}
-			websocket.Message.Send(ws, "OK")
 
 			err = json.Unmarshal(msg, &apievent)
 			if err != nil {
@@ -53,14 +50,14 @@ func HandleWSRR(c echo.Context, cfg *config.Config) error {
 
 			switch apievent.MsgType {
 			case "bootstrap":
-				for _, value := range LISTENERS {
-					err = websocket.Message.Send(value.ConnectionPS, msg)
-					if err != nil {
-						log.Println(err)
+				for index, value := range LISTENERS {
+					if index != apievent.Hostname {
+						err = websocket.Message.Send(value.ConnectionPS, msg)
+						if err != nil {
+							log.Println(err)
+						}
 					}
 				}
-
-				log.Println(LISTENERS)
 			case "add":
 				session := cfg.Db.NewSession()
 				event.APIParse(session, apievent)
@@ -77,22 +74,33 @@ func HandleWSRR(c echo.Context, cfg *config.Config) error {
 						log.Println(err)
 					}
 				}
-				log.Printf("websocket: Inserted event")
-			case "file":
-				apievent.MsgType = "ws"
+				log.Printf("ws: Inserted event")
+			case "init":
+				for _, value := range LISTENERS {
+					err = websocket.Message.Send(value.ConnectionPS, msg)
+					if err != nil {
+						log.Println(err)
+					}
+				}
 			default:
-				log.Println("test")
 			}
 
-			WebSocketChannel <- apievent
-
 			err = websocket.Message.Send(ws, "OK")
 			if err != nil {
-				gcConn(ws, "rr")
-				log.Println(LISTENERS)
-				return
+				log.Println(err)
 			}
 
+			/*
+				WebSocketChannel <- apievent
+
+				err = websocket.Message.Send(ws, "OK")
+				if err != nil {
+					//gcConn(ws, "rr")
+					log.Println(LISTENERS)
+					return
+				}
+			*/
+
 		}
 	}).ServeHTTP(c.Response(), c.Request())
 	return nil