/
conn.go
120 lines (100 loc) · 2.67 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
package server
import (
"errors"
"net"
"time"
"github.com/quic-go/quic-go"
)
var (
errClosed = errors.New("tcp: use of closed connection")
errShutdown = errors.New("tcp: protocol is shutdown")
)
type serverConn struct {
net.Conn
idleTimeout time.Duration
maxDeadline time.Time
}
func (c *serverConn) Write(p []byte) (int, error) {
c.updateDeadline()
return c.Conn.Write(p)
}
func (c *serverConn) Read(b []byte) (int, error) {
c.updateDeadline()
return c.Conn.Read(b)
}
// rstAvoidanceDelay is the amount of time we sleep after closing the
// write side of a TCP connection before closing the entire socket.
// By sleeping, we increase the chances that the client sees our FIN
// and processes its final data before they process the subsequent RST
// from closing a connection with known unread data.
// This RST seems to occur mostly on BSD systems. (And Windows?)
// This timeout is somewhat arbitrary (~latency around the planet).
const rstAvoidanceDelay = 500 * time.Millisecond
type closeWriter interface {
CloseWrite() error
}
var _ closeWriter = (*net.TCPConn)(nil)
// closeWrite flushes any outstanding data and sends a FIN packet (if
// client is connected via TCP), signalling that we're done. We then
// pause for a bit, hoping the client processes it before any
// subsequent RST.
//
// See https://golang.org/issue/3595
func (c *serverConn) closeWriteAndWait() {
if tcp, ok := c.Conn.(closeWriter); ok {
tcp.CloseWrite()
}
time.Sleep(rstAvoidanceDelay)
}
func (c *serverConn) Close() error {
c.closeWriteAndWait()
//time.Sleep(time.Millisecond * 20)
return c.Conn.Close()
}
func (c *serverConn) updateDeadline() {
switch {
case c.idleTimeout > 0:
idleDeadline := time.Now().Add(c.idleTimeout)
if idleDeadline.Unix() < c.maxDeadline.Unix() || c.maxDeadline.IsZero() {
c.Conn.SetDeadline(idleDeadline)
return
}
fallthrough
default:
c.Conn.SetDeadline(c.maxDeadline)
}
}
type quicConn struct {
quic.Stream
conn quic.Connection
idleTimeout time.Duration
maxDeadline time.Time
}
func (c *quicConn) Write(p []byte) (int, error) {
c.updateDeadline()
return c.Stream.Write(p)
}
func (c *quicConn) Read(b []byte) (int, error) {
c.updateDeadline()
return c.Stream.Read(b)
}
func (c *quicConn) updateDeadline() {
switch {
case c.idleTimeout > 0:
idleDeadline := time.Now().Add(c.idleTimeout)
if idleDeadline.Unix() < c.maxDeadline.Unix() || c.maxDeadline.IsZero() {
c.Stream.SetDeadline(idleDeadline)
return
}
fallthrough
default:
c.Stream.SetDeadline(c.maxDeadline)
}
}
// LocalAddr locla addr
func (c *quicConn) LocalAddr() net.Addr {
return c.conn.LocalAddr()
}
func (c *quicConn) RemoteAddr() net.Addr {
return c.conn.RemoteAddr()
}