forked from elazarl/goproxy
-
Notifications
You must be signed in to change notification settings - Fork 0
/
websocket.go
110 lines (99 loc) · 3.13 KB
/
websocket.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
package goproxy
import (
"bufio"
"crypto/tls"
"fmt"
"net"
"net/http"
"net/http/httptest"
"github.com/gorilla/websocket"
)
type ResponseHijack struct {
*httptest.ResponseRecorder
br *bufio.Reader
conn *tls.Conn
}
func (r *ResponseHijack) Hijack() (net.Conn, *bufio.ReadWriter, error) {
rw := bufio.NewReadWriter(r.br, bufio.NewWriter(r.conn))
return r.conn, rw, nil
}
func (proxy *ProxyHttpServer) websocket(ctx *ProxyCtx, req *http.Request, br *bufio.Reader, rawClientTls *tls.Conn) {
rw := &ResponseHijack{httptest.NewRecorder(), br, rawClientTls}
requestHeader := http.Header{}
if origin := req.Header.Get("Origin"); origin != "" {
requestHeader.Add("Origin", origin)
}
for _, prot := range req.Header[http.CanonicalHeaderKey("Sec-WebSocket-Protocol")] {
requestHeader.Add("Sec-WebSocket-Protocol", prot)
}
for _, cookie := range req.Header[http.CanonicalHeaderKey("Cookie")] {
requestHeader.Add("Cookie", cookie)
}
for _, target := range req.Header[http.CanonicalHeaderKey("X-Wise-Target")] {
requestHeader.Add("X-Wise-Target", target)
}
req.URL.Scheme = "wss"
req.URL.Host = req.Host
url := req.URL.String()
ctx.Logf("websocket: connect to %s", url)
connBackend, resp, err := websocket.DefaultDialer.Dial(url, requestHeader)
if err != nil {
ctx.Warnf("websocket: dial error %s", err)
return
}
defer connBackend.Close()
upgradeHeader := http.Header{}
if hdr := resp.Header.Get("Sec-Websocket-Protocol"); hdr != "" {
upgradeHeader.Set("Sec-Websocket-Protocol", hdr)
}
if hdr := resp.Header.Get("Set-Cookie"); hdr != "" {
upgradeHeader.Set("Set-Cookie", hdr)
}
connPub, err := websocket.Upgrade(rw, req, upgradeHeader, 1024, 1024)
if err != nil {
ctx.Warnf("websocket: can not upgrade %s", err)
}
defer connPub.Close()
errClient := make(chan error, 1)
errBackend := make(chan error, 1)
replicateWebsocketConn := func(dst, src *websocket.Conn, errc chan error, typ int) {
for {
msgType, msg, err := src.ReadMessage()
if err == nil && proxy.handleWebsocket != nil {
msg, err = proxy.handleWebsocket(url, typ, msgType, msg)
}
if err != nil {
m := websocket.FormatCloseMessage(websocket.CloseNormalClosure, fmt.Sprintf("%v", err))
if e, ok := err.(*websocket.CloseError); ok {
if e.Code != websocket.CloseNoStatusReceived {
m = websocket.FormatCloseMessage(e.Code, e.Text)
}
}
errc <- err
dst.WriteMessage(websocket.CloseMessage, m)
break
}
err = dst.WriteMessage(msgType, msg)
if err != nil {
errc <- err
break
}
}
}
go replicateWebsocketConn(connPub, connBackend, errClient, 1)
go replicateWebsocketConn(connBackend, connPub, errBackend, 0)
var message string
select {
case err = <-errClient:
message = "websocket: Error when copying from backend to client: %v"
case err = <-errBackend:
message = "websocket: Error when copying from client to backend: %v"
}
if e, ok := err.(*websocket.CloseError); !ok || e.Code == websocket.CloseAbnormalClosure {
ctx.Warnf(message, err)
}
return
}
func (proxy *ProxyHttpServer) SetWebsocketHandler(f func(url string, typ, msgType int, msg []byte) ([]byte, error)) {
proxy.handleWebsocket = f
}