Skip to content

Commit

Permalink
chore: migrate from gorilla/websocket to gobwas/ws
Browse files Browse the repository at this point in the history
  • Loading branch information
wwqgtxx committed Oct 6, 2023
1 parent d1e88a3 commit 5ff4473
Show file tree
Hide file tree
Showing 6 changed files with 186 additions and 142 deletions.
9 changes: 4 additions & 5 deletions component/tls/utls.go
Original file line number Diff line number Diff line change
Expand Up @@ -99,10 +99,9 @@ func copyConfig(c *tls.Config) *utls.Config {
}
}

// WebsocketHandshake basically calls UConn.Handshake inside it but it will only send
// http/1.1 in its ALPN.
// BuildWebsocketHandshakeState it will only send http/1.1 in its ALPN.
// Copy from https://github.com/XTLS/Xray-core/blob/main/transport/internet/tls/tls.go
func (c *UConn) WebsocketHandshake() error {
func (c *UConn) BuildWebsocketHandshakeState() error {
// Build the handshake state. This will apply every variable of the TLS of the
// fingerprint in the UConn
if err := c.BuildHandshakeState(); err != nil {
Expand All @@ -120,11 +119,11 @@ func (c *UConn) WebsocketHandshake() error {
if !hasALPNExtension { // Append extension if doesn't exists
c.Extensions = append(c.Extensions, &utls.ALPNExtension{AlpnProtocols: []string{"http/1.1"}})
}
// Rebuild the client hello and do the handshake
// Rebuild the client hello
if err := c.BuildHandshakeState(); err != nil {
return err
}
return c.Handshake()
return nil
}

func SetGlobalUtlsClient(Client string) {
Expand Down
4 changes: 3 additions & 1 deletion go.mod
Original file line number Diff line number Diff line change
Expand Up @@ -11,8 +11,8 @@ require (
github.com/go-chi/chi/v5 v5.0.10
github.com/go-chi/cors v1.2.1
github.com/go-chi/render v1.0.3
github.com/gobwas/ws v1.3.0
github.com/gofrs/uuid/v5 v5.0.0
github.com/gorilla/websocket v1.5.0
github.com/insomniacslk/dhcp v0.0.0-20230908212754-65c27093e38a
github.com/jpillora/backoff v1.0.0
github.com/klauspost/cpuid/v2 v2.2.5
Expand Down Expand Up @@ -69,6 +69,8 @@ require (
github.com/gaukas/godicttls v0.0.4 // indirect
github.com/go-ole/go-ole v1.3.0 // indirect
github.com/go-task/slim-sprig v0.0.0-20230315185526-52ccab3ef572 // indirect
github.com/gobwas/httphead v0.1.0 // indirect
github.com/gobwas/pool v0.2.1 // indirect
github.com/google/btree v1.1.2 // indirect
github.com/google/go-cmp v0.5.9 // indirect
github.com/google/pprof v0.0.0-20210407192527-94a9f03dee38 // indirect
Expand Down
9 changes: 7 additions & 2 deletions go.sum
Original file line number Diff line number Diff line change
Expand Up @@ -49,6 +49,12 @@ github.com/go-ole/go-ole v1.3.0 h1:Dt6ye7+vXGIKZ7Xtk4s6/xVdGDQynvom7xCFEdWr6uE=
github.com/go-ole/go-ole v1.3.0/go.mod h1:5LS6F96DhAwUc7C+1HLexzMXY1xGRSryjyPPKW6zv78=
github.com/go-task/slim-sprig v0.0.0-20230315185526-52ccab3ef572 h1:tfuBGBXKqDEevZMzYi5KSi8KkcZtzBcTgAUUtapy0OI=
github.com/go-task/slim-sprig v0.0.0-20230315185526-52ccab3ef572/go.mod h1:9Pwr4B2jHnOSGXyyzV8ROjYa2ojvAY6HCGYYfMoC3Ls=
github.com/gobwas/httphead v0.1.0 h1:exrUm0f4YX0L7EBwZHuCF4GDp8aJfVeBrlLQrs6NqWU=
github.com/gobwas/httphead v0.1.0/go.mod h1:O/RXo79gxV8G+RqlR/otEwx4Q36zl9rqC5u12GKvMCM=
github.com/gobwas/pool v0.2.1 h1:xfeeEhW7pwmX8nuLVlqbzVc7udMDrwetjEv+TZIz1og=
github.com/gobwas/pool v0.2.1/go.mod h1:q8bcK0KcYlCgd9e7WYLm9LpyS+YeLd8JVDW6WezmKEw=
github.com/gobwas/ws v1.3.0 h1:sbeU3Y4Qzlb+MOzIe6mQGf7QR4Hkv6ZD0qhGkBFL2O0=
github.com/gobwas/ws v1.3.0/go.mod h1:hRKAFb8wOxFROYNsT1bqfWnhX+b5MFeJM9r2ZSwg/KY=
github.com/gofrs/uuid/v5 v5.0.0 h1:p544++a97kEL+svbcFbCQVM9KFu0Yo25UoISXGNNH9M=
github.com/gofrs/uuid/v5 v5.0.0/go.mod h1:CDOjlDMVAtN56jqyRUZh58JT31Tiw7/oQyEXZV+9bD8=
github.com/golang/protobuf v1.3.1/go.mod h1:6lQm79b+lXiMfvg/cZm0SGofjICqVBUtrP5yJMmIC1U=
Expand All @@ -63,8 +69,6 @@ github.com/google/go-cmp v0.5.9/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeN
github.com/google/pprof v0.0.0-20210407192527-94a9f03dee38 h1:yAJXTCF9TqKcTiHJAE8dj7HMvPfh66eeA2JYW7eFpSE=
github.com/google/pprof v0.0.0-20210407192527-94a9f03dee38/go.mod h1:kpwsk12EmLew5upagYY7GY0pfYCcupk39gWOCRROcvE=
github.com/google/tink/go v1.6.1 h1:t7JHqO8Ath2w2ig5vjwQYJzhGEZymedQc90lQXUBa4I=
github.com/gorilla/websocket v1.5.0 h1:PPwGk2jz7EePpoHN/+ClbZu8SPxiqlu12wZP/3sWmnc=
github.com/gorilla/websocket v1.5.0/go.mod h1:YR8l580nyteQvAITg2hZ9XVh4b55+EU/adAjf1fMHhE=
github.com/hashicorp/yamux v0.1.1 h1:yrQxtgseBDrq9Y652vSRDvsKCJKOUD+GzTS4Y0Y8pvE=
github.com/hashicorp/yamux v0.1.1/go.mod h1:CtWFDAQgb7dxtzFs4tWbplKIe2jSi3+5vKbgIO0SLnQ=
github.com/ianlancetaylor/demangle v0.0.0-20200824232613-28f6c0f3b639/go.mod h1:aSSvb/t6k1mPoxDqO4vJh6VOCGPwU4O0C2/Eqndh1Sc=
Expand Down Expand Up @@ -237,6 +241,7 @@ golang.org/x/sys v0.0.0-20220715151400-c0bba94af5f8/go.mod h1:oPkhp1MJrh7nUepCBc
golang.org/x/sys v0.0.0-20220908164124-27713097b956/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
golang.org/x/sys v0.1.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
golang.org/x/sys v0.5.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
golang.org/x/sys v0.8.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
golang.org/x/sys v0.11.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
golang.org/x/sys v0.12.0 h1:CM0HF96J0hcLAwsHPJZjfdNzs0gftsLfgKt57wWHJ0o=
Expand Down
9 changes: 5 additions & 4 deletions hub/route/connections.go
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,8 @@ import (

"github.com/go-chi/chi/v5"
"github.com/go-chi/render"
"github.com/gorilla/websocket"
"github.com/gobwas/ws"
"github.com/gobwas/ws/wsutil"
)

func connectionRouter() http.Handler {
Expand All @@ -23,13 +24,13 @@ func connectionRouter() http.Handler {
}

func getConnections(w http.ResponseWriter, r *http.Request) {
if !websocket.IsWebSocketUpgrade(r) {
if !(r.Header.Get("Upgrade") == "websocket") {
snapshot := statistic.DefaultManager.Snapshot()
render.JSON(w, r, snapshot)
return
}

conn, err := upgrader.Upgrade(w, r, nil)
conn, _, _, err := ws.UpgradeHTTP(r, w)
if err != nil {
return
}
Expand All @@ -55,7 +56,7 @@ func getConnections(w http.ResponseWriter, r *http.Request) {
return err
}

return conn.WriteMessage(websocket.TextMessage, buf.Bytes())
return wsutil.WriteMessage(conn, ws.StateServerSide, ws.OpText, buf.Bytes())
}

if err := sendSnapshot(); err != nil {
Expand Down
36 changes: 16 additions & 20 deletions hub/route/server.go
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@ import (
"crypto/subtle"
"crypto/tls"
"encoding/json"
"net"
"net/http"
"runtime/debug"
"strings"
Expand All @@ -21,20 +22,15 @@ 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 (
serverSecret = ""
serverAddr = ""

uiPath = ""

upgrader = websocket.Upgrader{
CheckOrigin: func(r *http.Request) bool {
return true
},
}
)

type Traffic struct {
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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
}
Expand Down Expand Up @@ -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 {
Expand All @@ -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
}
Expand Down Expand Up @@ -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 {
Expand Down Expand Up @@ -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
}
Expand Down Expand Up @@ -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 {
Expand Down
Loading

0 comments on commit 5ff4473

Please sign in to comment.