forked from lqqyt2423/go-mitmproxy
/
websocket.go
67 lines (55 loc) · 1.27 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
package proxy
import (
"crypto/tls"
"net"
"net/http"
"net/http/httputil"
"strings"
log "github.com/sirupsen/logrus"
)
// 当前仅做了转发 websocket 流量
type webSocket struct{}
var defaultWebSocket webSocket
func (s *webSocket) ws(conn net.Conn, host string) {
log := log.WithField("in", "webSocket.ws").WithField("host", host)
defer conn.Close()
remoteConn, err := net.Dial("tcp", host)
if err != nil {
logErr(log, err)
return
}
defer remoteConn.Close()
transfer(log, conn, remoteConn)
}
func (s *webSocket) wss(res http.ResponseWriter, req *http.Request) {
log := log.WithField("in", "webSocket.wss").WithField("host", req.Host)
upgradeBuf, err := httputil.DumpRequest(req, false)
if err != nil {
log.Errorf("DumpRequest: %v\n", err)
res.WriteHeader(502)
return
}
cconn, _, err := res.(http.Hijacker).Hijack()
if err != nil {
log.Errorf("Hijack: %v\n", err)
res.WriteHeader(502)
return
}
defer cconn.Close()
host := req.Host
if !strings.Contains(host, ":") {
host = host + ":443"
}
conn, err := tls.Dial("tcp", host, nil)
if err != nil {
log.Errorf("tls.Dial: %v\n", err)
return
}
defer conn.Close()
_, err = conn.Write(upgradeBuf)
if err != nil {
log.Errorf("wss upgrade: %v\n", err)
return
}
transfer(log, conn, cconn)
}