connect/websocket.go

package main

import (
	"flag"
	"fmt"
	"log"
	"net/http"
	"strings"
	"time"

	"github.com/DeanThompson/ginpprof"
	"github.com/denisbrodbeck/machineid"
	"github.com/gin-gonic/gin"
	"github.com/google/uuid"
	"github.com/gorilla/websocket"
	"sikey.com/websocket/config"
	"sikey.com/websocket/repositories"
	"sikey.com/websocket/server"
	"x.sikey.com.cn/serverx/dbx"
	"x.sikey.com.cn/serverx/gid"
	"x.sikey.com.cn/serverx/rdbx"
)

var name = flag.String("n", "w303", "the name of the server")
var prot = flag.Int64("p", 8080, "the port of the server")
var nodeId = flag.Int64("i", 1, "the node id of the server")
var configFile = flag.String("f", "./etc/websocket.toml", "the config file")

func main() {
	time.Local = time.UTC

	flag.Parse()
	gid.SetNodeId(*nodeId)
	config.MustLoadConfig(*configFile)
	port := fmt.Sprintf(":%d", *prot)
	log.Println("Server start at ", port)

	app := newApp()
	app.Run(fmt.Sprintf(":%d", prot))
}

func newApp() *gin.Engine {
	app := gin.Default()
	ginpprof.Wrap(app)
	gin.SetMode(config.GetEnvironment())

	id := serverId()
	srv := &server.Server{
		ID: id,
		Upgrader: websocket.Upgrader{
			ReadBufferSize:  config.Websocket.ReadBufferSize,
			WriteBufferSize: config.Websocket.WriteBufferSize,
			CheckOrigin: func(r *http.Request) bool {
				return true
			},
		},
		Hub:          server.NewHub(id),
		Repositories: repositories.NewRepositories(dbx.GetConnect(), rdbx.GetConnect()),
	}

	app.GET("/websocket/endpoint", func(ctx *gin.Context) { srv.WebsocketHandler(ctx) })
	app.GET("/websocket/clients", func(ctx *gin.Context) {
		clients := srv.Hub.GetClients()
		ctx.JSON(http.StatusOK, gin.H{"clients": clients})
	})
	return app
}

func serverId() string {
	var id string
	id, err := machineid.ID()
	if err != nil {
		id = uuid.NewString()
	} else {
		id = strings.ToLower(id)
	}
	return id
}