Skip to content

Commit 6e786f2

Browse files
committed
Add proxy support
1 parent 98abd3f commit 6e786f2

File tree

2 files changed

+49
-16
lines changed

2 files changed

+49
-16
lines changed

hbdm/nws.go

Lines changed: 24 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@ import (
1111
"github.com/recws-org/recws"
1212
"github.com/tidwall/gjson"
1313
"log"
14+
"net/http"
1415
"net/url"
1516
"strings"
1617
"sync"
@@ -23,7 +24,7 @@ type NWS struct {
2324

2425
ctx context.Context
2526
cancel context.CancelFunc
26-
wsConn recws.RecConn
27+
conn recws.RecConn
2728

2829
wsURL string
2930
cid string // 客户端请求唯一ID
@@ -38,6 +39,21 @@ type NWS struct {
3839
liquidationOrdersCallback func(liquidationOrders *WSLiquidationOrders)
3940
}
4041

42+
// SetProxy 设置代理地址
43+
// porxyURL:
44+
// socks5://127.0.0.1:1080
45+
// https://127.0.0.1:1080
46+
func (ws *NWS) SetProxy(proxyURL string) (err error) {
47+
var purl *url.URL
48+
purl, err = url.Parse(proxyURL)
49+
if err != nil {
50+
return
51+
}
52+
log.Printf("[ws][%s] proxy url:%s", proxyURL, purl)
53+
ws.conn.Proxy = http.ProxyURL(purl)
54+
return
55+
}
56+
4157
// SubscribeOrders 订阅订单成交数据
4258
// symbol: BTC
4359
func (ws *NWS) SubscribeOrders(id string, symbol string) {
@@ -207,11 +223,11 @@ func (ws *NWS) subscribeHandler() error {
207223
}
208224

209225
func (ws *NWS) sendWSMessage(msg interface{}) error {
210-
return ws.wsConn.WriteJSON(msg)
226+
return ws.conn.WriteJSON(msg)
211227
}
212228

213229
func (ws *NWS) Start() {
214-
ws.wsConn.Dial(ws.wsURL, nil)
230+
ws.conn.Dial(ws.wsURL, nil)
215231
go ws.run()
216232
}
217233

@@ -220,11 +236,11 @@ func (ws *NWS) run() {
220236
for {
221237
select {
222238
case <-ctx.Done():
223-
go ws.wsConn.Close()
224-
log.Printf("Websocket closed %s", ws.wsConn.GetURL())
239+
go ws.conn.Close()
240+
log.Printf("Websocket closed %s", ws.conn.GetURL())
225241
return
226242
default:
227-
messageType, msg, err := ws.wsConn.ReadMessage()
243+
messageType, msg, err := ws.conn.ReadMessage()
228244
if err != nil {
229245
log.Printf("Read error: %v", err)
230246
time.Sleep(100 * time.Millisecond)
@@ -371,9 +387,9 @@ func NewNWS(wsURL string, accessKey string, secretKey string) *NWS {
371387
subscriptions: make(map[string]interface{}),
372388
}
373389
ws.ctx, ws.cancel = context.WithCancel(context.Background())
374-
ws.wsConn = recws.RecConn{
390+
ws.conn = recws.RecConn{
375391
KeepAliveTimeout: 10 * time.Second,
376392
}
377-
ws.wsConn.SubscribeHandler = ws.subscribeHandler
393+
ws.conn.SubscribeHandler = ws.subscribeHandler
378394
return ws
379395
}

hbdmswap/ws.go

Lines changed: 25 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,8 @@ import (
1010
"github.com/recws-org/recws"
1111
"github.com/tidwall/gjson"
1212
"log"
13+
"net/http"
14+
"net/url"
1315
"strings"
1416
"sync"
1517
"time"
@@ -21,7 +23,7 @@ type WS struct {
2123

2224
ctx context.Context
2325
cancel context.CancelFunc
24-
wsConn recws.RecConn
26+
conn recws.RecConn
2527

2628
wsURL string
2729
accessKey string
@@ -33,6 +35,21 @@ type WS struct {
3335
tradeCallback func(trade *WSTrade)
3436
}
3537

38+
// SetProxy 设置代理地址
39+
// porxyURL:
40+
// socks5://127.0.0.1:1080
41+
// https://127.0.0.1:1080
42+
func (ws *WS) SetProxy(proxyURL string) (err error) {
43+
var purl *url.URL
44+
purl, err = url.Parse(proxyURL)
45+
if err != nil {
46+
return
47+
}
48+
log.Printf("[ws][%s] proxy url:%s", proxyURL, purl)
49+
ws.conn.Proxy = http.ProxyURL(purl)
50+
return
51+
}
52+
3653
// SubscribeTicker 订阅 Market Ticker 数据
3754
// id: 订阅的编号
3855
// symbol: BTC_CQ
@@ -107,11 +124,11 @@ func (ws *WS) subscribeHandler() error {
107124
}
108125

109126
func (ws *WS) sendWSMessage(msg interface{}) error {
110-
return ws.wsConn.WriteJSON(msg)
127+
return ws.conn.WriteJSON(msg)
111128
}
112129

113130
func (ws *WS) Start() {
114-
ws.wsConn.Dial(ws.wsURL, nil)
131+
ws.conn.Dial(ws.wsURL, nil)
115132
go ws.run()
116133
}
117134

@@ -120,11 +137,11 @@ func (ws *WS) run() {
120137
for {
121138
select {
122139
case <-ctx.Done():
123-
go ws.wsConn.Close()
124-
log.Printf("Websocket closed %s", ws.wsConn.GetURL())
140+
go ws.conn.Close()
141+
log.Printf("Websocket closed %s", ws.conn.GetURL())
125142
return
126143
default:
127-
messageType, msg, err := ws.wsConn.ReadMessage()
144+
messageType, msg, err := ws.conn.ReadMessage()
128145
if err != nil {
129146
log.Printf("Read error: %v", err)
130147
time.Sleep(100 * time.Millisecond)
@@ -225,9 +242,9 @@ func NewWS(wsURL string, accessKey string, secretKey string) *WS {
225242
subscriptions: make(map[string]interface{}),
226243
}
227244
ws.ctx, ws.cancel = context.WithCancel(context.Background())
228-
ws.wsConn = recws.RecConn{
245+
ws.conn = recws.RecConn{
229246
KeepAliveTimeout: 10 * time.Second,
230247
}
231-
ws.wsConn.SubscribeHandler = ws.subscribeHandler
248+
ws.conn.SubscribeHandler = ws.subscribeHandler
232249
return ws
233250
}

0 commit comments

Comments
 (0)