modify websocketserver to broker with callback functions

websocket/clientHandler.go

@@ -0,0 +1,73 @@
+package websocket
+
+import (
+	"encoding/json"
+	"fmt"
+	"sync"
+
+	"github.com/gin-gonic/gin"
+	"github.com/tecamino/tecamino-dbm/websocket/models"
+	wsModels "github.com/tecamino/tecamino-dbm/websocket/models"
+	json_dataModels "github.com/tecamino/tecamino-json_data/models"
+)
+
+// serves as connection handler of websocket
+type ClientHandler struct {
+	sync.RWMutex
+	Clients models.Clients
+}
+
+// initaiates new conections with client map
+func NewConnectionHandler() *ClientHandler {
+	return &ClientHandler{
+		Clients: make(models.Clients),
+	}
+}
+
+// Connect a recieving websocket connection
+func (cH *ClientHandler) ConnectNewClient(id string, c *gin.Context) (client *models.Client, err error) {
+	if _, exists := cH.Clients[id]; exists {
+		return cH.Clients[id], nil
+	}
+
+	client, err = models.ConnectNewClient(id, c)
+	client.OnClose = func(code int, reason string) {
+		delete(cH.Clients, id)
+	}
+
+	if err != nil {
+		return nil, err
+	}
+	cH.Lock()
+	cH.Clients[id] = client
+	cH.Unlock()
+
+	return client, nil
+}
+
+// get client client
+func (c *ClientHandler) GetClient(id string) *wsModels.Client {
+	if client, ok := c.Clients[id]; ok {
+		return client
+	}
+	return nil
+}
+
+// sends json response to client
+func (c *ClientHandler) SendResponse(id string, r *json_dataModels.Response) error {
+	client, ok := c.Clients[id]
+	if !ok {
+		return fmt.Errorf("client not found for id %s", id)
+
+	}
+	b, err := json.Marshal(r)
+	if err != nil {
+		return err
+	}
+
+	if err := client.SendResponse(b, 5); err != nil {
+		return err
+	}
+
+	return nil
+}

websocket/models/client.go

@@ -0,0 +1,164 @@
+package models
+
+import (
+	"encoding/json"
+	"fmt"
+	"log"
+	"net/http"
+	"time"
+
+	"github.com/gin-gonic/gin"
+	"github.com/gorilla/websocket"
+)
+
+var Origins []string = []string{"*"}
+
+type Client struct {
+	Id        string
+	Connected bool            `json:"connected"`
+	Conn      *websocket.Conn `json:"-"`
+	OnOpen    func()
+	OnMessage func(data []byte)
+	OnClose   func(code int, reason string)
+	OnError   func(err error)
+	OnPing    func()
+	OnPong    func()
+	timeout   time.Duration
+}
+
+var upgrader = websocket.Upgrader{
+	CheckOrigin: func(r *http.Request) bool {
+		if len(Origins) == 0 {
+			return false
+		}
+		if Origins[0] == "*" {
+			return true
+		}
+		origin := r.Header.Get("Origin")
+		for _, o := range Origins {
+			if o == origin {
+				return true
+			}
+		}
+		return false
+	},
+	EnableCompression: false,
+}
+
+func ConnectNewClient(id string, c *gin.Context) (*Client, error) {
+	conn, err := upgrader.Upgrade(c.Writer, c.Request, nil)
+	if err != nil {
+		return nil, fmt.Errorf("websocket upgrade error: %w", err)
+	}
+
+	client := &Client{
+		Id:        id,
+		Connected: true,
+		Conn:      conn,
+		timeout:   5,
+	}
+
+	conn.SetPingHandler(func(appData string) error {
+		if client.OnPing != nil {
+			client.OnPing()
+		}
+		conn.SetWriteDeadline(time.Now().Add(client.timeout))
+		conn.SetReadDeadline(time.Now().Add(client.timeout))
+		return conn.WriteMessage(websocket.PongMessage, []byte(appData))
+	})
+
+	conn.SetPongHandler(func(appData string) error {
+		if client.OnPong != nil {
+			client.OnPong()
+		}
+		conn.SetReadDeadline(time.Now().Add(client.timeout))
+		return nil
+	})
+
+	// Start reading messages from client
+	go client.Listen(7)
+	return client, nil
+}
+
+func (c *Client) Listen(timeout uint) {
+	if timeout > 0 {
+		c.timeout = time.Duration(timeout) * time.Second
+	}
+
+	if c.OnOpen != nil {
+		c.OnOpen()
+	}
+
+	c.Conn.SetReadDeadline(time.Now().Add(c.timeout))
+	for c.Connected {
+		msgType, msg, err := c.Conn.ReadMessage()
+		if err != nil {
+			c.handleError(fmt.Errorf("read error (id:%s): %w", c.Id, err))
+			return
+		}
+		switch msgType {
+		case websocket.CloseMessage:
+			c.handleClose(1000, "Client closed")
+			return
+		case websocket.TextMessage:
+			if isPing := c.handleJsonPing(msg); isPing {
+				continue
+			}
+			if c.OnMessage != nil {
+				c.OnMessage(msg)
+			} else {
+				log.Printf("Received message but no handler set (id:%s): %s", c.Id, string(msg))
+			}
+		default:
+			log.Printf("Unhandled message type %d (id:%s)", msgType, c.Id)
+		}
+	}
+}
+
+func (c *Client) handleJsonPing(msg []byte) (isPing bool) {
+	var wsMsg WSMessage
+	err := json.Unmarshal(msg, &wsMsg)
+	if err == nil && wsMsg.IsPing() {
+		c.Conn.SetReadDeadline(time.Now().Add(c.timeout))
+		// Respond with pong JSON
+		c.Conn.SetWriteDeadline(time.Now().Add(c.timeout))
+		err = c.Conn.WriteMessage(websocket.TextMessage, GetPongByteSlice())
+		if err != nil {
+			c.handleError(fmt.Errorf("write pong error: %w", err))
+			return
+		}
+		if c.OnPing != nil {
+			c.OnPing()
+		}
+		isPing = true
+	}
+	return
+}
+
+func (c *Client) SendResponse(data []byte, timeout uint) error {
+	c.Conn.SetWriteDeadline(time.Now().Add(time.Duration(timeout) * time.Second))
+	return c.Conn.WriteMessage(websocket.TextMessage, data)
+}
+
+func (c *Client) Close(code int, reason string) {
+	c.handleClose(code, reason)
+}
+
+func (c *Client) handleClose(code int, text string) {
+	if !c.Connected {
+		return
+	}
+	c.Connected = false
+	if c.OnClose != nil {
+		c.OnClose(code, text)
+	}
+	c.Conn.WriteMessage(websocket.CloseMessage, websocket.FormatCloseMessage(code, text))
+	c.Conn.Close()
+}
+
+func (c *Client) handleError(err error) {
+	if c.OnError != nil {
+		c.OnError(err)
+	}
+	c.Close(websocket.CloseInternalServerErr, err.Error())
+}

