iloom/washing-service/internal/handler/ws.go

package handler

import (
	"log"
	"net/http"

	"github.com/gin-gonic/gin"
	gorillaWS "github.com/gorilla/websocket"
	"github.com/muyuqingfeng/iloom/shared/pkg/websocket"
)

var upgrader = gorillaWS.Upgrader{
	ReadBufferSize:  1024,
	WriteBufferSize: 1024,
	CheckOrigin: func(r *http.Request) bool {
		return true
	},
}

func HandleWebSocket(hub *websocket.Hub) gin.HandlerFunc {
	return func(c *gin.Context) {
		companyID := c.GetHeader("X-Company-ID")
		if companyID == "" {
			companyID = c.Query("company_id")
		}
		userID := c.GetHeader("X-User-ID")
		if userID == "" {
			userID = c.Query("user_id")
		}

		if companyID == "" {
			c.JSON(http.StatusBadRequest, gin.H{"error": "missing company_id"})
			return
		}

		conn, err := upgrader.Upgrade(c.Writer, c.Request, nil)
		if err != nil {
			log.Printf("websocket upgrade error: %v", err)
			return
		}

		client := websocket.NewClient(hub, conn, companyID, userID)
		client.Register()

		go client.WritePump()
		go client.ReadPump()
	}
}