forked from cretz/tor-dht-poc
-
Notifications
You must be signed in to change notification settings - Fork 0
/
conn.go
123 lines (101 loc) · 2.34 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
122
123
package websocket
import (
"io"
"net"
"sync"
"time"
"github.com/gorilla/websocket"
)
// GracefulCloseTimeout is the time to wait trying to gracefully close a
// connection before simply cutting it.
var GracefulCloseTimeout = 100 * time.Millisecond
var _ net.Conn = (*Conn)(nil)
// Conn implements net.Conn interface for gorilla/websocket.
type Conn struct {
*websocket.Conn
DefaultMessageType int
done func()
reader io.Reader
closeOnce sync.Once
}
func (c *Conn) Read(b []byte) (int, error) {
if c.reader == nil {
if err := c.prepNextReader(); err != nil {
return 0, err
}
}
for {
n, err := c.reader.Read(b)
switch err {
case io.EOF:
c.reader = nil
if n > 0 {
return n, nil
}
if err := c.prepNextReader(); err != nil {
return 0, err
}
// explicitly looping
default:
return n, err
}
}
}
func (c *Conn) prepNextReader() error {
t, r, err := c.Conn.NextReader()
if err != nil {
if wserr, ok := err.(*websocket.CloseError); ok {
if wserr.Code == 1000 || wserr.Code == 1005 {
return io.EOF
}
}
return err
}
if t == websocket.CloseMessage {
return io.EOF
}
c.reader = r
return nil
}
func (c *Conn) Write(b []byte) (n int, err error) {
if err := c.Conn.WriteMessage(c.DefaultMessageType, b); err != nil {
return 0, err
}
return len(b), nil
}
// Close closes the connection. Only the first call to Close will receive the
// close error, subsequent and concurrent calls will return nil.
// This method is thread-safe.
func (c *Conn) Close() error {
var err error
c.closeOnce.Do(func() {
if c.done != nil {
c.done()
// Be nice to GC
c.done = nil
}
c.Conn.WriteControl(websocket.CloseMessage, nil, time.Now().Add(GracefulCloseTimeout))
err = c.Conn.Close()
})
return err
}
func (c *Conn) SetDeadline(t time.Time) error {
if err := c.SetReadDeadline(t); err != nil {
return err
}
return c.SetWriteDeadline(t)
}
func (c *Conn) SetReadDeadline(t time.Time) error {
return c.Conn.SetReadDeadline(t)
}
func (c *Conn) SetWriteDeadline(t time.Time) error {
return c.Conn.SetWriteDeadline(t)
}
// NewConn creates a Conn given a regular gorilla/websocket Conn.
func NewConn(raw *websocket.Conn, done func()) *Conn {
return &Conn{
Conn: raw,
DefaultMessageType: websocket.BinaryMessage,
done: done,
}
}