websocket/models/clients.go

@@ -0,0 +1,3 @@
+package models
+
+type Clients map[string]*Client

websocket/models/wsMessage.go

@@ -0,0 +1,21 @@
+package models
+
+import "encoding/json"
+
+type WSMessage struct {
+	Type string `json:"type"`
+}
+
+func GetPongByteSlice() []byte {
+	b, err := json.Marshal(WSMessage{
+		Type: "pong",
+	})
+	if err != nil {
+		return []byte{}
+	}
+	return b
+}
+
+func (w WSMessage) IsPing() bool {
+	return w.Type == "ping"
+}

websocket/server.go

@@ -0,0 +1,49 @@
+package websocket
+
+import (
+	"fmt"
+	"sync"
+	"time"
+
+	"github.com/gin-contrib/cors"
+	"github.com/gin-gonic/gin"
+	"github.com/tecamino/tecamino-dbm/cert"
+	"github.com/tecamino/tecamino-logger/logging"
+)
+
+// server model for database manager websocket
+type Server struct {
+	Routes *gin.Engine
+	sync.RWMutex
+	Logger *logging.Logger
+}
+
+// initalizes new dbm server
+func NewServer(allowOrigins []string) *Server {
+	r := gin.Default()
+	r.Use(cors.New(cors.Config{
+		AllowOrigins:     allowOrigins,
+		AllowMethods:     []string{"GET", "POST", "PUT", "DELETE", "OPTIONS"},
+		AllowHeaders:     []string{"Origin", "Content-Type", "Accept"},
+		ExposeHeaders:    []string{"Content-Length"},
+		AllowCredentials: true,
+		MaxAge:           12 * time.Hour,
+	}))
+	return &Server{
+		Routes: r,
+	}
+}
+
+// serve dbm as http
+func (s *Server) ServeHttp(ip string, port uint) error {
+	return s.Routes.Run(fmt.Sprintf("%s:%d", ip, port))
+}
+
+// serve dbm as http
+func (s *Server) ServeHttps(ip string, port uint, cert cert.Cert) error {
+	// generate self signed tls certificate
+	if err := cert.GenerateSelfSignedCert(); err != nil {
+		return err
+	}
+	return s.Routes.RunTLS(fmt.Sprintf("%s:%d", ip, port), cert.CertFile, cert.KeyFile)
+}