初始化项目文件

api_iris/service/ws/websocket.go

@@ -0,0 +1,71 @@
+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
+}