/
websocket.go
153 lines (128 loc) · 3.53 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
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
package service
import (
"context"
"crypto/tls"
"log/slog"
"net"
"net/http"
"time"
"github.com/gorilla/websocket"
"github.com/DVKunion/SeaMoon/pkg/consts"
"github.com/DVKunion/SeaMoon/pkg/transfer"
"github.com/DVKunion/SeaMoon/pkg/tunnel"
)
const (
defaultTimeout = 10 * time.Second
defaultReadBufferSize = 32 * 1024
)
type WSService struct {
startAt time.Time
upGrader *websocket.Upgrader
}
func init() {
register(tunnel.WST, &WSService{})
}
func (s *WSService) Conn(ctx context.Context, t transfer.Type, sOpts ...Option) (net.Conn, error) {
// todo: useless ctx
var srvOpts = &Options{}
for _, o := range sOpts {
o(srvOpts)
}
wsDialer := &websocket.Dialer{
HandshakeTimeout: defaultTimeout,
ReadBufferSize: defaultReadBufferSize,
WriteBufferSize: defaultReadBufferSize,
EnableCompression: false,
}
if srvOpts.buffers != nil {
wsDialer.ReadBufferSize = srvOpts.buffers.ReadBufferSize
wsDialer.WriteBufferSize = srvOpts.buffers.WriteBufferSize
wsDialer.EnableCompression = srvOpts.buffers.EnableCompression
}
var requestHeader = http.Header{}
if srvOpts.tor {
requestHeader.Add("SM-ONION", "enable")
}
wsConn, _, err := wsDialer.Dial(t.Path(srvOpts.addr), requestHeader)
if err != nil {
return nil, err
}
return tunnel.WsWrapConn(wsConn), nil
}
func (s *WSService) Serve(ln net.Listener, sOpts ...Option) error {
var srvOpts = &Options{}
for _, o := range sOpts {
o(srvOpts)
}
s.upGrader = &websocket.Upgrader{
HandshakeTimeout: defaultTimeout,
ReadBufferSize: defaultReadBufferSize,
WriteBufferSize: defaultReadBufferSize,
EnableCompression: false,
CheckOrigin: func(r *http.Request) bool { return true },
}
if srvOpts.buffers != nil {
s.upGrader.ReadBufferSize = srvOpts.buffers.ReadBufferSize
s.upGrader.WriteBufferSize = srvOpts.buffers.WriteBufferSize
s.upGrader.EnableCompression = srvOpts.buffers.EnableCompression
}
if srvOpts.tlsConf != nil {
ln = tls.NewListener(ln, srvOpts.tlsConf)
}
mux := http.NewServeMux()
// websocket http proxy handler
mux.HandleFunc("/http", s.http)
// websocket socks5 proxy handler
mux.HandleFunc("/socks5", s.socks5)
s.startAt = time.Now()
// inject
mux.HandleFunc("/_health", s.health)
server := &http.Server{
Addr: srvOpts.addr,
Handler: mux,
}
return server.Serve(ln)
}
func (s *WSService) http(w http.ResponseWriter, r *http.Request) {
// means use http to connector
conn, err := s.upGrader.Upgrade(w, r, nil)
if err != nil {
return
}
t := tunnel.WsWrapConn(conn)
go func() {
if err := transfer.HttpTransport(t); err != nil {
slog.Error("connection error", "msg", err)
}
}()
}
func (s *WSService) socks5(w http.ResponseWriter, r *http.Request) {
onion := r.Header.Get("SM-ONION")
// means use socks5 to connector
conn, err := s.upGrader.Upgrade(w, r, nil)
if err != nil {
return
}
t := tunnel.WsWrapConn(conn)
go func() {
// 检测是否存在 onion 标识,代表着是否要开启 tor 转发
if onion != "" {
if err := transfer.TorTransport(t); err != nil {
slog.Error("connection error", "msg", err)
}
} else {
if err := transfer.Socks5Transport(t); err != nil {
slog.Error("connection error", "msg", err)
}
}
}()
}
func (s *WSService) health(w http.ResponseWriter, _ *http.Request) {
w.Header().Set("Content-Type", "text/plain")
w.WriteHeader(http.StatusOK)
_, err := w.Write([]byte("OK\n" + s.startAt.Format("2006-01-02 15:04:05") + "\n" + consts.Version + "-" + consts.Commit))
if err != nil {
slog.Error("server status error", "msg", err)
return
}
}