72 lines
1.5 KiB
Go
72 lines
1.5 KiB
Go
package ws
|
||
|
||
import (
|
||
websocket1 "github.com/gorilla/websocket"
|
||
"github.com/kataras/iris/v12/websocket"
|
||
"github.com/kataras/neffos"
|
||
"github.com/sirupsen/logrus"
|
||
"net/http"
|
||
"strconv"
|
||
)
|
||
|
||
func checkOrigin(_ *http.Request) bool {
|
||
return true
|
||
}
|
||
|
||
func PopItem(l []string, item string) []string {
|
||
var tmp []string
|
||
for _, v := range l {
|
||
if v == item {
|
||
continue
|
||
} else {
|
||
tmp = append(tmp, v)
|
||
}
|
||
}
|
||
return tmp
|
||
}
|
||
|
||
func SetupWebsocket() *neffos.Server {
|
||
Upgrader := websocket1.Upgrader{
|
||
ReadBufferSize: 1024,
|
||
WriteBufferSize: 1024,
|
||
CheckOrigin: checkOrigin,
|
||
}
|
||
var userList []string
|
||
ws := websocket.New(websocket.GorillaUpgrader(Upgrader), websocket.Events{
|
||
websocket.OnNativeMessage: func(conn *websocket.NSConn, message websocket.Message) error {
|
||
mg := websocket.Message{
|
||
Body: []byte(strconv.Itoa(len(userList))),
|
||
IsNative: true,
|
||
}
|
||
conn.Conn.Write(mg)
|
||
return nil
|
||
},
|
||
})
|
||
ws.OnConnect = func(c *websocket.Conn) error {
|
||
userList = append(userList, c.ID())
|
||
mg := websocket.Message{
|
||
Body: []byte(strconv.Itoa(len(userList))),
|
||
IsNative: true,
|
||
}
|
||
c.Write(mg)
|
||
for _, cli := range ws.GetConnections() {
|
||
cli.Write(mg)
|
||
}
|
||
return nil
|
||
}
|
||
ws.OnDisconnect = func(c *websocket.Conn) {
|
||
userList = PopItem(userList, c.ID())
|
||
mg := websocket.Message{
|
||
Body: []byte(strconv.Itoa(len(userList))),
|
||
IsNative: true,
|
||
}
|
||
for _, cli := range ws.GetConnections() {
|
||
cli.Write(mg)
|
||
}
|
||
}
|
||
ws.OnUpgradeError = func(err error) {
|
||
logrus.Errorln("ws初始化失败：", err)
|
||
}
|
||
return ws
|
||
}
|