Hertz WebSocket 扩展
1. WebSocket
websocket是一种全双工通信,且只需要一次握手。
2.代码
package main
import (
"context"
"flag"
"github.com/cloudwego/hertz/pkg/app"
"github.com/cloudwego/hertz/pkg/app/server"
"github.com/hertz-contrib/websocket"
"html/template"
"log"
)
var addr = flag.String("addr", "localhost:8080", "http service address")
var upgrader = websocket.HertzUpgrader{}
func home(_ context.Context, c *app.RequestContext) {
c.SetContentType("text/html;charset=utf-8") //设置HTTP响应的Content-Type头部为text/html
// 将homeTemplate的html渲染,并写入到http响应中
// 第一个参数是上下文对象,第二个参数是websocker连接的URI
homeTemplate.Execute(c, "ws://"+string(c.Host())+"/echo")
}
func echo(_ context.Context, c *app.RequestContext) {
// upgrader.Upgrade 将HTTP协议升级为WebSocket协议的连接请求。
// websocket.Conn 代表一个WebSocket连接
// 第二个参数HertzHandler,它在握手完成后接收一个websocket连接,通过劫持这个连接完成全双工通信,必须由用户实现。。
// func (u *HertzUpgrader) Upgrade(ctx *app.RequestContext, handler HertzHandler) error
err := upgrader.Upgrade(c, func(conn *websocket.Conn) {
for {
// 读取客户端发送的信息
mt, message, err := conn.ReadMessage()
if err != nil {
log.Println("read: ", err)
break
}
log.Printf("recv: %s", message)
// 向客户端发送信息
err = conn.WriteMessage(mt, message)
if err != nil {
log.Println("write: ", err)
break
}
}
})
if err != nil {
log.Println("upgrade:", err)
return
}
}
func main() {
flag.Parse()
h := server.Default(server.WithHostPorts(*addr))
h.NoHijackConnPool = true
h.GET("/", home)
h.GET("/echo", echo)
h.Spin()
}
var homeTemplate = template.Must(template.New("").Parse(`
<!DOCTYPE html>
<html>
<head>
<meta charset="utf-8">
<script>
window.addEventListener("load", function(evt) {
var output = document.getElementById("output");
var input = document.getElementById("input");
var ws;
var print = function(message) {
var d = document.createElement("div");
d.textContent = message;
output.appendChild(d);
output.scroll(0, output.scrollHeight);
};
document.getElementById("open").onclick = function(evt) {
if (ws) {
return false;
}
ws = new WebSocket("{{.}}");
ws.onopen = function(evt) {
print("OPEN");
}
ws.onclose = function(evt) {
print("CLOSE");
ws = null;
}
ws.onmessage = function(evt) {
print("RESPONSE: " + evt.data);
}
ws.onerror = function(evt) {
print("ERROR: " + evt.data);
}
return false;
};
document.getElementById("send").onclick = function(evt) {
if (!ws) {
return false;
}
print("SEND: " + input.value);
ws.send(input.value);
return false;
};
document.getElementById("close").onclick = function(evt) {
if (!ws) {
return false;
}
ws.close();
return false;
};
});
</script>
</head>
<body>
<table>
<tr><td valign="top" width="50%">
<p>Click "Open" to create a connection to the server,
"Send" to send a message to the server and "Close" to close the connection.
You can change the message and send multiple times.
<p>
<form>
<button id="open">Open</button>
<button id="close">Close</button>
<p><input id="input" type="text" value="Hello world!">
<button id="send">Send</button>
</form>
</td><td valign="top" width="50%">
<div id="output" style="max-height: 70vh;overflow-y: scroll;"></div>
</td></tr></table>
</body>
</html>
`))
运行成功后,在127.0.0.1:8080
即可访问。
参考:https://juejin.cn/post/7176901666791227450,详细可以在这里看。
3.总结
通过websocket.HertzUpgrader
类中的Upgrade
方法将http升级为websocket。
在借助 Upgrade
的 HertzHandler
实现劫持连接之后实现的操作,其中的操作需要用户实现。