使用 Go + WebSockets 创建实时聊天服务

安装第三方库

go get github.com/gorilla/websocket

服务端

src/main.go

package main

import (
    "github.com/gorilla/websocket"
    "log"
    "net/http"
)

// 使用字典结构更容易追加和删除内容
var clients = make(map[*websocket.Conn]bool) // 保存连接的客户端
var broadcast = make(chan Message)           // 消息广播通道

// 配置 upgrader
// 其中包含获取正常HTTP连接并将其升级到WebSocket的方法
var upgrader = websocket.Upgrader{
    CheckOrigin: func(r *http.Request) bool {
        return true
    },
}

// 定义 message 结构
type Message struct {
    Email    string `json:"email"`
    Username string `json:"username"`
    Message  string `json:"message"`
}

func handleConnections(w http.ResponseWriter, r *http.Request) {
    // 将初始GET请求升级到websocket
    ws, err := upgrader.Upgrade(w, r, nil)
    if err != nil {
        log.Fatal(err)
    }
    // 确保在函数返回时关闭连接
    defer ws.Close()
    // 将新的客户端连接添加到字典中
    clients[ws] = true

    for {
        var msg Message
        // 反序列化
        err := ws.ReadJSON(&msg)
        if err != nil {
            log.Printf("read error: %v", err)
            delete(clients, ws)
            break
        }
        // 将新接收的消息发送到广播通道
        broadcast <- msg
    }
}

func handleMessages() {
    for {
        // 从广播通道中获取消息
        msg := <-broadcast
        // 将其发送给当前连接的每个客户端
        for ws := range clients {
            err := ws.WriteJSON(msg)
            if err != nil {
                log.Printf("write error: %v", err)
                ws.Close()
                delete(clients, ws)
            }
        }
    }
}

func main() {
    // 创建简单的静态文件服务器
    http.Handle("/", http.FileServer(http.Dir("../public")))
    // 配置 websocket route
    http.HandleFunc("/ws", handleConnections)

    go handleMessages()

    log.Println("http server started on :8000")
    err := http.ListenAndServe(":8000", nil)
    if err != nil {
        log.Fatal("ListenAndServe error: ", err)
    }
}

客户端

public/index.html

    
    

    
    
    
    




    

        

            Simple Chat
        
    


    

        

            

                

                
            
        
    
    

        

            
        
        

            
                chat
                Send
            
        
    
    

        

            
        
        

            
        
        

            
                done
                Join
            
        
    

public/app.js

new Vue({
    el: '#app',
    data: {
        ws: null, // Our websocket
        newMsg: '', // Holds new messages to be sent to the server
        chatContent: '', // A running list of chat messages displayed on the screen
        email: null, // Email address used for grabbing an avatar
        username: null, // Our username
        joined: false // True if email and username have been filled in
    },
    created: function() {
        var self = this;
        this.ws = new WebSocket('ws://' + window.location.host + '/ws');
        this.ws.addEventListener('message', function(e) {
            var msg = JSON.parse(e.data);
            self.chatContent += '
'
                + '' // Avatar
                + msg.username
                + '
'
                + emojione.toImage(msg.message) + '
'; // Parse emojis

            var element = document.getElementById('chat-messages');
            element.scrollTop = element.scrollHeight; // Auto scroll to the bottom
        });
    },
    methods: {
        send: function () {
            if (this.newMsg != '') {
                this.ws.send(
                    JSON.stringify({
                            email: this.email,
                            username: this.username,
                            // 使用一些jQuery技巧来避免任何传入消息中的HTML和JavaScript
                            // 这可以防止任何类型的注入攻击
                            message: $('
').html(this.newMsg).text()
                        }
                    ));
                // 重置 newMsg
                this.newMsg = '';
            }
        },
        join: function () {
            if (!this.email) {
                Materialize.toast('You must enter an email', 2000);
                return
            }
            if (!this.username) {
                Materialize.toast('You must choose a username', 2000);
                return
            }
            this.email = $('
').html(this.email).text();
            this.username = $('
').html(this.username).text();
            this.joined = true;
        },
        gravatarURL: function(email) {
            return 'http://www.gravatar.com/avatar/' + CryptoJS.MD5(email);
        }
    }
});

public/style.css

body {
    display: flex;
    min-height: 100vh;
    flex-direction: column;
}

main {
    flex: 1 0 auto;
}

#chat-messages {
    min-height: 10vh;
    height: 60vh;
    width: 100%;
    overflow-y: scroll;
}

目录结构

运行

进入到 src 目录下执行：

$ go run main.go

打开浏览器，输入：
http://localhost:8000

参考地址：https://scotch.io/bar-talk/build-a-realtime-chat-server-with-go-and-websockets