golang：websocket消息的定向人投送

golang：websocket消息的定向人投送

后端

package main
import (
    "encoding/json"
    "fmt"
    "net/http"
"github.com/gorilla/websocket"
//"github.com/gofrs/uuid"
)
//客户端管理
type ClientManager struct {
    //客户端 map 储存并管理所有的长连接client，在线的为true，不在的为false
    clients       map[*Client]bool
    //创建一个切片，来保存连接的用户
    receivePeople map[string]string
    //web端发送来的的message我们用broadcast来接收，并最后分发给所有的client
    broadcast chan []byte
    //谁发来的信息
    cid string
    //谁jie信息
    uid string
//新创建的长连接client
register chan *Client
//新注销的长连接client
unregister chan *Client
}
//客户端 Client
type Client struct {
    //用户id
    id string
    //连接的socket
    socket *websocket.Conn
    //发送的消息
    send chan []byte
}
//会把Message格式化成json
type Message struct {
    //消息struct
    Sender    string json:"sender,omitempty"    //发送者
    Recipient string json:"recipient,omitempty" //接收者
    Content   string json:"content,omitempty"   //内容
}
//创建客户端管理者
var manager = ClientManager{
    broadcast:     make(chan []byte),
    register:      make(chan Client),
    unregister:    make(chan *Client),
    clients:       make(map[Client]bool),
    cid:           "",
    receivePeople: make(map[string]string),
}
//socket 设置
var (
    upgrader = websocket.Upgrader{
        //
        ReadBufferSize: 1023,
        //
        WriteBufferSize: 1023,
        //允许跨域
        CheckOrigin: func(r *http.Request) bool {
            return true
        },
    }
)
func (manager *ClientManager) start() {
    fmt.Println("start")
    for {
        select {
        //如果有新的连接接入,就通过channel把连接传递给conn
        case conn := <-manager.register:
            //把客户端的连接设置为true
            manager.clients[conn] = true
            //把返回连接成功的消息json格式化
            jsonMessage, _ := json.Marshal(&Message{Content: "id:" + manager.cid + " has connected."})
            //调用客户端的send方法，发送消息
            manager.send(jsonMessage, conn)
            //如果连接断开了
        case conn := <-manager.unregister:
            //判断连接的状态，如果是true,就关闭send，删除连接client的值
            if _, ok := manager.clients[conn]; ok {
                close(conn.send)
                delete(manager.clients, conn)
                jsonMessage, _ := json.Marshal(&Message{Content: "id:" + manager.cid + " has disconnected."})
                manager.send(jsonMessage, conn)
            }
            //广播
        case message := <-manager.broadcast:
            //遍历已经连接的客户端，把消息发送给他们
            //fmt.Println(message)
            for conn := range manager.clients {
                //判断发送给谁（不发送自己）
                //接收者id： conn.id
                //发送者id： manager.cid
                if conn.id == manager.cid {
                    continue
                }
                //判断是否是接受人
                key, ok := manager.receivePeople[manager.cid]
                if !ok {
                    continue
                }
                if key != conn.id {
                    continue
                }
                select {
                case conn.send <- message:
                default:
                    close(conn.send)
                    delete(manager.clients, conn)
                }
            }
        }
    }
}
//定义客户端管理的send方法
func (manager *ClientManager) send(message []byte, ignore *Client) {
    for conn := range manager.clients {
        //不给屏蔽的连接发送消息
        if conn != ignore {
            conn.send <- message
        }
    }
}
//定义客户端结构体的read方法
func (c *Client) read() {
    defer func() {
        //结构体cid赋值
        manager.cid = c.id
        //触发关闭
        manager.unregister <- c
        c.socket.Close()
    }()
for {
    //读取消息
    _, message, err := c.socket.ReadMessage()
    //如果有错误信息，就注销这个连接然后关闭
    if err != nil {
        manager.unregister <- c
        c.socket.Close()
        break
    }
    //如果没有错误信息就把信息放入broadcast
    jsonMessage, _ := json.Marshal(&Message{Sender: c.id, Content: string(message)})
    //结构体cid赋值
    manager.cid = c.id
    //触发消息发送
    manager.broadcast <- jsonMessage
}
}
func (c *Client) write() {
    defer func() {
        c.socket.Close()
    }()
for {
    select {
    //从send里读消息
    case message, ok := <-c.send:
        //如果没有消息
        if !ok {
            c.socket.WriteMessage(websocket.CloseMessage, []byte{})
            return
        }
        //有消息就写入，发送给web端
        c.socket.WriteMessage(websocket.TextMessage, message)
    }
}
}
func main() {
    fmt.Println("Starting application...")
    //开一个goroutine执行开始程序
    go manager.start()
    //注册默认路由为 /ws ，并使用wsHandler这个方法
    http.HandleFunc("/ws", wsHandler)
    //监听本地的8182端口
    http.ListenAndServe(":8182", nil)
}
func wsHandler(res http.ResponseWriter, req *http.Request) {
    //将http协议升级成websocket协议
    conn, err := (&websocket.Upgrader{CheckOrigin: func(r *http.Request) bool { return true }}).Upgrade(res, req, nil)
    if err != nil {
        http.NotFound(res, req)
        return
    }
    //可以用传来的参数识别身份
    req.ParseForm()
    roomid := req.Form["roomid"][0]
    uid := req.Form["uid"][0]
    fmt.Println(roomid, "\n", uid)
    manager.receivePeople[roomid] = uid
    //这里是随机生成id
    //每一次连接都会新开一个client
    client := &Client{id: roomid, socket: conn, send: make(chan []byte)}
    //注册一个新的链接
    manager.cid = client.id
manager.register <- client

//启动协程收web端传过来的消息
go client.read()
//启动协程把消息返回给web端
go client.write()
}

