-
Notifications
You must be signed in to change notification settings - Fork 3
/
conn.go
121 lines (96 loc) · 2.01 KB
/
conn.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
package fastws
import (
"github.com/clubpay/ronykit/kit"
"github.com/clubpay/ronykit/kit/errors"
"github.com/clubpay/ronykit/kit/utils"
"github.com/gobwas/ws"
"github.com/gobwas/ws/wsutil"
"github.com/panjf2000/gnet/v2"
)
type wsConn struct {
utils.SpinLock
id uint64
kv map[string]string
c gnet.Conn
r *wsutil.Reader
w *wsutil.Writer
handshakeDone bool
rpcOutFactory kit.OutgoingRPCFactory
}
var _ kit.Conn = (*wsConn)(nil)
func newWebsocketConn(
id uint64, c gnet.Conn,
rpcOutFactory kit.OutgoingRPCFactory,
) *wsConn {
wsc := &wsConn{
w: wsutil.NewWriter(c, ws.StateServerSide, ws.OpText),
r: wsutil.NewReader(c, ws.StateServerSide),
id: id,
kv: map[string]string{},
c: c,
rpcOutFactory: rpcOutFactory,
}
return wsc
}
func (c *wsConn) Close() {
_ = c.c.Close()
}
func (c *wsConn) ConnID() uint64 {
return c.id
}
func (c *wsConn) ClientIP() string {
addr := c.c.RemoteAddr()
if addr == nil {
return ""
}
return addr.String()
}
func (c *wsConn) Write(data []byte) (int, error) {
c.Lock()
defer c.Unlock()
n, err := c.w.Write(data)
if err != nil {
return n, err
}
err = c.w.Flush()
return n, err
}
func (c *wsConn) WriteEnvelope(e *kit.Envelope) error {
outC := c.rpcOutFactory()
outC.InjectMessage(e.GetMsg())
outC.SetID(e.GetID())
e.WalkHdr(func(key string, val string) bool {
outC.SetHdr(key, val)
return true
})
data, err := outC.Marshal()
if err != nil {
return errors.Wrap(kit.ErrEncodeOutgoingMessageFailed, err)
}
_, err = c.Write(data)
outC.Release()
return err
}
func (c *wsConn) Stream() bool {
return true
}
func (c *wsConn) Walk(f func(key string, val string) bool) {
c.Lock()
defer c.Unlock()
for k, v := range c.kv {
if !f(k, v) {
return
}
}
}
func (c *wsConn) Get(key string) string {
c.Lock()
v := c.kv[key]
c.Unlock()
return v
}
func (c *wsConn) Set(key string, val string) {
c.Lock()
c.kv[key] = val
c.Unlock()
}