-
Notifications
You must be signed in to change notification settings - Fork 2.1k
/
tcp_socket.go
90 lines (73 loc) · 1.98 KB
/
tcp_socket.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
package privval
import (
"net"
"time"
)
// timeoutError can be used to check if an error returned from the netp package
// was due to a timeout.
type timeoutError interface {
Timeout() bool
}
// tcpTimeoutListener implements net.Listener.
var _ net.Listener = (*tcpTimeoutListener)(nil)
// tcpTimeoutListener wraps a *net.TCPListener to standardise protocol timeouts
// and potentially other tuning parameters.
type tcpTimeoutListener struct {
*net.TCPListener
acceptDeadline time.Duration
connDeadline time.Duration
period time.Duration
}
// timeoutConn wraps a net.Conn to standardise protocol timeouts / deadline resets.
type timeoutConn struct {
net.Conn
connDeadline time.Duration
}
// newTCPTimeoutListener returns an instance of tcpTimeoutListener.
func newTCPTimeoutListener(
ln net.Listener,
acceptDeadline, connDeadline time.Duration,
period time.Duration,
) tcpTimeoutListener {
return tcpTimeoutListener{
TCPListener: ln.(*net.TCPListener),
acceptDeadline: acceptDeadline,
connDeadline: connDeadline,
period: period,
}
}
// newTimeoutConn returns an instance of newTCPTimeoutConn.
func newTimeoutConn(
conn net.Conn,
connDeadline time.Duration) *timeoutConn {
return &timeoutConn{
conn,
connDeadline,
}
}
// Accept implements net.Listener.
func (ln tcpTimeoutListener) Accept() (net.Conn, error) {
err := ln.SetDeadline(time.Now().Add(ln.acceptDeadline))
if err != nil {
return nil, err
}
tc, err := ln.AcceptTCP()
if err != nil {
return nil, err
}
// Wrap the conn in our timeout wrapper
conn := newTimeoutConn(tc, ln.connDeadline)
return conn, nil
}
// Read implements net.Listener.
func (c timeoutConn) Read(b []byte) (n int, err error) {
// Reset deadline
c.Conn.SetReadDeadline(time.Now().Add(c.connDeadline))
return c.Conn.Read(b)
}
// Write implements net.Listener.
func (c timeoutConn) Write(b []byte) (n int, err error) {
// Reset deadline
c.Conn.SetWriteDeadline(time.Now().Add(c.connDeadline))
return c.Conn.Write(b)
}