package controller import ( "encoding/json" "fmt" "github.com/gin-gonic/gin" "github.com/gorilla/websocket" "net/http" "strconv" "trading-go/model" "trading-go/response" ) var UP = websocket.Upgrader{ ReadBufferSize: 1024, WriteBufferSize: 1024, } type client struct { conn *websocket.Conn pip chan *model.Message } var Conns map[int64]*client func init() { Conns = make(map[int64]*client) } // 发送消息 func send(conn *websocket.Conn, uid int64) { pip := Conns[uid].pip defer func() { delete(Conns, uid) }() for { data := <-pip msg, err := json.Marshal(data) if err != nil { fmt.Println("link closed") break } err = conn.WriteMessage(data.MsgType, msg) if err != nil { fmt.Println(err.Error()) break } } } // 接收消息 func reception(conn *websocket.Conn, uid int64) { defer func() { delete(Conns, uid) }() for { var msg model.Message err := conn.ReadJSON(&msg) if err != nil { fmt.Println(err.Error()) break } if v, ok := Conns[msg.To]; ok { v.pip <- &msg } else { } } } func Chat(w http.ResponseWriter, rq *http.Request, uid int64) { // 升级为websocket conn, err := UP.Upgrade(w, rq, nil) if err != nil { response.Fail(w, err.Error(), http.StatusUpgradeRequired) return } pip := make(chan *model.Message, 1024) client := client{ conn: conn, pip: pip, } Conns[uid] = &client go send(conn, uid) go reception(conn, uid) response.Success(w, "success", nil) } func LinkToServer(c *gin.Context) { uid, err := strconv.Atoi(c.Query("uid")) if err != nil { response.Fail(c.Writer, "failed", 500) return } Chat(c.Writer, c.Request, int64(uid)) }