WebSocket编程示例

1、聊天室小示例

这里使用go语言做演示

需要安装第三方包：

• cmd中：go get -u -v github.com/gorilla/websocket

也可以使用：go get -u -v golang.org/x/net/websocket

在同一级目录下新建四个go文件

connection.go | data.go | hub.go | server.go

2、server.go文件代码

package main

import (
    "fmt"
    "github.com/gorilla/mux"
    "net/http"
    "text/template"
)

func main(){
    router := mux.NewRouter()

    // 开启一个 协程 运行 调度
    go h.run()

    // 注册 websocket 连接处理函数
    router.HandleFunc("/ws", myws)

    router.HandleFunc("/chat", chat)

    if err := http.ListenAndServe("127.0.0.1:8080", router); err != nil {
        fmt.Println("err:", err)
        return
    }
}

func chat(w http.ResponseWriter, r *http.Request){
    // 解析指定文件生成模板对象
    tmpl, err := template.ParseFiles("./other/websocket/local.html")
    if err != nil {
        fmt.Println("create template failed, err:", err)
        return
    }
    // 利用给定数据渲染模板，并将结果写入w
    tmpl.Execute(w, nil)
}

3、hub.go文件代码

package main

import "encoding/json"

// 初始化 调度 实例
var h = hub{
    c: make(map[*connection]bool),
    u: make(chan *connection),
    b: make(chan []byte),
    r: make(chan *connection),
}

// 定义 调度 结构体
type hub struct {
    c map[*connection]bool
    b chan []byte
    r chan *connection
    u chan *connection
}

// 定义 调度 运行函数
func (h *hub) run() {
    for {
        select {
        // 检查 是否 可以从 r 管道 获取 连接
        case c := <-h.r:
            // 将连接 存入 连接 map 并设置状态 为 true
            h.c[c] = true
            // 写入 连接 的 基本信息 如 来源IP、当前连接类型、全局用户列表
            c.data.Ip = c.ws.RemoteAddr().String()
            c.data.Type = "handshake"
            c.data.UserList = user_list
            data_b, _ := json.Marshal(c.data)
            // 将 连接后 的数据 格式化为 json 后 写入 连接的 发送管道
            c.sc <- data_b
        // 检查 是否 可以从 u 管道 获取 连接
        case c := <-h.u:
            // 检查 连接 map 中 是否存在 当前连接，存在则 删除连接，并 关闭通信管道
            if _, ok := h.c[c]; ok {
                delete(h.c, c)
                close(c.sc)
            }
        // 检查 是否 可以从 b 管道 获取 数据    
        case data := <-h.b:
            // 将 从管道 中获取 到的 数据 发送给 每一个 连接
            for c := range h.c {
                select {
                case c.sc <- data:
                default:
                    // 此处表示 如果 发送消息不成功时，删除 这个连接，并关闭这个连接的 发送管道
                    delete(h.c, c)
                    close(c.sc)
                }
            }
        }
    }
}

4、data.go文件代码

package main

type Data struct {
    Ip       string   `json:"ip"`
    User     string   `json:"user"`
    From     string   `json:"from"`
    Type     string   `json:"type"`
    Content  string   `json:"content"`
    UserList []string `json:"user_list"`
}

5、connection.go文件代码

package main

import (
    "encoding/json"
    "fmt"
    "github.com/gorilla/websocket"
    "net/http"
)

// 定义连接结构体
type connection struct {
    // websocket 连接指针 对象
    ws   *websocket.Conn
    // 连接的 消息管道
    sc   chan []byte
    // 数据指针 对象
    data *Data
}

// 定义 升级为 WebSocket 的 配置
// Upgrader 指定用于将 HTTP 连接升级为 WebSocket 连接的参数。并发调用升级程序的方法是安全的。
var wu = &websocket.Upgrader{
    ReadBufferSize: 512,
    WriteBufferSize: 512,
    CheckOrigin: func(r *http.Request) bool { return true },
}

// 定义 websocket 连接处理函数
func myws(w http.ResponseWriter, r *http.Request) {
    // 将 http 请求升级为 websocket
    ws, err := wu.Upgrade(w, r, nil)
    if err != nil {
        return
    }
    // 实例化一个连接实例
    c := &connection{sc: make(chan []byte, 256), ws: ws, data: &Data{}}
    // 将 连接 写入 调度的 r 管道
    h.r <- c

    // 开启一个 协程 处理 连接的 写操作
    go c.writer()

    // 从连接中读取 消息
    c.reader()

    // 使用 defer 处理 websocket 断开时的操作
    defer func() {
        c.data.Type = "logout"
        user_list = del(user_list, c.data.User)
        c.data.UserList = user_list
        c.data.Content = c.data.User
        data_b, _ := json.Marshal(c.data)
        h.b <- data_b
        h.u <- c
    }()
}

// 定义 连接 写入 消息 的方法
func (c *connection) writer() {
    for message := range c.sc {
        // 从 连接的 sc 管道 中 逐个 取出 消息，并 写入 websocket 的消息中
        c.ws.WriteMessage(websocket.TextMessage, message)
    }
    c.ws.Close()
}

// 定义用户列表
var user_list = []string{}

