TRA-3437 switch fiber and ikisocket with gin-gonic and gorilla websocket (#136)

* WIP

* WIP

* WIP

* Update socket_server_handlers.go and socket_routes.go

* Fix stuck sockets

* Update go.mod, go.sum, and 5 more files...

* Update socket_routes.go

* Update Dockerfile, go.sum, and fiber_middleware.go

* fix analyze

Co-authored-by: RamiBerm <rami.berman@up9.com>

agent/pkg/routes/entries_routes.go

@@ -1,25 +1,25 @@
 package routes
 
 import (
-	"github.com/gofiber/fiber/v2"
+	"github.com/gin-gonic/gin"
 	"mizuserver/pkg/controllers"
 )
 
 // EntriesRoutes defines the group of har entries routes.
-func EntriesRoutes(fiberApp *fiber.App) {
-	routeGroup := fiberApp.Group("/api")
+func EntriesRoutes(ginApp *gin.Engine) {
+	routeGroup := ginApp.Group("/api")
 
-	routeGroup.Get("/entries", controllers.GetEntries)        // get entries (base/thin entries)
-	routeGroup.Get("/entries/:entryId", controllers.GetEntry) // get single (full) entry
-	routeGroup.Get("/exportEntries", controllers.GetFullEntries)
-	routeGroup.Get("/uploadEntries", controllers.UploadEntries)
-	routeGroup.Get("/resolving", controllers.GetCurrentResolvingInformation)
+	routeGroup.GET("/entries", controllers.GetEntries)        // get entries (base/thin entries)
+	routeGroup.GET("/entries/:entryId", controllers.GetEntry) // get single (full) entry
+	routeGroup.GET("/exportEntries", controllers.GetFullEntries)
+	routeGroup.GET("/uploadEntries", controllers.UploadEntries)
+	routeGroup.GET("/resolving", controllers.GetCurrentResolvingInformation)
 
-	routeGroup.Get("/har", controllers.GetHARs)
+	routeGroup.GET("/har", controllers.GetHARs)
 
-	routeGroup.Get("/resetDB", controllers.DeleteAllEntries)     // get single (full) entry
-	routeGroup.Get("/generalStats", controllers.GetGeneralStats) // get general stats about entries in DB
+	routeGroup.GET("/resetDB", controllers.DeleteAllEntries)     // get single (full) entry
+	routeGroup.GET("/generalStats", controllers.GetGeneralStats) // get general stats about entries in DB
 
-	routeGroup.Get("/tapStatus", controllers.GetTappingStatus) // get tapping status
-	routeGroup.Get("/analyzeStatus", controllers.AnalyzeInformation)
+	routeGroup.GET("/tapStatus", controllers.GetTappingStatus) // get tapping status
+	routeGroup.GET("/analyzeStatus", controllers.AnalyzeInformation)
 }

agent/pkg/routes/metadata_routes.go

@@ -1,13 +1,13 @@
 package routes
 
 import (
-	"github.com/gofiber/fiber/v2"
+	"github.com/gin-gonic/gin"
 	"mizuserver/pkg/controllers"
 )
 
 // MetadataRoutes defines the group of metadata routes.
-func MetadataRoutes(fiberApp *fiber.App) {
-	routeGroup := fiberApp.Group("/metadata")
+func MetadataRoutes(app *gin.Engine) {
+	routeGroup := app.Group("/metadata")
 
-	routeGroup.Get("/version", controllers.GetVersion)
+	routeGroup.GET("/version", controllers.GetVersion)
 }

agent/pkg/routes/not_found_route.go

@@ -1,12 +1,15 @@
 package routes
 
-import "github.com/gofiber/fiber/v2"
+import (
+	"github.com/gin-gonic/gin"
+	"net/http"
+)
 
 // NotFoundRoute defines the 404 Error route.
-func NotFoundRoute(fiberApp *fiber.App) {
-	fiberApp.Use(
-		func(c *fiber.Ctx) error {
-			return c.Status(fiber.StatusNotFound).JSON(fiber.Map{
+func NotFoundRoute(app *gin.Engine) {
+	app.Use(
+		func(c *gin.Context) {
+			c.JSON(http.StatusNotFound, map[string]interface{}{
 				"error": true,
 				"msg":   "sorry, endpoint is not found",
 			})

agent/pkg/routes/socket_routes.go

@@ -1,31 +1,118 @@
 package routes
 
 import (
-	"github.com/antoniodipinto/ikisocket"
-	"github.com/gofiber/fiber/v2"
+	"errors"
+	"fmt"
+	"github.com/gin-gonic/gin"
+	"github.com/gorilla/websocket"
+	"github.com/up9inc/mizu/shared/debounce"
+	"net/http"
+	"sync"
+	"time"
 )
 
 type EventHandlers interface {
-	WebSocketConnect(ep *ikisocket.EventPayload)
-	WebSocketDisconnect(ep *ikisocket.EventPayload)
-	WebSocketClose(ep *ikisocket.EventPayload)
-	WebSocketError(ep *ikisocket.EventPayload)
-	WebSocketMessage(ep *ikisocket.EventPayload)
+	WebSocketConnect(socketId int, isTapper bool)
+	WebSocketDisconnect(socketId int, isTapper bool)
+	WebSocketMessage(socketId int, message []byte)
 }
 
-func WebSocketRoutes(app *fiber.App, eventHandlers EventHandlers) {
-	app.Get("/ws", ikisocket.New(func(kws *ikisocket.Websocket) {
-		kws.SetAttribute("is_tapper", false)
-	}))
-
-	app.Get("/wsTapper", ikisocket.New(func(kws *ikisocket.Websocket) {
-		// Tapper clients are handled differently, they don't need to receive new message broadcasts.
-		kws.SetAttribute("is_tapper", true)
-	}))
-
-	ikisocket.On(ikisocket.EventMessage, eventHandlers.WebSocketMessage)
-	ikisocket.On(ikisocket.EventConnect, eventHandlers.WebSocketConnect)
-	ikisocket.On(ikisocket.EventDisconnect, eventHandlers.WebSocketDisconnect)
-	ikisocket.On(ikisocket.EventClose, eventHandlers.WebSocketClose) // This event is called when the server disconnects the user actively with .Close() method
-	ikisocket.On(ikisocket.EventError, eventHandlers.WebSocketError) // On error event
+type SocketConnection struct {
+	connection *websocket.Conn
+	lock *sync.Mutex
+	eventHandlers EventHandlers
+	isTapper bool
+}
+
+var websocketUpgrader = websocket.Upgrader{
+	ReadBufferSize:  1024,
+	WriteBufferSize: 1024,
+}
+
+var websocketIdsLock = sync.Mutex{}
+var connectedWebsockets map[int]*SocketConnection
+var connectedWebsocketIdCounter = 0
+
+func init() {
+	websocketUpgrader.CheckOrigin = func(r *http.Request) bool { return true } // like cors for web socket
+	connectedWebsockets = make(map[int]*SocketConnection, 0)
+}
+
+func WebSocketRoutes(app *gin.Engine, eventHandlers EventHandlers) {
+	app.GET("/ws", func(c *gin.Context) {
+		websocketHandler(c.Writer, c.Request, eventHandlers, false)
+	})
+	app.GET("/wsTapper", func(c *gin.Context) {
+		websocketHandler(c.Writer, c.Request, eventHandlers, true)
+	})
+}
+
+func websocketHandler(w http.ResponseWriter, r *http.Request, eventHandlers EventHandlers, isTapper bool) {
+	conn, err := websocketUpgrader.Upgrade(w, r, nil)
+	if err != nil {
+		fmt.Println("Failed to set websocket upgrade: %+v", err)
+		return
+	}
+
+	websocketIdsLock.Lock()
+
+	connectedWebsocketIdCounter++
+	socketId := connectedWebsocketIdCounter
+	connectedWebsockets[socketId] = &SocketConnection{connection: conn, lock: &sync.Mutex{}, eventHandlers: eventHandlers, isTapper: isTapper}
+
+	websocketIdsLock.Unlock()
+
+	defer func() {
+		socketCleanup(socketId, connectedWebsockets[socketId])
+	}()
+
+	eventHandlers.WebSocketConnect(socketId, isTapper)
+
+	for {
+		_, msg, err := conn.ReadMessage()
+		if err != nil {
+			fmt.Printf("Conn err: %v\n", err)
+			break
+		}
+		eventHandlers.WebSocketMessage(socketId, msg)
+	}
+}
+
+func socketCleanup(socketId int, socketConnection *SocketConnection) {
+	err := socketConnection.connection.Close()
+	if err != nil {
+		fmt.Printf("Error closing socket connection for socket id %d: %v\n", socketId, err)
+	}
+
+	websocketIdsLock.Lock()
+	connectedWebsockets[socketId] = nil
+	websocketIdsLock.Unlock()
+
+	socketConnection.eventHandlers.WebSocketDisconnect(socketId, socketConnection.isTapper)
+}
+
+var db = debounce.NewDebouncer(time.Second * 5, func() {
+	fmt.Println("Successfully sent to socket")
+})
+
+func SendToSocket(socketId int, message []byte) error {
+	socketObj := connectedWebsockets[socketId]
+	if socketObj == nil {
+		return errors.New("Socket is disconnected")
+	}
+
+	var sent = false
+	time.AfterFunc(time.Second * 5, func() {
+		if !sent {
+			fmt.Println("Socket timed out")
+			socketCleanup(socketId, socketObj)
+		}
+	})
+
+	socketObj.lock.Lock() // gorilla socket panics from concurrent writes to a single socket
+	err := socketObj.connection.WriteMessage(1, message)
+	socketObj.lock.Unlock()
+
+	sent = true
+	return err
 }