前端页面

<html>
<head>
    <title>Golang Chat</title>
    <meta http-equiv="Content-Type" content="text/html; charset=utf-8"/>
    <script src="https://cdn.bootcdn.net/ajax/libs/jquery/3.5.1/jquery.js"></script>
    <script type="text/javascript">
        $(function () {

            var conn;
            var msg = $("#msg");
            var log = $("#log");

            function appendLog(msg) {
                var d = log[0]
                var doScroll = d.scrollTop == d.scrollHeight - d.clientHeight;
                msg.appendTo(log)
                if (doScroll) {
                    d.scrollTop = d.scrollHeight - d.clientHeight;
                }
            }

            $("#form").submit(function () {
                if (!conn) {
                    return false;
                }
                if (!msg.val()) {
                    return false;
                }
                conn.send(msg.val());
                msg.val("");
                return false
            });

            function getQueryVariable(variable) {
                var query = window.location.search.substring(1);
                var vars = query.split("&");
                for (var i = 0; i < vars.length; i++) {
                    var pair = vars[i].split("=");
                    if (pair[0] == variable) {
                        return pair[1];
                    }
                }
                return 0;
            }

            if (window["WebSocket"]) {
                //conn = new WebSocket("ws://localhost:8011/ws");
                conn = new WebSocket("ws://localhost:8182/ws?roomid=" + getQueryVariable('roomid') + "&uid=" + getQueryVariable('uid'));
                conn.onopen = function () {
                    // layer.load(1,{shade:0.1});
                    console.log('连接成功');
                };
                conn.onerror = function () {
                    console.log('服务器连接失败!');
                    //layer.msg('服务器连接失败!',{shade:0.1,icon:2});
                };
                conn.onclose = function (evt) {
                    appendLog($("<div><b>Connection Closed.</b></div>"))
                }
                conn.onmessage = function (evt) {
                    appendLog($("<div/>").text(evt.data))
                }
            } else {
                appendLog($("<div><b>WebSockets Not Support.</b></div>"))
            }


        });
    </script>
    <style type="text/css">
        html {
            overflow: hidden;
        }

        body {
            overflow: hidden;
            padding: 0;
            margin: 0;
            width: 100%;
            height: 100%;
            background: gray;
        }

        #log {
            background: white;
            margin: 0;
            padding: 0.5em 0.5em 0.5em 0.5em;
            position: absolute;
            top: 0.5em;
            left: 0.5em;
            right: 0.5em;
            bottom: 3em;
            overflow: auto;
        }

        #form {
            padding: 0 0.5em 0 0.5em;
            margin: 0;
            position: absolute;
            bottom: 1em;
            left: 0px;
            width: 100%;
            overflow: hidden;
        }

    </style>
</head>
<body>
<div id="log"></div>
<form id="form">
    <input type="submit" value="发送"/>
    <input type="text" id="msg" size="128"/>
</form>
</body>
</html>

原创地址：在此基础上加了一些用户储存。