nginx-ui/server/modules/agent_server/server.go

package agent_server

import (
	"github.com/gorilla/websocket"
	"log"
	"net/http"
	"nginx-ui/server/models"
)

var upgrader = websocket.Upgrader{
	ReadBufferSize:  4096,
	WriteBufferSize: 4096,
}

type Server struct {
}

func NewServer() *Server {
	go AgentHub.run()
	return &Server{}
}

func SetMessageHandler(key string, handler func(c *WsClient, message *models.AgentData) (interface{}, error)) {
	AgentHub.SetMessageHandler(key, handler)
}

func (s *Server) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Println(err)
		return
	}
	token := r.Header.Get("Token")
	client := NewWsClient(token, conn)
	AgentHub.register <- client
	go client.readPump()
}