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 {
	hub *Hub
}

var AgentHub = newHub()

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

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 := &WsClient{
		Token:     token,
		hub:       s.hub,
		conn:      conn,
		send:      make(chan []byte, 1024),
		callbacks: make(map[string]chan *models.AgentData),
	}
	client.hub.register <- client
	go client.readPump()
}