updated ws feature

cmd/ipbl/ipbl.go

@@ -10,6 +10,7 @@ import (
 	"git.paulbsd.com/paulbsd/ipbl/src/database"
 	"git.paulbsd.com/paulbsd/ipbl/src/models"
 	"git.paulbsd.com/paulbsd/ipbl/src/routers"
+	"git.paulbsd.com/paulbsd/ipbl/src/ws"
 	"git.paulbsd.com/paulbsd/ipbl/utils"
 	_ "github.com/lib/pq"
 )
@@ -36,6 +37,7 @@ func main() {
 		go models.ScanIP(&cfg)
 	}
 	//go zmqrouter.Init(&cfg)
+	go ws.Init(&cfg)
 	go func() { err = routers.RunServer(&ctx, &cfg) }()
 	if err != nil {
 		log.Fatalln(err)

src/routers/funcs.go

@@ -155,6 +155,5 @@ func RegisterRoutes(e *echo.Echo, ctx *context.Context, cfg *config.Config) {
 		return Result(c, err, disc)
 	})
 	e.File("/test.html", "/home/paul/test.html")
-	e.GET("/wsps", ws.HandleWSPubSub)
-	e.GET("/wsrr", ws.HandleWSReqRep)
+	e.GET("/ws", ws.HandleWS)
 }

src/ws/main.go

@@ -1 +1,53 @@
 package ws
+
+import (
+	"encoding/json"
+	"fmt"
+	"log"
+
+	"git.paulbsd.com/paulbsd/ipbl/src/config"
+	"git.paulbsd.com/paulbsd/ipbl/src/models"
+	"github.com/labstack/echo/v4"
+	"golang.org/x/net/websocket"
+)
+
+var LISTENERS map[string]*websocket.Conn
+var CHAN int
+
+func Init(cfg *config.Config) {
+	LISTENERS = make(map[string]*websocket.Conn)
+}
+
+func HandleWS(c echo.Context) error {
+	websocket.Handler(func(ws *websocket.Conn) {
+		defer ws.Close()
+
+		fmt.Printf("rr: %v\n", ws)
+
+		for {
+			// Read
+			var event models.APIEvent
+			var msg []byte
+			err := websocket.Message.Receive(ws, &msg)
+			if err != nil {
+				return
+			}
+			err = json.Unmarshal(msg, &event)
+			if err != nil {
+				log.Println(err)
+			}
+			if event.MsgType == "bootstrap" {
+				LISTENERS[event.Hostname] = ws
+			}
+			fmt.Println(LISTENERS)
+
+			// Write
+			err = websocket.Message.Send(ws, "OK")
+			if err != nil {
+				return
+			}
+
+		}
+	}).ServeHTTP(c.Response(), c.Request())
+	return nil
+}

src/ws/pubsub.go

@@ -1,27 +0,0 @@
-package ws
-
-import (
-	"github.com/labstack/echo/v4"
-	"golang.org/x/net/websocket"
-)
-
-func HandleWSPubSub(c echo.Context) error {
-	websocket.Handler(func(ws *websocket.Conn) {
-		defer ws.Close()
-		for {
-			// Write
-			err := websocket.Message.Send(ws, "Hello, Client!")
-			if err != nil {
-				return
-			}
-
-			// Read
-			msg := ""
-			err = websocket.Message.Receive(ws, &msg)
-			if err != nil {
-				return
-			}
-		}
-	}).ServeHTTP(c.Response(), c.Request())
-	return nil
-}

src/ws/reqrep.go

@@ -1,31 +0,0 @@
-package ws
-
-import (
-	"fmt"
-
-	"github.com/labstack/echo/v4"
-	"golang.org/x/net/websocket"
-)
-
-func HandleWSReqRep(c echo.Context) error {
-	websocket.Handler(func(ws *websocket.Conn) {
-		defer ws.Close()
-		for {
-			// Read
-			msg := ""
-			err := websocket.Message.Receive(ws, &msg)
-			if err != nil {
-				return
-			}
-			fmt.Println(msg)
-
-			// Write
-			err = websocket.Message.Send(ws, "OK")
-			if err != nil {
-				return
-			}
-
-		}
-	}).ServeHTTP(c.Response(), c.Request())
-	return nil
-}