// 定义 连接 读取 消息 的方法
func (c *connection) reader() {
    for {
        // 从 websocket 连接中 读取 消息
        _, message, err := c.ws.ReadMessage()
        if err != nil {
            // 如果 读取消息 发生错误，则 将连接写入 调度的 u 管道
            h.u <- c
            break
        }
        // 将消息解析为 json 格式 ，并绑定 到 连接实例的 data 属性
        json.Unmarshal(message, &c.data)

        // 检查消息 类型
        switch c.data.Type {
        // 登录 消息
        case "login":
            c.data.User = c.data.Content
            c.data.From = c.data.User
            // 将 登录用户 加入 全局 用户列表（这里是演示，正式情况需要验证用户是否已经登录）
            user_list = append(user_list, c.data.User)
            // 将 目前 登录的用户 绑定到 数据 上
            c.data.UserList = user_list
            data_b, _ := json.Marshal(c.data)
            // 格式化 消息数据 为 json 格式后 写入 调度的 b 管道
            h.b <- data_b
        // 普通 聊天消息
        case "user":
            c.data.Type = "user"
            data_b, _ := json.Marshal(c.data)
            // 格式化 消息数据 为 json 格式后 写入 调度的 b 管道
            h.b <- data_b
        // 退出消息
        case "logout":
            c.data.Type = "logout"
            user_list = del(user_list, c.data.User)
            data_b, _ := json.Marshal(c.data)
            // 格式化 消息数据 为 json 格式后 写入 调度的 b 管道
            h.b <- data_b
            // 并将 当前连接 写入 调度的 u 管道
            h.u <- c
        // 其他未定义的 消息
        default:
            fmt.Print("========default================")
        }
    }
}

// 删除用户
func del(slice []string, user string) []string {
    count := len(slice)
    if count == 0 {
        return slice
    }
    if count == 1 && slice[0] == user {
        return []string{}
    }
    var n_slice = []string{}
    for i := range slice {
        if slice[i] == user && i == count {
            return slice[:count]
        } else if slice[i] == user {
            n_slice = append(slice[:i], slice[i+1:]...)
            break
        }
    }
    fmt.Println(n_slice)
    return n_slice
}

6、local.html文件代码

<!DOCTYPE html>
<html>
<head>
    <title></title>
    <meta http-equiv="content-type" content="text/html;charset=utf-8">
    <style>
        p {
            text-align: left;
            padding-left: 20px;
        }
    </style>
</head>
<body>
<div style="width: 800px;height: 600px;margin: 30px auto;text-align: center">
    <h1>演示聊天室</h1>
    <div style="width: 800px;border: 1px solid gray;height: 300px;">
        <div style="width: 200px;height: 300px;float: left;text-align: left;">
            <p><span>当前在线:</span><span id="user_num">0</span></p>
            <div id="user_list" style="overflow: auto;">
            </div>
        </div>
        <div id="msg_list" style="width: 598px;border:  1px solid gray; height: 300px;overflow: scroll;float: left;">
        </div>
    </div>
    <br>
    <textarea id="msg_box" rows="6" cols="50" onkeydown="confirm(event)"></textarea><br>
    <input type="button" value="发送" onclick="send()">
</div>
</body>
</html>
<script type="text/javascript">
    var uname = prompt('请输入用户名', 'user' + uuid(8, 16));
    var ws = new WebSocket("ws://127.0.0.1:8080/ws");
    ws.onopen = function () {
        var data = "系统消息：建立连接成功";
        listMsg(data);
    };
    ws.onmessage = function (e) {
        var msg = JSON.parse(e.data);
        var sender, user_name, name_list, change_type;
        switch (msg.type) {
            case 'system':
                sender = '系统消息: ';
                break;
            case 'user':
                sender = msg.from + ': ';
                break;
            case 'handshake':
                var user_info = {'type': 'login', 'content': uname};
                sendMsg(user_info);
                return;
            case 'login':
            case 'logout':
                user_name = msg.content;
                name_list = msg.user_list;
                change_type = msg.type;
                dealUser(user_name, change_type, name_list);
                return;
        }
        var data = sender + msg.content;
        listMsg(data);
    };
    ws.onerror = function () {
        var data = "系统消息 : 出错了,请退出重试.";
        listMsg(data);
    };
    function confirm(event) {
        var key_num = event.keyCode;
        if (13 == key_num) {
            send();
        } else {
            return false;
        }
    }
    function send() {
        var msg_box = document.getElementById("msg_box");
        var content = msg_box.value;
        var reg = new RegExp("\r\n", "g");
        content = content.replace(reg, "");
        var msg = {'content': content.trim(), 'type': 'user'};
        sendMsg(msg);
        msg_box.value = '';
    }
    function listMsg(data) {
        var msg_list = document.getElementById("msg_list");
        var msg = document.createElement("p");
        msg.innerHTML = data;
        msg_list.appendChild(msg);
        msg_list.scrollTop = msg_list.scrollHeight;
    }
    function dealUser(user_name, type, name_list) {
        var user_list = document.getElementById("user_list");
        var user_num = document.getElementById("user_num");
        while(user_list.hasChildNodes()) {
            user_list.removeChild(user_list.firstChild);
        }
        for (var index in name_list) {
            var user = document.createElement("p");
            user.innerHTML = name_list[index];
            user_list.appendChild(user);
        }
        user_num.innerHTML = name_list.length;
        user_list.scrollTop = user_list.scrollHeight;
        var change = type == 'login' ? '上线' : '下线';
        var data = '系统消息: ' + user_name + ' 已' + change;
        listMsg(data);
    }
    function sendMsg(msg) {
        var data = JSON.stringify(msg);
        ws.send(data);
    }
    function uuid(len, radix) {
        var chars = '0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz'.split('');
        var uuid = [], i;
        radix = radix || chars.length;
        if (len) {
            for (i = 0; i < len; i++) uuid[i] = chars[0 | Math.random() * radix];
        } else {
            var r;
            uuid[8] = uuid[13] = uuid[18] = uuid[23] = '-';
            uuid[14] = '4';
            for (i = 0; i < 36; i++) {
                if (!uuid[i]) {
                    r = 0 | Math.random() * 16;
                    uuid[i] = chars[(i == 19) ? (r & 0x3) | 0x8 : r];
                }
            }
        }
        return uuid.join('');
    }
</script>