chore: migrate from gorilla/websocket to gobwas/ws

This commit is contained in:
wwqgtxx
2023-10-06 17:44:36 +08:00
parent d1e88a30cb
commit 5ff4473083
6 changed files with 187 additions and 143 deletions

View File

@ -5,6 +5,7 @@ import (
"crypto/subtle"
"crypto/tls"
"encoding/json"
"net"
"net/http"
"runtime/debug"
"strings"
@ -21,7 +22,8 @@ import (
"github.com/go-chi/chi/v5/middleware"
"github.com/go-chi/cors"
"github.com/go-chi/render"
"github.com/gorilla/websocket"
"github.com/gobwas/ws"
"github.com/gobwas/ws/wsutil"
)
var (
@ -29,12 +31,6 @@ var (
serverAddr = ""
uiPath = ""
upgrader = websocket.Upgrader{
CheckOrigin: func(r *http.Request) bool {
return true
},
}
)
type Traffic struct {
@ -166,7 +162,7 @@ func authentication(next http.Handler) http.Handler {
}
// Browser websocket not support custom header
if websocket.IsWebSocketUpgrade(r) && r.URL.Query().Get("token") != "" {
if r.Header.Get("Upgrade") == "websocket" && r.URL.Query().Get("token") != "" {
token := r.URL.Query().Get("token")
if !safeEuqal(token, serverSecret) {
render.Status(r, http.StatusUnauthorized)
@ -197,10 +193,10 @@ func hello(w http.ResponseWriter, r *http.Request) {
}
func traffic(w http.ResponseWriter, r *http.Request) {
var wsConn *websocket.Conn
if websocket.IsWebSocketUpgrade(r) {
var wsConn net.Conn
if r.Header.Get("Upgrade") == "websocket" {
var err error
wsConn, err = upgrader.Upgrade(w, r, nil)
wsConn, _, _, err = ws.UpgradeHTTP(r, w)
if err != nil {
return
}
@ -230,7 +226,7 @@ func traffic(w http.ResponseWriter, r *http.Request) {
_, err = w.Write(buf.Bytes())
w.(http.Flusher).Flush()
} else {
err = wsConn.WriteMessage(websocket.TextMessage, buf.Bytes())
err = wsutil.WriteMessage(wsConn, ws.StateServerSide, ws.OpText, buf.Bytes())
}
if err != nil {
@ -240,10 +236,10 @@ func traffic(w http.ResponseWriter, r *http.Request) {
}
func memory(w http.ResponseWriter, r *http.Request) {
var wsConn *websocket.Conn
if websocket.IsWebSocketUpgrade(r) {
var wsConn net.Conn
if r.Header.Get("Upgrade") == "websocket" {
var err error
wsConn, err = upgrader.Upgrade(w, r, nil)
wsConn, _, _, err = ws.UpgradeHTTP(r, w)
if err != nil {
return
}
@ -280,7 +276,7 @@ func memory(w http.ResponseWriter, r *http.Request) {
_, err = w.Write(buf.Bytes())
w.(http.Flusher).Flush()
} else {
err = wsConn.WriteMessage(websocket.TextMessage, buf.Bytes())
err = wsutil.WriteMessage(wsConn, ws.StateServerSide, ws.OpText, buf.Bytes())
}
if err != nil {
@ -323,10 +319,10 @@ func getLogs(w http.ResponseWriter, r *http.Request) {
return
}
var wsConn *websocket.Conn
if websocket.IsWebSocketUpgrade(r) {
var wsConn net.Conn
if r.Header.Get("Upgrade") == "websocket" {
var err error
wsConn, err = upgrader.Upgrade(w, r, nil)
wsConn, _, _, err = ws.UpgradeHTTP(r, w)
if err != nil {
return
}
@ -385,7 +381,7 @@ func getLogs(w http.ResponseWriter, r *http.Request) {
_, err = w.Write(buf.Bytes())
w.(http.Flusher).Flush()
} else {
err = wsConn.WriteMessage(websocket.TextMessage, buf.Bytes())
err = wsutil.WriteMessage(wsConn, ws.StateServerSide, ws.OpText, buf.Bytes())
}
if err != nil {