websocket学习笔记 go语言使用 gorilla/websocket实现实时聊天项目
完整代码：https://github.com/diguacheng/webchat
扩展包 gorilla/websocket的几个重要函数
- 协议升级 ，1）先初始化upgrader，再调用其upgrade方法将http协议升级为websicket协议。2）也可以直接使用upgrade函数 进行协议升级
type Upgrader struct {
HandshakeTimeout time.Duration
ReadBufferSize, WriteBufferSize int
WriteBufferPool BufferPool
Subprotocols []string
Error func(w http.ResponseWriter, r *http.Request, status int, reason error)
CheckOrigin func(r *http.Request) bool
EnableCompression bool
}
func Upgrade(w http.ResponseWriter, r *http.Request, responseHeader http.Header, readBufSize, writeBufSize int) (*Conn, error) {
u := Upgrader{ReadBufferSize: readBufSize, WriteBufferSize: writeBufSize}
u.Error = func(w http.ResponseWriter, r *http.Request, status int, reason error) {
// don't return errors to maintain backwards compatibility
}
u.CheckOrigin = func(r *http.Request) bool {
// allow all connections by default
return true
}
return u.Upgrade(w, r, responseHeader)
}
func Upgrade(w http.ResponseWriter, r *http.Request, responseHeader http.Header, readBufSize, writeBufSize int) (*Conn, error) {
u := Upgrader{ReadBufferSize: readBufSize, WriteBufferSize: writeBufSize}
u.Error = func(w http.ResponseWriter, r *http.Request, status int, reason error) {
// don't return errors to maintain backwards compatibility
}
u.CheckOrigin = func(r *http.Request) bool {
// allow all connections by default
return true
}
return u.Upgrade(w, r, responseHeader)
}
-
服务端发送数据
使用
func (c *Conn) WriteMessage(messageType int, data []byte) error方法。或者使用
func (c *Conn) NextWriter(messageType int) (io.WriteCloser, error)方法获得writer,再调用func (w *messageWriter) Write(p []byte) (int, error)将数据写入。 -
服务端接受数据
使用
func (c *Conn) ReadMessage() (messageType int, p []byte, err error)方法。
其他函数方法可参考文档
项目实战 web实时聊天室
项目的架构如图所示：
// Client is a middleman between the websocket connection and the hub.
type Client struct {
// The websocket connection.
conn *websocket.Conn
// Buffered channel of outbound messages.
send chan []byte
Data *Data
}
// Data 定义传输的数据，
type Data struct {
IP string `json:"ip"` // 本地ip
Type string `json:"type"` // 类型
User string `json:"user"` // 接受信息的用户名
Form string `json:"from"` // 发送信息的用户名
FormIP string `json:"FromIp"`// 发送信息的IP
Content string `json:"content"` // 信息的内容。
}
// Hub maintains the set of active clients and broadcasts messages to the
// clients.
type Hub struct {
// Registered clients. // 一个client代表一个user
clients map[*Client]bool
// Inbound messages from the clients.
broadcast chan []byte
// Register requests from the clients.
register chan *Client
// Unregister requests from clients.
unregister chan *Client
}
hub.go
package main
// Hub maintains the set of active clients and broadcasts messages to the
// clients.
type Hub struct {
// Registered clients.
clients map[*Client]bool
// Inbound messages from the clients.
broadcast chan []byte
// Register requests from the clients.
register chan *Client
// Unregister requests from clients.
unregister chan *Client
}
func newHub() *Hub {
return &Hub{
broadcast: make(chan []byte),
register: make(chan *Client),
unregister: make(chan *Client),
clients: make(map[*Client]bool),
}
}
var h = newHub()
func (h *Hub) run() {
for {
select {
case client := <-h.register:
h.clients[client] = true // client注册，
case client := <-h.unregister:
// client注销，将该client从map中删除，并关闭send通道
if _, ok := h.clients[client]; ok {
delete(h.clients, client)
close(client.send)
}
case message := <-h.broadcast:
// 将每一个user发送的到broadcast通道的信息，通过hub发送到所有client的通道。
for client := range h.clients {
select {
case client.send <- message:
default:
close(client.send)
delete(h.clients, client)
}
}
}
}
}
client.go
package main
import (
"encoding/json"
"log"
"math/rand"
"net/http"
"strconv"
"time"
"github.com/gorilla/websocket"
)
var randnum = rand.New(rand.NewSource(time.Now().Unix()))
const (
// Time allowed to write a message to the peer.
writeWait = 10 * time.Second
// Time allowed to read the next pong message from the peer.
pongWait = 60 * time.Second
// Send pings to peer with this period. Must be less than pongWait.
pingPeriod = (pongWait * 9) / 10
// Maximum message size allowed from peer.
maxMessageSize = 512
)
var upgrader = websocket.Upgrader{
ReadBufferSize: 1024,
WriteBufferSize: 1024,
}
// Client is a middleman between the websocket connection and the hub.
type Client struct {
// The websocket connection.
conn *websocket.Conn
// Buffered channel of outbound messages.
send chan []byte
Data *Data
}
// readPump pumps messages from the websocket connection to the hub.
//
// The application runs readPump in a per-connection goroutine. The application
// ensures that there is at most one reader on a connection by executing all
// reads from this goroutine.
func (c *Client) readPump() {
defer func() {
h.unregister <- c
c.conn.Close()
}()
c.conn.SetReadLimit(maxMessageSize)
c.conn.SetReadDeadline(time.Now().Add(pongWait))
c.conn.SetPongHandler(func(string) error { c.conn.SetReadDeadline(time.Now().Add(pongWait)); return nil })
for {
_, message, err := c.conn.ReadMessage()
if err != nil {
if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseAbnormalClosure) {
log.Printf("error: %v", err)
}
break
}
// 从websocket收到的数据是网页端用户输入的数据，属于data中的content
c.Data.Content = string(message)
// 将从c.data编码成json,传入到hub的broadcast通道。
message, _ = json.Marshal(c.Data)
h.broadcast <- message
}
}
// writePump pumps messages from the hub to the websocket connection.
//
// A goroutine running writePump is started for each connection. The
// application ensures that there is at most one writer to a connection by
// executing all writes from this goroutine.
func (c *Client) writePump() {
ticker := time.NewTicker(pingPeriod)// ticker是一个定时发送的定时器。
defer func() {
ticker.Stop()
c.conn.Close()
}()
for {
select {
case message, ok := <-c.send:
c.conn.SetWriteDeadline(time.Now().Add(writeWait))
if !ok {
// The hub closed the channel.
c.conn.WriteMessage(websocket.CloseMessage, []byte{})
return
}
// 将从send收到的信息，更新到该client的data中。
data := &Data{}
_ = json.Unmarshal(message, data)
c.Data.Type = "broadcast"
c.Data.Form = data.User
c.Data.Content = data.Content
c.Data.FormIP = data.IP
message, _ = json.Marshal(c.Data)
//发送数据，将数据发送到websocket协议的另一端。
c.conn.WriteMessage(websocket.TextMessage,message)
case <-ticker.C:
//定时器,websocket的心跳机制。
c.conn.SetWriteDeadline(time.Now().Add(writeWait))
if err := c.conn.WriteMessage(websocket.PingMessage, nil); err != nil {
return
}
}
}
}
// serveWs handles websocket requests from the peer.
func serveWs(w http.ResponseWriter, r *http.Request) {
conn, err := upgrader.Upgrade(w, r, nil)
if err != nil {
log.Println(err)
return
}
data := &Data{
IP: r.RemoteAddr,
Type: "handshake",
User: "user" + strconv.Itoa(randnum.Intn(10)),
Content: "xxx上线了",
}
client := &Client{conn: conn, send: make(chan []byte, 256), Data: data}
message, _ := json.Marshal(client.Data)
conn.WriteMessage(websocket.TextMessage, message)
h.register <- client
// Allow collection of memory referenced by the caller by doing all work in
// new goroutines.
go client.writePump()
go client.readPump()
}
main.go
package main
import (
"flag"
"log"
"net/http"
)
var addr = flag.String("addr", ":8080", "http service address")
func serveHome(w http.ResponseWriter, r *http.Request) {
log.Println(r.URL)
if r.URL.Path != "/" {
http.Error(w, "Not found", http.StatusNotFound)
return
}
if r.Method != "GET" {
http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
return
}
http.ServeFile(w, r, "home.html")
}
func main() {
flag.Parse()
go h.run()
http.HandleFunc("/", serveHome)
http.HandleFunc("/ws", func(w http.ResponseWriter, r *http.Request) {
serveWs( w, r)
})
err := http.ListenAndServe(*addr, nil)
if err != nil {
log.Fatal("ListenAndServe: ", err)
}
}
home.html
<html lang="en">
<head>
<title>Chat Exampletitle>
<script type="text/javascript">
window.onload = function () {
var conn;
var msg = document.getElementById("msg");
var log = document.getElementById("log");
function appendLog(item) {
var doScroll = log.scrollTop > log.scrollHeight - log.clientHeight - 1;
log.appendChild(item);
if (doScroll) {
log.scrollTop = log.scrollHeight - log.clientHeight;
}
}
document.getElementById("form").onsubmit = function () {
if (!conn) {
return false;
}
if (!msg.value) {
return false;
}
conn.send(msg.value);
msg.value = "";
return false;
};
if (window["WebSocket"]) {
conn = new WebSocket("ws://" + document.location.host + "/ws");
conn.onclose = function (evt) {
var item = document.createElement("div");
item.innerHTML = "Connection closed.";
appendLog(item);
};
conn.onmessage = function (evt) {
var messages = evt.data.split('\n');
for (var i = 0; i < messages.length; i++) {
var item = document.createElement("div");
item.innerText = messages[i];
appendLog(item);
}
};
} else {
var item = document.createElement("div");
item.innerHTML = "Your browser does not support WebSockets.";
appendLog(item);
}
};
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="Send" />
<input type="text" id="msg" size="64" autofocus />
form>
body>
html>
运行效果如图所示：
