2020package transport
2121
2222import (
23- "bytes"
2423 "context"
2524 "crypto/tls"
26- "errors"
27- "log"
2825 "net"
2926 "net/http"
3027 "time"
3128
32- "github.com/gorilla/websocket"
29+ websocket "github.com/gobwas/ws"
30+ "github.com/gobwas/ws/wsutil"
3331)
3432
35- const (
36- // Time allowed to write a message to the peer.
37- writeWait = 10 * time .Second
38- // Time allowed to read the next pong message from the peer.
39- pongWait = 60 * time .Second
40-
41- // Send pings to peer with this period. Must be less than pongWait.
42- pingPeriod = (pongWait * 9 ) / 10
43-
44- // Maximum message size allowed from peer.
45- maxMessageSize = 512
46- )
47-
48- var (
49- newline = []byte {'\n' }
50- space = []byte {' ' }
51- errUpgradingConn = errors .New ("encountered error upgrading connection to websocket protocol" )
52- errStartingServer = errors .New ("encountered error starting http server" )
53- )
54-
55- var upgrader = websocket.Upgrader {
56- ReadBufferSize : 1024 ,
57- WriteBufferSize : 1024 ,
58- CheckOrigin : func (r * http.Request ) bool { return true },
59- }
60-
6133type WebSocket struct {
62- Bind string
63- TLS * tls.Config
64- CORSOrigin string
65- Parallel bool
34+ Bind string
35+ TLS * tls.Config
36+ CORSOrigin string
37+ Parallel bool
38+ ReadDeadline , WriteDeadline time.Duration //Set custom timeout for future read and write calls
6639}
6740
41+ func (ws * WebSocket ) WithReadDealine () bool { return ws .ReadDeadline != 0 }
42+ func (ws * WebSocket ) WithWriteDealine () bool { return ws .WriteDeadline != 0 }
43+
6844func (ws * WebSocket ) Run (ctx context.Context , resolver Resolver ) error {
6945 srv := http.Server {
7046 Addr : ws .Bind ,
7147 Handler : http .HandlerFunc (func (w http.ResponseWriter , r * http.Request ) {
72-
73- wsconn , err := upgrader .Upgrade (w , r , nil )
48+ wsconn , _ , _ , err := websocket .UpgradeHTTP (r , w )
7449 if err != nil {
75- log .Println (err )
7650 return
7751 }
7852
79- log .Println ("successfully upgraded connection" )
80-
8153 defer func () {
8254 wsconn .Close ()
8355 }()
8456
85- wsconn .SetReadLimit (maxMessageSize )
86- wsconn .SetReadDeadline (time .Now ().Add (pongWait ))
87- wsconn .SetPongHandler (func (string ) error {
88- wsconn .SetReadDeadline (time .Now ().Add (pongWait ))
89- return nil
90- })
57+ if ws .WithReadDealine () {
58+ wsconn .SetReadDeadline (time .Now ().Add (ws .ReadDeadline * time .Second ))
59+ }
60+
61+ if ws .WithWriteDealine () {
62+ wsconn .SetWriteDeadline (time .Now ().Add (ws .WriteDeadline * time .Second ))
63+ }
9164
9265 for {
66+
9367 // read message from connection
94- messageType , message , err := wsconn . ReadMessage ( )
68+ _ , reader , err := wsutil . NextReader ( wsconn , websocket . StateServerSide )
9569 if err != nil {
96- if websocket .IsUnexpectedCloseError (err , websocket .CloseGoingAway , websocket .CloseAbnormalClosure ) {
97- log .Printf ("error: %v" , err )
98- }
99- break
70+ return
10071 }
10172
102- message = bytes .TrimSpace (bytes .Replace (message , newline , space , - 1 ))
73+ // create writer object that implements io.WriterCloser interface
74+ writer := wsutil .NewWriter (wsconn , websocket .StateServerSide , websocket .OpText )
10375
104- wsconn . SetWriteDeadline ( time . Now (). Add ( writeWait ) )
76+ resolver . Resolve ( ctx , reader , writer , ws . Parallel )
10577
106- // create writer object that implements io.WriterCloser interface
107- // messageType is same as the messageType recieved from the connection
108- w , err := wsconn .NextWriter (messageType )
109- if err != nil {
78+ if err := writer .Flush (); err != nil {
11079 return
11180 }
11281
113- resolver .Resolve (ctx , bytes .NewBuffer (message ), w , ws .Parallel )
11482 }
83+
11584 }),
11685
11786 BaseContext : func (l net.Listener ) context.Context {
@@ -125,7 +94,6 @@ func (ws *WebSocket) Run(ctx context.Context, resolver Resolver) error {
12594 }()
12695
12796 if err := srv .ListenAndServe (); err != http .ErrServerClosed {
128- log .Println (err )
12997 return err
13098 }
13199 return nil
0 commit comments