Go WebSocket 200 行代碼開發一個簡易聊天室
日期 : 2023-01-31 18:03:26作者 : 熊人
lib
效果圖
服務端 main.go
package main
import (
_ "embed"
"encoding/json"
"github.com/lxzan/gws"
"log"
"net/http"
"time"
)
const PingInterval = 15 * time.Second // 客戶端心跳間隔
//go:embed index.html
var html []byte
func main() {
var handler = NewWebSocket()
var upgrader = gws.NewUpgrader(func(c *gws.Upgrader) {
c.CompressEnabled = true
c.EventHandler = handler
// 在 querystring 里面傳入用戶名
// 把 Sec-WebSocket-Key 作為連接的 key
// 刷新頁面的時候, 會觸發上一個連接的 OnClose/OnError 事件, 這時候需要對比 key 并刪除 map 里存儲的連接
c.CheckOrigin = func(r *gws.Request) bool {
var name = r.URL.Query().Get("name")
if name == "" {
return false
}
r.SessionStorage.Store("name", name)
r.SessionStorage.Store("key", r.Header.Get("Sec-WebSocket-Key"))
return true
}
})
http.HandleFunc("/connect", func(writer http.ResponseWriter, request *http.Request) {
socket, err := upgrader.Accept(writer, request)
if err != nil {
log.Printf("Accept: " + err.Error())
return
}
socket.Listen()
})
http.HandleFunc("/index.html", func(writer http.ResponseWriter, request *http.Request) {
_, _ = writer.Write(html)
})
if err := http.ListenAndServe(":3000", nil); err != nil {
log.Fatalf("%+v", err)
}
}
func NewWebSocket() *WebSocket {
return &WebSocket{sessions: gws.NewConcurrentMap(16)}
}
type WebSocket struct {
sessions *gws.ConcurrentMap // 使用內置的 ConcurrentMap 存儲連接, 可以減少鎖沖突
}
func (c *WebSocket) getName(socket *gws.Conn) string {
name, _ := socket.SessionStorage.Load("name")
return name.(string)
}
func (c *WebSocket) getKey(socket *gws.Conn) string {
name, _ := socket.SessionStorage.Load("key")
return name.(string)
}
// 根據用戶名獲取 WebSocket 連接
func (c *WebSocket) GetSocket(name string) (*gws.Conn, bool) {
if v0, ok0 := c.sessions.Load(name); ok0 {
if v1, ok1 := v0.(*gws.Conn); ok1 {
return v1, true
}
}
return nil, false
}
// RemoveSocket 移除 WebSocket 連接
func (c *WebSocket) RemoveSocket(socket *gws.Conn) {
name := c.getName(socket)
key := c.getKey(socket)
if mSocket, ok := c.GetSocket(name); ok {
if mKey := c.getKey(mSocket); mKey == key {
c.sessions.Delete(name)
}
}
}
func (c *WebSocket) OnOpen(socket *gws.Conn) {
name := c.getName(socket)
if v, ok := c.sessions.Load(name); ok {
var conn = v.(*gws.Conn)
conn.Close(1000, []byte("connection replaced"))
}
socket.SetDeadline(time.Now().Add(3 * PingInterval))
c.sessions.Store(name, socket)
log.Printf("%s connected\n", name)
}
func (c *WebSocket) OnError(socket *gws.Conn, err error) {
name := c.getName(socket)
c.RemoveSocket(socket)
log.Printf("onerror, name=%s, msg=%s\n", name, err.Error())
}
func (c *WebSocket) OnClose(socket *gws.Conn, code uint16, reason []byte) {
name := c.getName(socket)
c.RemoveSocket(socket)
log.Printf("onclose, name=%s, code=%d, msg=%s\n", name, code, string(reason))
}
func (c *WebSocket) OnPing(socket *gws.Conn, payload []byte) {}
func (c *WebSocket) OnPong(socket *gws.Conn, payload []byte) {}
type Input struct {
To string `json:"to"`
Text string `json:"text"`
}
func (c *WebSocket) OnMessage(socket *gws.Conn, message *gws.Message) {
defer message.Close()
// chrome websocket 不支持 ping 方法, 所以在 text frame 里面模擬 ping
if b := message.Bytes(); len(b) == 4 && string(b) == "ping" {
socket.WriteMessage(gws.OpcodeText, []byte("pong"))
socket.SetDeadline(time.Now().Add(3 * PingInterval))
return
}
var input = &Input{}
_ = json.Unmarshal(message.Bytes(), input)
if v, ok := c.sessions.Load(input.To); ok {
v.(*gws.Conn).WriteMessage(gws.OpcodeText, message.Bytes())
}
}
客戶端 index.html
<!DOCTYPE html>
<html lang="en">
<head>
<meta charset="UTF-8">
<title>ChatRoom</title>
<style>
#app {
width: 400px;
margin: 50px auto 0;
}
.form {
margin: 10px auto;
}
#app input {
width: 300px;
height: 20px;
float: right;
}
#app span {
height: 26px;
line-height: 26px;
}
textarea {
width: 400px;
}
</style>
</head>
<body>
<div id="app">
<div class="form"><span>From</span> <input type="text" id="from"></div>
<div class="form"><span>To</span> <input type="text" id="to"></div>
<div><textarea id="text" cols="30" rows="10"></textarea></div>
<button onclick="connect()">Connect</button>
<button onclick="send()">Send</button>
</div>
<script>
function connect() {
let from = document.getElementById("from").value;
window.ws = new WebSocket(`ws://127.0.0.1:3000/connect?name=${from}`);
window.ws.onclose = function (event) {
console.log(event);
}
if (window.interval !== undefined) {
clearInterval(window.interval)
}
window.interval = setInterval(function () {
window.ws.send("ping");
}, 15 * 1000)
}
function send() {
let to = document.getElementById("to").value;
let text = document.getElementById("text").value;
ws.send(JSON.stringify({to, text}));
}
</script>
</body>
</html>
